"""QAT on the validated surrogate: per layer, straight-through ternary codes + block scales + norm
gains, distilled toward teacher trajectories with the eval's own metric (KL + hidden-state
distance). Writes a new GGUF at --out.

    .venv/bin/python qat.py --teacher /root/models/teacher --king /root/models/king-ix \
        --data /root/train_samples.json --out /root/models/chal-v2 [--layers all|7]
"""
import argparse
import json
import shutil
import sys
from pathlib import Path

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

from common.gguf import read_header
from common.gguf.dequant import hadamard, rotation
from common.gguf.types import PQ2_0, TYPES
from model_dedup.gguf_checkpoint import GgufCheckpoint, canonical

ap = argparse.ArgumentParser()
ap.add_argument("--teacher", default="/root/models/teacher")
ap.add_argument("--king", default="/root/models/king-ix")
ap.add_argument("--data", default="/root/train_samples.json")
ap.add_argument("--out", default="/root/models/chal-v2")
ap.add_argument("--layers", default="7,8,15,16,23,24,31,32,39,40,47,48,55,56,63")
ap.add_argument("--steps", type=int, default=25)
ap.add_argument("--lr", type=float, default=2e-3)
ap.add_argument("--alpha", type=float, default=1.0)
ap.add_argument("--seed", type=int, default=0)
args = ap.parse_args()

torch.manual_seed(args.seed)
ck = GgufCheckpoint(args.king, "cpu")
g = ck.g
hd = ck._heads
block_size = int(g.kv["prism.hadamard.block_size"].value)
H = torch.from_numpy(hadamard(block_size)).float()

# ---- training data
records = json.loads(Path(args.data).read_text())
print("training samples:", len(records))

# ---- the surrogate model and the king's state dict
model = AutoModelForCausalLM.from_pretrained(args.teacher, torch_dtype=torch.bfloat16, device_map="cuda")
model.eval()
sd_king = torch.load("/dev/shm/king_hf.pt", map_location="cpu", weights_only=True)

sys.path.insert(0, "/root/ingot/src")
from eval_service import metrics  # noqa: E402
from eval_service.config import LAYERS as CAP_LAYERS  # noqa: E402

heads = metrics.heads({"teacher": Path(args.teacher), "king": Path(args.king)}, ["cuda:0"])["cuda:0"]
head_t, head_k = heads["teacher"], heads["king"]

# ---- gguf byte io
src_gguf = Path(args.king) / "Ingot-IX-27B-Q2.gguf"
out_dir = Path(args.out)
out_dir.mkdir(parents=True, exist_ok=True)
dst = out_dir / "chal-v2.gguf"
shutil.copy(src_gguf, dst)


def gguf_bytes(tname):
    t = g.tensors[tname]
    n = int(np.prod(t.shape))
    size = n // TYPES[t.ggml_type].block_elems * TYPES[t.ggml_type].block_bytes
    return np.memmap(dst, np.uint8, "r+", g.data_start + t.offset, (size,))


def read_blocks(tname):
    t = g.tensors[tname]
    kind = TYPES[t.ggml_type]
    n = int(np.prod(t.shape))
    size = n // kind.block_elems * kind.block_bytes
    raw = np.memmap(dst, np.uint8, "r", g.data_start + t.offset, (size,))
    blocks = raw.reshape(-1, 34)
    d = blocks[:, :2].copy().view("<f2").astype(np.float32)
    qs = blocks[:, 2:]
    codes = np.stack([(qs >> (2 * k)) & 3 for k in range(4)], axis=-1).reshape(-1, 128)
    return codes, d


def unrotate(w, tname):
    """get()'s transform: HF-layout tensor from the stored one. w: (rows, width)."""
    folded = rotation(g, tname)
    if folded is not None:
        block, signs = folded
        rows, width = w.shape
        w = (w.reshape(rows, width // block, block) @ H).reshape(rows, width)
        w = w * torch.from_numpy(signs).float()
    return ck._hf_order(tname, w, folded is not None)


def hf_name(tname):
    c = canonical(tname)
    return c if c.startswith(("lm_head", "norm.")) else "model." + c


def hf_module(hf):
    parts = hf.split(".")
    m = model
    for p in parts[:-1]:
        m = getattr(m, p)
    return m


def layer_tensors(il):
    out = []
    for kind in ("attn_qkv", "attn_gate", "ssm_alpha", "ssm_beta", "ssm_out", "attn_q", "attn_k", "attn_v", "attn_output", "ffn_gate", "ffn_up", "ffn_down"):
        tname = f"blk.{il}.{kind}.weight"
        if tname in g.tensors and g.tensors[tname].ggml_type == PQ2_0:
            out.append(tname)
    return out


def train_layer(il):
    torch.cuda.empty_cache()
    print(f"--- layer {il}")
    tensors = layer_tensors(il)
    zs, ss = {}, {}
    for tname in tensors:
        codes, d = read_blocks(tname)
        t = g.tensors[tname]
        rows, width = int(np.prod(t.shape) // t.shape[0]), t.shape[0]
        nblk = rows * width // 128
        zs[tname] = torch.tensor(codes, dtype=torch.float32, device="cuda", requires_grad=True).reshape(nblk, 128)
        ss[tname] = torch.tensor(d, dtype=torch.float32, device="cuda", requires_grad=True).reshape(nblk)
    hooks = []
    for tname in tensors:
        hf = hf_name(tname)
        module = hf_module(hf)
        t = g.tensors[tname]
        rows, width = int(np.prod(t.shape) // t.shape[0]), t.shape[0]

        def make_pre(mod, tname=tname, rows=rows, width=width):
            def pre(mod, inp):
                z, s = zs[tname], ss[tname]
                c = torch.clamp(torch.round(z), 0, 2)
                W = (s * (c - 1.0)).reshape(rows, width)
                mod._orig_weight = mod.weight
                mod.weight = unrotate(W, tname).bfloat16()
            return pre

        def make_post(mod):
            def post(mod, inp, out):
                mod.weight = mod._orig_weight
            return post

        module.register_forward_pre_hook(make_pre(module))
        module.register_forward_hook(make_post(module))
        hooks.append(module)

    opt = torch.optim.Adam(list(zs.values()) + list(ss.values()), lr=args.lr)
    for step in range(args.steps):
        torch.cuda.empty_cache()
        total = torch.zeros((), device="cuda")
        n = 0
        for rec in records:
            tokens = rec["prompt"] + rec["answer"]
            if len(tokens) > 2600:
                continue
            first = len(rec["prompt"]) - 1
            ids = torch.tensor([tokens], device="cuda")
            out = model(input_ids=ids, output_hidden_states=True, use_cache=False)
            hs = [out.hidden_states[l][0, first:].float() for l in CAP_LAYERS]
            h_k = torch.stack(hs, 1)  # (rows, 8, 5120)
            st = metrics.load_states(rec["states"], tokens[:-1], first).float().to("cuda")
            m = min(h_k.shape[0], st.shape[0])
            h_k, st = h_k[:m], st[:m]
            lp_t = head_t.logprobs(st[:, -1])
            lp_k = head_k.logprobs(h_k[:, -1])
            kl = (lp_t.exp() * (lp_t - lp_k)).sum(-1)
            a = st[1:]
            b = h_k[1:]
            dist = ((b - a).norm(dim=-1) / a.norm(dim=-1)).mean(1)
            loss = kl.mean() + args.alpha * dist.mean()
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(list(zs.values()) + list(ss.values()), 1.0)
            opt.step()
            total += loss.detach()
            n += 1
            del out, hs, h_k, st, lp_t, lp_k
        print(f"  step {step}: loss {total.item() / n:.5f} over {n} samples", flush=True)

    for tname in tensors:
        z, s = zs[tname].detach().cpu().numpy(), ss[tname].detach().cpu().numpy()
        c = np.clip(np.round(z), 0, 2).astype(np.uint8)
        rows, width = int(np.prod(g.tensors[tname].shape) // g.tensors[tname].shape[0]), g.tensors[tname].shape[0]
        nblk = rows * width // 128
        raw = gguf_bytes(tname).reshape(-1, 34)
        raw[:, :2] = s.astype(np.float16).reshape(-1, 1).astype("<f2").view(np.uint8).reshape(-1, 2)
        qs = np.zeros((nblk, 32), np.uint8)
        for j in range(128):
            qs[:, j // 4] |= (c[:, j] << (2 * (j % 4))).astype(np.uint8)
        raw[:, 2:] = qs
        print(f"  wrote {tname}", flush=True)
    for mod in hooks:
        mod._forward_pre_hooks.clear()
        mod._forward_hooks.clear()
    del zs, ss
    torch.cuda.empty_cache()


for il in (int(x) for x in args.layers.split(",") if x):
    train_layer(il)

print("saved", dst)