"""Challenger v1: the king with every PQ2_0 block's fp16 scale refit by least squares against the
teacher's rotated weights (codes untouched). Works entirely in the stored (rotated) layout, where
stored = W_hf · signs · H per block for both folded weights and latent embeddings.

    .venv/bin/python scale_refit.py [--out /root/models/chal-v1]
"""
import argparse
import shutil
import sys
from pathlib import Path

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

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, Heads

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("--max-tensors", type=int, default=0, help="limit for testing")
args = ap.parse_args()

ck = GgufCheckpoint(args.king, "cpu")
g = ck.g
heads = ck._heads
block_size = int(g.kv["prism.hadamard.block_size"].value)
H = hadamard(block_size)
print("block_size", block_size, "heads", heads)

teacher_shards = {}
import json

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

def teacher_tensor(hf_name: str) -> np.ndarray:
    raw_name = (
        hf_name
        if hf_name == "lm_head.weight"
        else "model.language_model." + hf_name[len("model.") :]
    )
    shard = weight_map[raw_name]
    if shard not in teacher_shards:
        import safetensors

        teacher_shards[shard] = safetensors.safe_open(str(Path(args.teacher) / shard), "pt")
    t = teacher_shards[shard].get_tensor(raw_name)
    return t.to(dtype=torch.float32).numpy()

import torch  # noqa: E402


def tile(w: np.ndarray, dim: int, head: int) -> np.ndarray:
    """Inverse of GgufCheckpoint._untile: HF grouped [n_k, rep, head] -> stored tiled [rep, n_k, head]."""
    shape = list(w.shape)
    x = w.reshape(*shape[:dim], heads.n_k, heads.rep, head, *shape[dim + 1 :])
    return x.transpose(dim, dim + 1).reshape(shape)


def hf_to_stored(name: str, w: np.ndarray) -> np.ndarray:
    """Teacher tensor, HF layout -> the king's stored layout (row order, then rotation)."""
    if name.startswith("layers.") and heads.rep > 1:
        kind = name.split(".")[2]
        if kind == "in_proj_qkv" and name.endswith("in_proj_qkv.weight"):
            qk = 2 * heads.n_k * heads.k_dim
            w = np.concatenate([w[:qk], tile(w[qk:], 0, heads.v_dim)])
        elif kind == "in_proj_z":
            w = tile(w, 0, heads.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, heads.v_dim).T
    signs = rot_signs(name)
    if signs is None:
        return w
    rows, width = w.shape
    nblk = width // block_size
    # stored = (W · signs) · H per col block, as get()'s inverse: W = (stored · H) · signs
    w = (w.reshape(rows, nblk, block_size) * signs.reshape(nblk, block_size)) @ H
    return w.reshape(rows, width)


def rot_signs(name: str) -> np.ndarray | None:
    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 np.asarray(values[at : at + w_], np.float32)
        at += w_
    raise ValueError(f"no sign vector of width {width} for {name}")


# the gguf tensor name each HF name maps to (dedup canonical, with the .weight suffix)
from model_dedup.gguf_checkpoint import canonical as gguf_canonical  # noqa: E402

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

done = 0
with open(dst, "r+b") as fh:
    for name, t in g.tensors.items():
        if t.ggml_type != PQ2_0:
            continue
        if args.max_tensors and done >= args.max_tensors:
            break
        canon = gguf_canonical(name)
        if canon is None:
            print("skip (no hf name):", name)
            continue
        hf_name = "model." + canon if not canon.startswith(("lm_head", "norm.")) else canon
        # canonical gives e.g. layers.0.self_attn.q_proj.weight; HF: model.layers.0...
        w_t = teacher_tensor(hf_name)
        w_t = hf_to_stored(name, w_t)
        rows, width = w_t.shape
        nblk = width // 128  # the PQ2_0 block is 128 wide; the Hadamard block is not
        # read raw blocks
        size = rows * nblk * 34
        raw = np.memmap(dst, np.uint8, "r+", g.data_start + t.offset, (size,))
        blocks = raw.reshape(-1, 34)
        qs = blocks[:, 2:]
        codes = np.stack([(qs >> (2 * k)) & 3 for k in range(4)], axis=-1).reshape(rows, nblk, 128)
        if (codes == 3).any():
            raise SystemExit(f"{name}: a code 3 in the king")
        c_vals = codes.astype(np.float32) - 1.0
        w_t_b = w_t.reshape(rows, nblk, 128)
        d_orig = blocks[:, :2].copy().view("<f2").astype(np.float32).reshape(rows, nblk)
        support = (c_vals * c_vals).sum(-1)
        num = (c_vals * w_t_b).sum(-1)
        # a block of mostly-zero codes has no stable least-squares scale: ridge + a support floor,
        # and never more than double the king's own scale (its codes were fit with it)
        d_new = num / (support + 4.0)
        d_new = np.where(support >= 8, d_new, d_orig)
        d_new = np.clip(d_new, -2.0 * np.abs(d_orig), 2.0 * np.abs(d_orig))
        d_new = np.clip(d_new, -5000.0, 5000.0)
        d_new = d_new.astype(np.float16)
        blocks[:, :2] = d_new.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}")