"""GPU scale refit: per PQ2_0 block, d* = <c, teacher_rotated>/(support + 4), capped at 2x the
king's own scale, written back to a copy of the GGUF. All block math on CUDA."""
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 common.gguf import read_header
from common.gguf.dequant import hadamard, rotation
from common.gguf.types import PQ2_0
from model_dedup.gguf_checkpoint import GgufCheckpoint, canonical

ap = argparse.ArgumentParser()
ap.add_argument("--king", default="/root/models/king-ix")
ap.add_argument("--teacher", default="/root/models/teacher")
ap.add_argument("--out", default="/root/models/chal-v1")
ap.add_argument("--device", default="cuda")
args = ap.parse_args()

dev = args.device
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().to(dev)

idx = json.loads((Path(args.teacher) / "model.safetensors.index.json").read_text())
weight_map = idx["weight_map"]

import safetensors  # noqa: E402

_shards = {}


def teacher_tensor(hf_name):
    raw_name = "lm_head.weight" if hf_name == "lm_head.weight" else "model.language_model." + hf_name[len("model.") :]
    shard = weight_map[raw_name]
    if shard not in _shards:
        _shards[shard] = safetensors.safe_open(str(Path(args.teacher) / shard), "pt")
    t = _shards[shard].get_tensor(raw_name)
    return t.float().to(dev)


def tile(w, dim, head):
    shape = list(w.shape)
    x = w.reshape(*shape[:dim], hd.n_k, hd.rep, head, *shape[dim + 1 :])
    return x.transpose(dim, dim + 1).reshape(shape)


def rot_signs(name):
    kv = g.kv
    listed = [
        kv["prism.hadamard." + key].value
        for key in ("weight_names", "inverse_weight_names")
        if "prism.hadamard." + key in kv
    ]
    if not any(name in names for names in listed):
        return None
    widths = kv["prism.hadamard.sign_widths"].value
    values = kv["prism.hadamard.sign_values"].value
    width = g.tensors[name].shape[0]
    at = 0
    for w_ in widths:
        if w_ == width:
            return torch.from_numpy(np.asarray(values[at : at + w_], np.float32)).to(dev)
        at += w_
    raise ValueError(f"no sign vector of width {width} for {name}")


def hf_to_stored(name, w):
    if name.startswith("layers.") and hd.rep > 1:
        kind = name.split(".")[2]
        if kind == "in_proj_qkv" and name.endswith("in_proj_qkv.weight"):
            qk = 2 * hd.n_k * hd.k_dim
            w = torch.cat([w[:qk], tile(w[qk:], 0, hd.v_dim)])
        elif kind == "in_proj_z":
            w = tile(w, 0, hd.v_dim)
        elif kind in ("in_proj_a", "in_proj_b"):
            w = tile(w, 0, 1)
        elif kind == "out_proj" and not (rot_signs(name) is not None and ck._grouped):
            w = tile(w.T, 0, hd.v_dim).T
    signs = rot_signs(name)
    if signs is None:
        return w
    rows, width = w.shape
    nblk = width // block_size
    w = (w.reshape(rows, nblk, block_size) * signs.reshape(nblk, block_size)) @ H
    return w.reshape(rows, width)


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

done = 0
with open(dst, "r+b") as fh:
    for name, t in g.tensors.items():
        if t.ggml_type != PQ2_0:
            continue
        canon = canonical(name)
        if canon is None:
            continue
        hf_name = canon if canon.startswith(("lm_head", "norm.")) else "model." + canon
        w_t = hf_to_stored(name, teacher_tensor(hf_name))
        rows, width = w_t.shape
        nblk = rows // 128  # blocks tile the numpy rows (the gguf column axis)
        size = rows * width // 128 * 34
        raw = np.memmap(dst, np.uint8, "r+", g.data_start + t.offset, (size,))
        blocks = raw.reshape(-1, 34)
        qs = torch.from_numpy(blocks[:, 2:].copy()).to(dev)
        codes = torch.stack([(qs >> (2 * k)) & 3 for k in range(4)], -1).reshape(-1, 128)
        c_vals = (codes.float() - 1.0).reshape(width, nblk, 128).permute(1, 0, 2)
        d_orig = torch.from_numpy(blocks[:, :2].copy().view("<f2").astype(np.float32)).to(dev)
        d_orig = d_orig.reshape(width, nblk).T
        w_b = w_t.reshape(nblk, 128, width).permute(0, 2, 1)
        support = (c_vals * c_vals).sum(-1)
        num = (c_vals * w_b).sum(-1)
        d_new = num / (support + 4.0)
        d_new = torch.where(support >= 8, d_new, d_orig)
        cap = 2.0 * d_orig.abs()
        d_new = torch.clamp(d_new, -cap, cap).clamp(-5000.0, 5000.0)
        blocks[:, :2] = d_new.T.reshape(-1).cpu().numpy().astype(np.float16).astype("<f2").view(np.uint8).reshape(-1, 2)
        done += 1
        if done % 50 == 0:
            print(f"refit {done} tensors ({name})", flush=True)

print(f"refit {done} tensors -> {dst}")