import sys

sys.path.insert(0, "/tmp/opencode/tl")
import numpy as np
import torch

import export_full as E
import int4hf as H
from transformers import AutoModelForCausalLM

GREEDY_FLOAT = [4612, 2619, 260, 647, 31844, 635, 393, 260, 1496, 2789, 3557, 21075, 31843, 1100,
                4838, 289, 826, 351, 600, 17360, 1086, 31843, 1591, 1124, 31844, 700, 2815, 260,
                1421, 3476, 288, 600, 2103, 31843, 1100, 393, 900, 6055, 291, 2598]


def deq(w):
    q, s = E.quantize(w)
    sr = np.round(s * 2 ** 15) / 2 ** 15
    return torch.tensor(q.astype(np.float32) * sr[:, None].astype(np.float32))


base = AutoModelForCausalLM.from_pretrained("/tmp/opencode/tl", dtype=torch.float32).eval()
t = E.load()


def build(groups):
    m = H.int4_model()
    with torch.no_grad():
        if "head" in groups:
            m.lm_head.weight.copy_(torch.tensor(t["lm_head.weight"] * t["model.norm.weight"][None, :]))
        if "embed" in groups:
            m.model.embed_tokens.weight.copy_(torch.tensor(t["model.embed_tokens.weight"]))
        for l in range(8):
            L = m.model.layers[l]
            ga = t[f"model.layers.{l}.input_layernorm.weight"]
            gf = t[f"model.layers.{l}.post_attention_layernorm.weight"]
            pre = f"model.layers.{l}."
            if "qkv" in groups:
                L.self_attn.q_proj.weight.copy_(torch.tensor(t[pre + "self_attn.q_proj.weight"] * ga[None, :]))
                L.self_attn.k_proj.weight.copy_(torch.tensor(t[pre + "self_attn.k_proj.weight"] * ga[None, :]))
                L.self_attn.v_proj.weight.copy_(torch.tensor(t[pre + "self_attn.v_proj.weight"] * ga[None, :]))
            if "o" in groups:
                L.self_attn.o_proj.weight.copy_(torch.tensor(t[pre + "self_attn.o_proj.weight"]))
            if "gateup" in groups:
                L.mlp.gate_proj.weight.copy_(torch.tensor(t[pre + "mlp.gate_proj.weight"] * gf[None, :]))
                L.mlp.up_proj.weight.copy_(torch.tensor(t[pre + "mlp.up_proj.weight"] * gf[None, :]))
            if "down" in groups:
                L.mlp.down_proj.weight.copy_(torch.tensor(t[pre + "mlp.down_proj.weight"]))
    return m


def greedy(m, n=40):
    ids = torch.tensor([[1]])
    out = []
    with torch.no_grad():
        for _ in range(n):
            lg = m(ids).logits[0, -1]
            k = int(lg.argmax())
            out.append(k)
            ids = torch.cat([ids, torch.tensor([[k]])], 1)
    return out


def hit(out):
    n = 0
    for a, b in zip(out, GREEDY_FLOAT):
        if a != b:
            break
        n += 1
    return n


print("base int4: ", hit(greedy(build(set()))), "/40", flush=True)
for g in (["head"], ["embed"], ["qkv"], ["o"], ["gateup"], ["down"]):
    print(f"solo {g[0]:8s} en float:", hit(greedy(build(set(g)))), "/40", flush=True)