"""Identity round-trip sanity v2 (chunked): rebuild PQ2_0 scales from the king's own
weights (codes unchanged) along the correct block axis; write back as a challenger."""
import hashlib
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.types import PQ2_0

KING = Path("/root/models/king-ix/Ingot-IX-27B-Q2.gguf")
OUT_DIR = Path("/root/models/chal-identity")
OUT = OUT_DIR / KING.name
CHUNK = 1 << 20  # blocks per chunk

g = read_header(KING)
pq = [(name, t) for name, t in g.tensors.items() if t.ggml_type == PQ2_0]
print(f"PQ2_0 tensors: {len(pq)}")


def codes_d(blocks):
    d = blocks[:, :2].copy().view("<f2").astype(np.float32).reshape(-1)
    qs = blocks[:, 2:]
    codes = np.stack([(qs >> (2 * k)) & 3 for k in range(4)], axis=-1).reshape(-1, 128).astype(np.float32) - 1.0
    return codes, d


def refit_scales(blocks):
    codes, d = codes_d(blocks)
    den = (codes * codes).sum(1)
    num = (codes * codes * d[:, None]).sum(1)
    return np.where(den > 0, num / np.maximum(den, 1e-9), d)


# ---- Part A: axis diagnostic on small PQ2_0 tensors ----
print("=== AXIS DIAGNOSTIC ===")
shown = 0
for name, t in pq:
    n_blocks = int(np.prod(t.shape)) // 128
    if n_blocks > 200_000:
        continue
    b = np.asarray(np.memmap(KING, np.uint8, "r", g.data_start + t.offset, (n_blocks * 34,)).reshape(-1, 34))
    d_orig = b[:, :2].copy().view("<f2").astype(np.float32).reshape(-1)
    d_ok = refit_scales(b).astype(np.float16).astype(np.float32)
    err = np.abs(d_ok - d_orig)
    scale = max(np.abs(d_orig).max(), 1e-9)
    # v1's wrong axis on a square grid, for the failure record
    side = int(round(np.sqrt(n_blocks)))
    bad = None
    if side * side == n_blocks:
        codes, _ = codes_d(b)
        dT = d_orig.reshape(side, side).T.reshape(-1)
        den = (codes * codes).sum(1)
        num = (codes * codes * dT[:, None]).sum(1)
        d_bad = np.where(den > 0, num / np.maximum(den, 1e-9), dT).astype(np.float16).astype(np.float32)
        bad = (np.abs(d_bad - d_orig).max() / scale)
    print(f"  {name[:46]:<46} blocks={n_blocks:>8} maxerr={err.max()/scale:.2e} bits_flipped={(d_ok.astype('<f2').view(np.uint16) != d_orig.astype('<f2').view(np.uint16)).sum()}"
          + (f"   v1-axis maxerr={bad:.2e}" if bad is not None else ""))
    shown += 1
    if shown >= 8:
        break

# ---- Part B: write identity challenger ----
OUT_DIR.mkdir(parents=True, exist_ok=True)
shutil.copy(KING, OUT)
patched = changed = 0
with open(OUT, "r+b") as fh:
    for name, t in pq:
        n_blocks = int(np.prod(t.shape)) // 128
        raw = np.memmap(OUT, np.uint8, "r+", g.data_start + t.offset, (n_blocks * 34,))
        for s in range(0, n_blocks, CHUNK):
            e = min(s + CHUNK, n_blocks)
            b = np.asarray(raw[s * 34:e * 34]).reshape(-1, 34)
            d_new = refit_scales(b).astype(np.float16).astype("<f2").view(np.uint8).reshape(-1, 2)
            changed += int((d_new != b[:, :2]).sum() // 2)
            b[:, :2] = d_new
        patched += 1
        if patched % 100 == 0:
            print(f"  patched {patched}/{len(pq)}", flush=True)

print(f"\npatched {patched} tensors; scales differing from original: {changed}")


def h(p):
    m = hashlib.sha256()
    with open(p, "rb") as f:
        for c in iter(lambda: f.read(1 << 24), b""):
            m.update(c)
    return m.hexdigest()[:16]


print(f"king sha256     {h(KING)}")
print(f"identity sha256 {h(OUT)}")
print("WROTE", OUT)