# SPDX-License-Identifier: Apache-2.0
"""Host-only study (round 8): can a channel PERMUTATION make bfp8 activations accurate enough?

bfp8_b shares one exponent per 16 consecutive channels of a token (one row of a 16x16 face), so a few
large-magnitude channels force a coarse step on their 15 neighbours.  Permuting channels so that channels of
similar magnitude share a group is exact algebra for the model (the weights are permuted to match on the host),
but it changes how much the bfp8 rounding hurts.

Fake-quantizes (CPU fp32 reference model, forward pre-hooks):
  ln1 = qkv input (the model already runs it in bfp8), ln2 = fc1 input, g = fc2 input (GELU output)
with identity or magnitude-sorted channel orders, and reports the 157-d output PCC vs the unquantized model over
the 32 acc.py crops (calibration of the orders on 8 DIFFERENT crops: the same hands with other bbox jitter).
usage: bfp8_perm_study.py [n_eval]"""
from __future__ import annotations

import json
import sys
from pathlib import Path

import torch

_REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO))
sys.path.insert(0, str(_REPO / "scripts"))
from hamer import Hamer, load_hamer_checkpoint  # noqa: E402
from hamer.preprocess import preprocess  # noqa: E402
from acc import crops  # noqa: E402

torch.set_num_threads(16)
import os
TRUNC = os.environ.get('BFP8_TRUNC') == '1'  # round toward zero (packer-truncation hypothesis)


def bfp8(x: torch.Tensor, perm: torch.Tensor | None) -> torch.Tensor:
    """bfp8_b fake quantization along the last dim (groups of 16 after `perm`), round to nearest."""
    if perm is not None:
        x = x[..., perm]
    s = x.shape
    g = x.reshape(*s[:-1], s[-1] // 16, 16)
    mx = g.abs().amax(-1, keepdim=True).clamp_min(1e-30)
    e = torch.floor(torch.log2(mx))
    step = torch.pow(2.0, e - 6)
    r = torch.trunc(g / step) if TRUNC else torch.round(g / step)
    q = (r.clamp(-127, 127) * step).reshape(s)
    if perm is not None:
        q = q[..., torch.argsort(perm)]
    return q


def calib_crops():
    media = _REPO.parent / "media"
    boxes = json.loads((media / "bboxes.json").read_text())
    out = []
    for name in sorted(boxes):
        for e in boxes[name]:
            x1, y1, x2, y2 = e["bbox"]
            w, h = x2 - x1, y2 - y1
            bb = (x1 - 0.05 * w, y1 + 0.05 * h, x2 - 0.05 * w, y2 + 0.05 * h)
            x, _ = preprocess(media / name, bbox_xyxy=bb, is_right=bool(e.get("is_right", 1)))
            out.append(x)
    return out


def main():
    n_eval = int(sys.argv[1]) if len(sys.argv) > 1 else 32
    hf = Path.home() / ".cache/huggingface/hub/models--changh95--hamer-weights/snapshots/e8217b46d54b4e8afaaf0dd6efb6b62341c97d32"
    ref = Hamer().eval()
    load_hamer_checkpoint(ref, hf / "hamer.ckpt", hf / "mano_mean_params.npz")
    blocks = ref.backbone.blocks
    L = len(blocks)

    # ---- calibration: per-channel mean |x| of ln1 / ln2 / g per layer
    stats = {k: [torch.zeros(1280 if k != "g" else 5120) for _ in range(L)] for k in ("ln1", "ln2", "g")}
    hooks = []
    for i, b in enumerate(blocks):
        for k, mod in (("ln1", b.attn.qkv), ("ln2", b.mlp.fc1), ("g", b.mlp.fc2)):
            def pre(m, inp, i=i, k=k):
                stats[k][i] += inp[0].detach().abs().reshape(-1, inp[0].shape[-1]).mean(0)
            hooks.append(mod.register_forward_pre_hook(pre))
    with torch.no_grad():
        for x in calib_crops():
            ref(x)
    for h in hooks:
        h.remove()
    # per-layer orders (hidden channels: free per layer); the stream orders must be ONE global order
    order = {k: [torch.argsort(stats[k][i], descending=True) for i in range(L)] for k in stats}
    # data-free order of the GELU outputs (round 8): E[relu(z)] for z ~ N(mu, sigma^2), mu = b1 + W1 beta2,
    # sigma = |W1 (.) gamma2| (LN output ~ unit variance per channel): needs the weights only
    def wf_order(b):
        W = b.mlp.fc1.weight.detach().double()          # [5120, 1280]
        g2, b2 = b.norm2.weight.detach().double(), b.norm2.bias.detach().double()
        mu = b.mlp.fc1.bias.detach().double() + W @ b2
        sg = (W * g2).norm(dim=1).clamp_min(1e-12)
        a = mu / sg
        nd = torch.distributions.Normal(0.0, 1.0)
        er = sg * torch.exp(-0.5 * a * a) / (2 * torch.pi) ** 0.5 + mu * nd.cdf(a)
        return torch.argsort(er.float(), descending=True)
    order["gw"] = [wf_order(b) for b in blocks]
    # power-of-two per-channel scales of the LN outputs (fold 2^-k into gamma / beta, 2^k into the next weight's K
    # rows: exact in bf16 / bfp8): data-free |h_c| ~ sqrt(gamma_c^2 + beta_c^2), or calibrated mean |h_c|
    def p2(v):
        return torch.pow(2.0, torch.round(torch.log2(v.clamp_min(1e-12))))
    scale = {"ln1": [p2((b.norm1.weight.detach() ** 2 + b.norm1.bias.detach() ** 2).sqrt()) for b in blocks],
             "ln2": [p2((b.norm2.weight.detach() ** 2 + b.norm2.bias.detach() ** 2).sqrt()) for b in blocks]}
    scale_c = {k: [p2(stats[k][i] / stats[k][i].median()) for i in range(L)] for k in ("ln1", "ln2")}
    # rank agreement with the calibrated order (Spearman)
    def spear(o1, o2):
        r1 = torch.empty_like(o1); r1[o1] = torch.arange(len(o1)); r2 = torch.empty_like(o2); r2[o2] = torch.arange(len(o2))
        return float(torch.corrcoef(torch.stack([r1.double(), r2.double()]))[0, 1])
    sp = [spear(order["g"][i], order["gw"][i]) for i in range(L)]
    print(f"[calib] data-free vs calibrated g order, Spearman per layer: min {min(sp):.3f} median {sorted(sp)[L // 2]:.3f}")
    # ONE stream order serves ln1 and ln2 (both are in the residual stream's channel order)
    tot = sum(s / s.mean() for s in stats["ln1"]) + sum(s / s.mean() for s in stats["ln2"])
    g1 = torch.argsort(tot, descending=True)
    glob = {"ln1": g1, "ln2": g1}
    # ratio of the largest to the median channel magnitude (outlier severity)
    for k in stats:
        r = [float(stats[k][i].max() / stats[k][i].median()) for i in range(L)]
        print(f"[calib] {k}: max/median channel |x| per layer: min {min(r):.1f} median {sorted(r)[L // 2]:.1f} max {max(r):.1f}")

    xs = crops()[:n_eval]
    with torch.no_grad():
        base = [ref(x) for _, x in xs]

    def pcc(u, v):
        u = u.reshape(-1).double(); v = v.reshape(-1).double()
        u = u - u.mean(); v = v - v.mean()
        return float((u @ v) / (u.norm() * v.norm()).clamp_min(1e-12))

    def run(cfg: dict, tag: str):
        hs = []
        for i, b in enumerate(blocks):
            for k, mod in (("ln1", b.attn.qkv), ("ln2", b.mlp.fc1), ("g", b.mlp.fc2)):
                mode = cfg.get(k)
                if mode is None:
                    continue
                if mode in ("scale", "scalec"):
                    sc = (scale if mode == "scale" else scale_c)[k][i]
                    hs.append(mod.register_forward_pre_hook(lambda m, inp, sc=sc: (bfp8(inp[0] / sc, None) * sc,)))
                    continue
                perm = None if mode == "id" else (glob[k] if mode == "glob" else (order["gw"][i] if mode == "wf" else order[k][i]))
                hs.append(mod.register_forward_pre_hook(lambda m, inp, perm=perm: (bfp8(inp[0], perm),)))
        with torch.no_grad():
            outs = [ref(x) for _, x in xs]
        for h in hs:
            h.remove()
        p = [pcc(o, r) for o, r in zip(outs, base)]
        cam = [float(o[0, -3] - r[0, -3]) for o, r in zip(outs, base)]
        betas = max(float((o[0, -13:-3] - r[0, -13:-3]).abs().max()) for o, r in zip(outs, base))
        print(f"[{tag:34s}] PCC mean {sum(p) / len(p):.7f} min {min(p):.7f}  cam[0] signed {sum(cam) / len(cam):+.5f}"
              f" abs {sum(abs(c) for c in cam) / len(cam):.5f}  betas_max {betas:.4f}", flush=True)

    if "trunc" in sys.argv:
        run({"ln1": "id"}, "ln1 bfp8")
        run({"ln1": "id", "ln2": "id"}, "ln1+ln2 bfp8")
        return
    if "scale" in sys.argv:
        run({"ln1": "id"}, "ln1 bfp8 (served today)")
        run({"ln1": "scale"}, "ln1 bfp8, p2 scale (gamma/beta)")
        run({"ln1": "scalec"}, "ln1 bfp8, p2 scale (calibrated)")
        run({"ln1": "id", "ln2": "id"}, "ln1+ln2 bfp8")
        run({"ln1": "scale", "ln2": "scale"}, "ln1+ln2 bfp8, p2 scale (gamma/beta)")
        run({"ln1": "scalec", "ln2": "scalec"}, "ln1+ln2 bfp8, p2 scale (calibrated)")
        return
    if "gonly" in sys.argv:
        run({"ln1": "id", "g": "layer"}, "ln1 + g bfp8, per-layer sorted g")
        run({"ln1": "id", "g": "wf"}, "ln1 + g bfp8, data-free sorted g")
        return
    run({"ln1": "id"}, "ln1 bfp8 (served today)")
    run({"ln1": "glob"}, "ln1 bfp8, global sorted order")
    run({"ln1": "id", "ln2": "id"}, "ln1+ln2 bfp8")
    run({"ln1": "glob", "ln2": "glob"}, "ln1+ln2 bfp8, global order (ln2 stats)")
    run({"ln1": "id", "g": "id"}, "ln1 + g bfp8")
    run({"ln1": "id", "g": "layer"}, "ln1 + g bfp8, per-layer sorted g")
    run({"ln1": "glob", "ln2": "glob", "g": "layer"}, "all bfp8, sorted")


if __name__ == "__main__":
    main()