import torch
from exl2conv.fasttensors import STFile
from exl2conv.architecture import ExLlamaV2ArchParams
import os, glob, json

class ExLlamaV2Config:

    debug_mode = False
    model_dir: str = None                       # Directory containing model files

    max_seq_len: int = 2048                     # Maximum sequence length. Sequences longer than this will throw an exception
    max_batch_size: int = 1                     # Maximum size of batches to process
    max_input_len: int = 2048                   # Maximum length of input IDs in a single forward pass. Sequences longer than this will be processed in multiple steps
    max_attention_size: int = 2048 ** 2         # Sequences will be processed in chunks to keep the size of the attention weights matrix <= this

    scale_pos_emb: float = 1.0                  # Factor by which to scale positional embeddings, e.g. for 4096-token sequence use a scaling factor of 2.0, requires finetuned model or LoRA
    scale_alpha_value: float = 1.0              # Alpha value for NTK RoPE scaling. Similar to compress_pos_emb but works without finetuned model

    no_flash_attn: bool = False                 # Implementation will automatically use flash-attn-2 when available
    fasttensors: bool = False                   # Experimental, Linux only

    # Loaded/set by .prepare():

    architecture: str
    arch: ExLlamaV2ArchParams = None

    model_config: str
    tensor_file_map: dict
    tensor_files: list

    tokenizer_path: str

    bos_token_id: int
    eos_token_id: int
    pad_token_id: int

    hidden_size: int
    initializer_range: int
    intermediate_size: int
    num_attention_heads: int
    num_key_value_heads: int
    num_key_value_groups: int
    num_hidden_layers: int
    norm_eps: float
    vocab_size: int
    rotary_embedding_base: float = 10000.0      # Constant for all Llama models, nodified by .prepare() if scale_alpha_value != 1.0
    head_dim: int = 128                         # Constant for all Llama models, except 3b
    num_experts: int = None
    num_experts_per_token: int = None

    checkpoint_fused_mlp: bool = False


    def __init__(self):
        pass


    # Set low-mem options

    def set_low_mem(self):

        self.max_input_len = 1024
        self.max_attention_size = 1024 ** 2


    # Populate config with required files from model_dir

    def prepare(self, no_tensors = False):

        assert self.model_dir is not None, "No model_dir specified in ExLlamaV2Config"
        assert os.path.exists(self.model_dir), "Can't find " + self.model_dir

        # Load config.json

        self.model_config = os.path.join(self.model_dir, "config.json")
        assert os.path.exists(self.model_config), "Can't find " + self.model_config

        with open(self.model_config, encoding = "utf8") as f:
            read_config = json.load(f)

        # Model architecture

        assert len(read_config["architectures"]) == 1, "Multiple architectures defined in config.json"
        arch_string = read_config["architectures"][0]
        self.arch = ExLlamaV2ArchParams(arch_string, read_config)

        # Vocab params

        self.bos_token_id = read_config.get("bos_token_id", 1)
        self.eos_token_id = read_config.get("eos_token_id", 2)
        self.pad_token_id = read_config.get("pad_token_id", 0)
        self.vocab_size = read_config["vocab_size"]

        # Standard params

        self.initializer_range = read_config["initializer_range"]
        self.num_hidden_layers = read_config["num_hidden_layers"]

        # Norm params

        self.norm_eps = read_config[self.arch.norm_eps_key]

        # Attn params

        self.num_attention_heads = read_config["num_attention_heads"]

        if "num_key_value_heads" in read_config:
            self.num_key_value_heads = read_config["num_key_value_heads"]
            self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads
        else:
            self.num_key_value_heads = self.num_attention_heads
            self.num_key_value_groups = 1

        # MLP params

        self.intermediate_size = read_config["intermediate_size"]
        self.num_experts = read_config.get("num_local_experts", None)
        self.num_experts_per_token = read_config.get("num_experts_per_tok", None)

        # Positional embeddings

        self.rotary_embedding_base = read_config["rope_theta"] if "rope_theta" in read_config else 10000.0

        if "max_sequence_length" in read_config: self.max_seq_len = read_config["max_sequence_length"]
        elif "max_position_embeddings" in read_config: self.max_seq_len = read_config["max_position_embeddings"]

        rs = read_config.get("rope_scaling", None)
        if rs and "factor" in rs:
            factor = rs["factor"]
            scaling_type = rs.get("type", None)
            if scaling_type == "linear":
                self.scale_pos_emb = factor
            # elif scaling_type == "yarn":
            #     self.scale_alpha_value = factor

        # Model dimensions

        self.hidden_size = read_config["hidden_size"]
        self.head_dim = read_config.get("head_dim", self.hidden_size // self.num_attention_heads)

        # Create map of model tensors

        if no_tensors: return

        self.tensor_file_map = {}

        st_pattern = os.path.join(self.model_dir, "*.safetensors")
        self.tensor_files = glob.glob(st_pattern)

        if len(self.tensor_files) == 0:
            raise ValueError(f" ## No .safetensors files found in {self.model_dir}")

        for st_file in self.tensor_files:
            f = STFile.open(st_file, fast = self.fasttensors)
            for key in f.get_dict():
                self.tensor_file_map[key] = st_file

        # For loading checkpoints with fused MLP layers

        if "model.layers.0.mlp.down_proj.weight" not in self.tensor_file_map and \
            "model.layers.0.mlp.swiglu.w12.weight" in self.tensor_file_map:
            self.checkpoint_fused_mlp = True
            self.arch.make_fused_mlp()

        # Make sure we found all the layers we need

        expect_keys = self.arch.expect_keys.copy()

        for layer_idx in range(self.num_hidden_layers):
            for ks in self.arch.layer_keys:
                prefixes = [f"model.layers.{layer_idx}.{k}" for k in ks]
                expect_keys.append(prefixes)

        for prefixes in expect_keys:
            for prefix in prefixes:
                if any(key.startswith(prefix) for key in self.tensor_file_map):
                    break
            else:
                raise ValueError(f" ## Could not find {prefix}.* in model")

