import json
import torch
from safetensors.torch import load_file
from transformers import AutoModel, AutoTokenizer


def _dequantize(sd, cfg):
    """Dekuantisasi tensor int8/int4 (per-group) kembali ke float16."""
    bits = cfg["bits"]
    group_size = cfg["group_size"]
    shapes = cfg["shapes"]
    out = {}
    for name, shape in shapes.items():
        n_orig = int(torch.tensor(shape).prod())
        n_pad = ((n_orig + group_size - 1) // group_size) * group_size
        scale = sd[name + ".scale"].float()
        zero = sd[name + ".zero"].float()
        n_groups = n_pad // group_size
        scale_e = scale.reshape(n_groups, 1).expand(n_groups, group_size).reshape(-1)
        zero_e = zero.reshape(n_groups, 1).expand(n_groups, group_size).reshape(-1)
        if bits == 8:
            q = sd[name].float()
            w = (q - zero_e) * scale_e
        else:
            q = sd[name].to(torch.uint8).cpu()
            lo = (q & 0x0F).to(torch.int8)
            hi = (q >> 4).to(torch.int8)
            vals = torch.stack([hi, lo], dim=1).reshape(-1)
            vals = vals - 8
            w = (vals.float() - zero_e) * scale_e
        w = w[:n_orig].reshape(shape)
        out[name] = w.half()
    return out


def load_quantized_model(repo_id, bits, device=None):
    """Muat model semantic-lite-2 versi terkuantisasi (8 atau 4 bit).

    Args:
        repo_id: id repo Hugging Face (misal "ukung/semantic-lite-2")
        bits: 8 atau 4
        device: "cuda" atau "cpu" (default otomatis)
    """
    assert bits in (8, 4), "bits harus 8 atau 4"
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"

    # 1. Muat model fp16 dasar (struktur + bias/norm kecil)
    model = AutoModel.from_pretrained(repo_id, trust_remote_code=True)
    model = model.to(device)

    # 2. Muat tensor terkuantisasi
    from huggingface_hub import hf_hub_download
    q_path = hf_hub_download(repo_id, f"model_q{bits}.safetensors")
    c_path = hf_hub_download(repo_id, f"model_q{bits}_config.json")
    sd = load_file(q_path)
    cfg = json.load(open(c_path))
    deq = _dequantize(sd, cfg)

    # 3. Ganti weight 2D dengan versi dekuantisasi
    with torch.no_grad():
        for name, p in model.named_parameters():
            if name in deq:
                p.data = deq[name].to(device)
    return model


def encode(model, tokenizer, texts, device=None):
    """Encode teks jadi vektor embedding (sudah ternormalisasi L2)."""
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
    tok = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
    tok = {k: v.to(device) for k, v in tok.items()}
    with torch.no_grad():
        out = model(**tok)
    return out


if __name__ == "__main__":
    import sys
    bits = int(sys.argv[1]) if len(sys.argv) > 1 else 8
    repo = "ukung/semantic-lite-2"
    model = load_quantized_model(repo, bits)
    tok = AutoTokenizer.from_pretrained(repo)
    v = encode(model, tok, ["contoh kalimat bahasa Indonesia"])
    print(f"Q{bits} loaded, output shape: {tuple(v.shape)}")