# SPDX-License-Identifier: Apache-2.0
"""Round 9: robustness sweep -- device vs CPU fp32 reference on out-of-distribution versions of the 8 annotated
hands (dark, bright, low contrast, gaussian noise, salt-and-pepper, blocky) plus pure random uint8 crops (what the
server warm-up gate uses).  Same served TtHamer path as acc.py; model knobs from the environment.
usage: robust.py [--ref-cache PATH]    prints per-group mean / min PCC and ``ROBUST {json}``"""
from __future__ import annotations
import argparse, json, os, sys
from pathlib import Path
_REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO))


def inputs():
    import torch
    from hamer.preprocess import preprocess_pil, normalize_raw
    from PIL import Image
    media = _REPO.parent / "media"
    boxes = json.loads((media / "bboxes.json").read_text())
    g = torch.Generator().manual_seed(7)
    out = []
    for name in sorted(boxes):
        img = Image.open(media / name).convert("RGB")
        for hi, e in enumerate(boxes[name]):
            u, _ = preprocess_pil(img, tuple(e["bbox"]), is_right=bool(e.get("is_right", 1)), raw=True)
            f = u.float()
            var = {
                "dark": f * 0.25,
                "bright": f * 1.8,
                "lowcon": 128 + (f - 128) * 0.3,
                "noise25": f + 25 * torch.randn(f.shape, generator=g),
                "saltpep": torch.where(torch.rand(f.shape, generator=g) < 0.05,
                                       255 * (torch.rand(f.shape, generator=g) < 0.5).float(), f),
                "blocky": torch.nn.functional.interpolate(torch.nn.functional.avg_pool2d(f, 16), scale_factor=16),
            }
            for k, v in var.items():
                out.append((f"{k}:{name}#{hi}", v.round().clamp(0, 255).to(torch.uint8)))
    for i in range(8):
        out.append((f"random:{i}", torch.randint(0, 256, (1, 3, 256, 192), generator=g, dtype=torch.uint8)))
    return [(n, u, normalize_raw(u)) for n, u in out]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ref-cache", default=str(Path(os.environ.get("TMPDIR", "/tmp")) / "hamer_robust_ref.pt"))
    a = ap.parse_args()
    os.environ.setdefault("TT_FUSED_HASH", "0")
    import torch
    from hamer import Hamer, build_from_reference, load_hamer_checkpoint
    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")
    xs = inputs()
    cache = Path(a.ref_cache)
    refs = None
    if cache.exists():
        d = torch.load(cache)
        if d.get("names") == [n for n, _, _ in xs]:
            refs = d["out"]
    if refs is None:
        with torch.no_grad():
            refs = [ref(x) for _, _, x in xs]
        torch.save({"names": [n for n, _, _ in xs], "out": refs}, cache)

    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))

    tt = build_from_reference(ref, device_id=0, strict=True)
    raw = getattr(tt.fused_opts, "input", "norm") == "u8"
    try:
        with torch.no_grad():
            tt(xs[0][1] if raw else xs[0][2]); tt(xs[0][1] if raw else xs[0][2])
            res = {}
            for (n, u, x), r in zip(xs, refs):
                o = tt(u if raw else x)
                res.setdefault(n.split(":")[0], []).append(pcc(o, r))
    finally:
        tt.close()
    summ = {k: {"mean": sum(v) / len(v), "min": min(v)} for k, v in res.items()}
    for k, v in summ.items():
        print(f"[robust] {k:8s} mean {v['mean']:.6f} min {v['min']:.6f}")
    print("ROBUST " + json.dumps(summ))


if __name__ == "__main__":
    main()