# SPDX-License-Identifier: Apache-2.0
"""Host-only (no device) sensitivity check of the ViT weight-precision change.

Quantizes ViT weights to bfp8 on the HOST (ttnn.from_torch(..., bfloat8_b, TILE) -> to_torch, no
device opened) and runs the fp32 torch CPU model on the same 32 crops as acc.py, PCC vs the
unquantized fp32 reference.  It isolates the effect of the declared weight-precision change:
  baseline config : fc1 weights bfp8
  optimized config: qkv + proj + fc1 + fc2 weights bfp8
It does NOT measure device numerics (bf16 activations, LN/SDPA/gelu kernels, fp32 vs bf16 DST);
those need acc.py on a healthy chip.  Run with TT_VISIBLE_DEVICES=none.
"""
import copy, json, sys
from pathlib import Path

_REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO)); sys.path.insert(0, str(_REPO / "scripts"))
import torch
import ttnn
from acc import crops
from hamer import Hamer, load_hamer_checkpoint


def q8(w):
    return ttnn.to_torch(ttnn.from_torch(w.detach().float(), dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT)).float()


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 quantized(ref, names):
    m = copy.deepcopy(ref)
    for blk in m.backbone.blocks if hasattr(m, "backbone") else m.vit.blocks:
        mods = {"qkv": blk.attn.qkv, "proj": blk.attn.proj, "fc1": blk.mlp.fc1, "fc2": blk.mlp.fc2}
        for n in names:
            # torch Linear weight is [N, K]; device layout is [K, N] -> quantize in that orientation
            mods[n].weight.data = q8(mods[n].weight.data.t().contiguous()).t().contiguous()
    return m


def main():
    hf = Path.home() / ".cache/huggingface/hub/models--changh95--hamer-weights/snapshots/e8217b46d54b4e8afaaf0dd6efb6b62341c97d32"
    torch.manual_seed(0)
    ref = Hamer().eval()
    load_hamer_checkpoint(ref, hf / "hamer.ckpt", hf / "mano_mean_params.npz")
    xs = crops()
    cfgs = {"baseline(fc1)": ["fc1"], "optimized(qkv,proj,fc1,fc2)": ["qkv", "proj", "fc1", "fc2"]}
    with torch.no_grad():
        refs = [ref(x) for _, x in xs]
        for tag, names in cfgs.items():
            m = quantized(ref, names)
            p = [pcc(m(x), r) for (_, x), r in zip(xs, refs)]
            print("HOSTQ", json.dumps({"cfg": tag, "n": len(p), "mean": sum(p) / len(p), "min": min(p),
                                       "below_0.9999": sum(v < 0.9999 for v in p)}), flush=True)


if __name__ == "__main__":
    main()