"""Build the king's full HF-layout state dict (BF16) by dequantizing the GGUF with the repo's
own code, adding the ssm tensors the dedup ignores. Saves /dev/shm/king_hf.pt."""
import sys

sys.path.insert(0, "/root/ingot/src")
import torch
import numpy as np

from common.gguf import read_header
from common.gguf.dequant import tensor as deq_tensor
from model_dedup.gguf_checkpoint import GgufCheckpoint

GGUF_DIR = "/root/models/king-ix"

ck = GgufCheckpoint(GGUF_DIR, "cpu")
g = ck.g

def raw(key):
    w = torch.from_numpy(deq_tensor(ck.file, g, key))
    return w.float()

sd = {}
for name in ck.names():
    sd[name] = ck.get(name)

n = int(g.kv["qwen35.block_count"].value)
hd = ck._heads


def untile(w: torch.Tensor, dim: int, head: int) -> torch.Tensor:
    """GgufCheckpoint._untile: stored tiled [rep, n_k, head] -> HF grouped [n_k, rep, head]."""
    shape = list(w.shape)
    x = w.reshape(*shape[:dim], hd.rep, hd.n_k, head, *shape[dim + 1 :])
    return x.transpose(dim, dim + 1).contiguous().reshape(shape)


for il in range(n):
    if f"blk.{il}.ssm_a" not in g.tensors:  # full-attention layers have no ssm tensors
        continue
    p = f"layers.{il}.linear_attn."
    a = raw(f"blk.{il}.ssm_a")
    sd[p + "A_log"] = untile(torch.log(-a).unsqueeze(-1), 0, 1).squeeze(-1)
    dt = raw(f"blk.{il}.ssm_dt.bias")
    sd[p + "dt_bias"] = untile(dt.unsqueeze(-1), 0, 1).squeeze(-1)
    conv = raw(f"blk.{il}.ssm_conv1d.weight").unsqueeze(1)  # (channels, 1, 4)
    qk_channels = 2 * hd.n_k * hd.k_dim
    v = untile(conv[qk_channels:], 0, hd.v_dim)
    sd[p + "conv1d.weight"] = torch.cat([conv[:qk_channels], v], dim=0)
    sd[p + "norm.weight"] = raw(f"blk.{il}.ssm_norm.weight")

# check against the teacher's shapes
from transformers import AutoModelForCausalLM  # noqa: E402

teacher = AutoModelForCausalLM.from_pretrained(
    "/root/models/teacher", torch_dtype=torch.bfloat16, device_map="cpu"
)
td = teacher.state_dict()
bad = []
for k, v in sd.items():
    hits = [x for x in td if x.endswith(k)]
    if not hits:
        bad.append((k, "no key"))
        continue
    if tuple(td[hits[0]].shape) != tuple(v.shape):
        bad.append((k, f"shape {tuple(v.shape)} vs {tuple(td[hits[0]].shape)}"))
print("mismatches:", len(bad), bad[:10])

# overlay the king's tensors onto the teacher model (MTP/visual stay the teacher's)
out = {}
for k, v in td.items():
    hit = next((x for x in sd if k.endswith(x)), None)
    out[k] = (sd[hit].to(torch.bfloat16) if hit is not None else v.to(torch.bfloat16))
print("saving", len(out), "tensors")
torch.save(out, "/dev/shm/king_hf.pt")
print("saved /dev/shm/king_hf.pt")