from dataclasses import dataclass, field
from pathlib import Path
import sys

import torch
from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM, modeling_utils


@dataclass
class AWQConfig:
    ckpt: str = field(
        default=None,
        metadata={
            "help": "Load quantized model. The path to the local AWQ checkpoint."
        },
    )
    wbits: int = field(default=16, metadata={"help": "#bits to use for quantization"})
    groupsize: int = field(
        default=-1,
        metadata={"help": "Groupsize to use for quantization; default uses full row."},
    )


def load_awq_quantized(model_name, awq_config: AWQConfig, device):
    print("Loading AWQ quantized model...")

    try:
        from tinychat.utils import load_quant
        from tinychat.modules import make_quant_norm, make_quant_attn, make_fused_mlp
    except ImportError as e:
        print(f"Error: Failed to import tinychat. {e}")
        print("Please double check if you have successfully installed AWQ")
        print("See https://github.com/lm-sys/FastChat/blob/main/docs/awq.md")
        sys.exit(-1)

    config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained(
        model_name, use_fast=False, trust_remote_code=True
    )

    def skip(*args, **kwargs):
        pass

    torch.nn.init.kaiming_uniform_ = skip
    torch.nn.init.kaiming_normal_ = skip
    torch.nn.init.uniform_ = skip
    torch.nn.init.normal_ = skip
    modeling_utils._init_weights = False

    torch.set_default_dtype(torch.half)
    model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)

    if any(name in find_awq_ckpt(awq_config) for name in ["llama", "vicuna"]):
        model = load_quant.load_awq_llama_fast(
            model,
            find_awq_ckpt(awq_config),
            awq_config.wbits,
            awq_config.groupsize,
            device,
        )
        make_quant_attn(model, device)
        make_quant_norm(model)
        make_fused_mlp(model)
    else:
        model = load_quant.load_awq_model(
            model,
            find_awq_ckpt(awq_config),
            awq_config.wbits,
            awq_config.groupsize,
            device,
        )
    return model, tokenizer


def find_awq_ckpt(awq_config: AWQConfig):
    if Path(awq_config.ckpt).is_file():
        return awq_config.ckpt

    for ext in ["*.pt", "*.safetensors"]:
        matched_result = sorted(Path(awq_config.ckpt).glob(ext))
        if len(matched_result) > 0:
            return str(matched_result[-1])

    print("Error: AWQ checkpoint not found")
    sys.exit(1)