"""End-to-end norm tuning on the surrogate: every RMSNorm weight (F32, free in the GGUF) trained
against teacher trajectories with the eval's metric (KL + hidden-state distance), the same loss
the eval uses. Writes a new GGUF.

    .venv/bin/python qat_norms.py --data /root/train_samples.json --out /root/models/chal-norms
"""
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.types import TYPES
from model_dedup.gguf_checkpoint import GgufCheckpoint

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-norms")
ap.add_argument("--steps", type=int, default=15)
ap.add_argument("--lr", type=float, default=3e-3)
ap.add_argument("--alpha", type=float, default=1.0)
ap.add_argument("--max-tokens", type=int, default=2200)
args = ap.parse_args()

ck = GgufCheckpoint(args.king, "cpu")
g = ck.g
records = json.loads(Path(args.data).read_text())
print("samples:", len(records))

model = AutoModelForCausalLM.from_pretrained(args.teacher, torch_dtype=torch.bfloat16, device_map="cuda")
model.eval()
model.gradient_checkpointing_enable()
sd_king = torch.load("/dev/shm/king_hf.pt", map_location="cpu", weights_only=True)
model.load_state_dict(sd_king, strict=True)

from eval_service import metrics  # noqa: E402
from eval_service.config import 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"]

params = []
for name, p in model.named_parameters():
    p.requires_grad_(False)
    if any(k in name for k in ("layernorm.weight", "norm.weight")):
        p.requires_grad_(True)
        params.append(p)
print("trainable norms:", len(params))

opt = torch.optim.Adam(params, lr=args.lr)
for step in range(args.steps):
    torch.cuda.empty_cache()
    total, n = 0.0, 0
    for rec in records:
        tokens = rec["prompt"] + rec["answer"]
        if len(tokens) > args.max_tokens:
            continue
        first = len(rec["prompt"]) - 1
        ids = torch.tensor([tokens], device="cuda")
        out = model(input_ids=ids, output_hidden_states=True, use_cache=False)
        h_k = torch.stack([out.hidden_states[l][0, first:].float() for l in LAYERS], 1)
        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).mean()
        a, b = st[1:], h_k[1:]
        dist = ((b - a).norm(dim=-1) / a.norm(dim=-1)).mean()
        loss = kl + args.alpha * dist
        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(params, 1.0)
        opt.step()
        total += loss.item()
        n += 1
        del out, h_k, st, lp_t, lp_k, ids, loss
    print(f"step {step}: loss {total / n:.5f} over {n} samples", flush=True)

# ---- write the norms back (attn/post/output norms stored as hf + 1; ssm_norm as-is)
out_dir = Path(args.out)
out_dir.mkdir(parents=True, exist_ok=True)
dst = out_dir / "chal-norms.gguf"
shutil.copy(Path(args.king) / "Ingot-IX-27B-Q2.gguf", dst)

def write_tensor(tname, values):
    t = g.tensors[tname]
    size = int(np.prod(t.shape)) * 4
    raw = np.memmap(dst, np.uint8, "r+", g.data_start + t.offset, (size,))
    raw.view("<f4")[:] = values.astype(np.float32).ravel()

n = int(g.kv["qwen35.block_count"].value)
with torch.no_grad():
    for il in range(n):
        for gguf_kind, hf_kind, plus1 in (
            ("attn_norm", "input_layernorm", True),
            ("post_attention_norm", "post_attention_layernorm", True),
            ("ssm_norm", "linear_attn.norm", False),
        ):
            if f"blk.{il}.{gguf_kind}.weight" not in g.tensors:
                continue
            hf = f"model.layers.{il}.{hf_kind}.weight"
            mod = model
            for part in hf.split(".")[:-1]:
                mod = getattr(mod, part)
            w = mod.weight.detach().cpu().float()
            write_tensor(f"blk.{il}.{gguf_kind}.weight", w.numpy() + (1.0 if plus1 else 0.0))
    mod = model.model
    w = mod.norm.weight.detach().cpu().float()
    write_tensor("output_norm.weight", w.numpy() + 1.0)
print("saved", dst)