# SPDX-License-Identifier: Apache-2.0
"""Error attribution ViT vs MANO head on the 32 acc.py crops: the device ViT feature map fed to the CPU fp32
head, and the device full model, both vs the CPU fp32 reference; prints the signed camera-scale error of each."""
import os, sys, statistics
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
sys.path.insert(0, str(Path(__file__).resolve().parent))
os.environ.setdefault("TT_FUSED_HASH", "0")
import torch, ttnn
from hamer import Hamer, build_from_reference, load_hamer_checkpoint
from hamer import vit as tv
from acc import crops

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()
tt = build_from_reference(ref, device_id=0, strict=True)
rows = []
with torch.no_grad():
    for n, x in xs:
        fr = ref.backbone(x)                                         # [1, 1280, 16, 12]
        o_ref = ref.head(fr)
        tok = tt._upload_tokens(x)
        f = tv.forward(tok, tt._vit_params, fused=True)
        ft = ttnn.to_torch(f).float().reshape(1, 16, 12, 1280).permute(0, 3, 1, 2)
        ttnn.deallocate(f)
        o_vit = ref.head(ft)                                         # device ViT + CPU head
        o_tt = tt(x)
        fe = (ft - fr)
        rows.append((float(o_vit[0, 154] - o_ref[0, 154]), float(o_tt[0, 154] - o_ref[0, 154]),
                     float((fe * fr.sign()).mean() / fr.abs().mean())))
tt.close()
for i, name in enumerate(("devViT+cpuHead", "device full", "feat mag-bias")):
    v = [r[i] for r in rows]
    print(f"{name:16s} mean {statistics.mean(v):+.5f}  pos {sum(t > 0 for t in v)}/{len(v)}")