# SPDX-License-Identifier: Apache-2.0
"""Host-only: end-to-end sensitivity of the real-checkpoint output (32 acc.py crops, fp32 CPU) to the
ViT MLP GELU implementation.  PCC vs the exact-erf reference for several approximations."""
import sys, math
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
from acc import crops
from hamer import Hamer, load_hamer_checkpoint

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")
cs = crops()
stats = {"n": 0, "lt-3": 0, "gt3": 0, "lt-4": 0, "lt-5.5": 0, "absmax": 0.0}

class Act(torch.nn.Module):
    def __init__(self, f): super().__init__(); self.f = f
    def forward(self, x):
        if self.f is None:
            stats["n"] += x.numel(); stats["lt-3"] += int((x < -3.1).sum()); stats["gt3"] += int((x > 3.1).sum())
            stats["lt-4"] += int((x < -4).sum()); stats["lt-5.5"] += int((x < -5.5).sum()); stats["absmax"] = max(stats["absmax"], float(x.abs().max()))
            return torch.nn.functional.gelu(x)
        return self.f(x)

def set_act(f):
    for b in ref.backbone.blocks:
        b.mlp.act = Act(f)

def run():
    with torch.no_grad():
        return [ref(x).reshape(-1).double() for _, x in cs]

def pcc(u, v):
    u = u - u.mean(); v = v - v.mean(); return float((u @ v) / (u.norm() * v.norm()))

set_act(None); base = run(); print("fc1-out stats", {k: (v / stats["n"] if k not in ("n", "absmax") else v) for k, v in stats.items()})
erf_gelu = lambda x: torch.nn.functional.gelu(x)
def clamp_tails(lo, hi):
    return lambda x: torch.where(x > hi, x, torch.where(x < lo, torch.zeros_like(x), torch.nn.functional.gelu(x)))
_C8 = [8.040229593e-11, -6.988785106e-09, 2.712364733e-07, -6.289864814e-06, 9.886907355e-05, -1.133757411e-03,
       9.877118282e-03, -6.641025841e-02, 3.989226222e-01]
def poly8(x):
    """hamer/kernels/gelu SFPU recipe in fp32: x * (0.5 + x * P(x^2)), x > 4 -> x, x < -4 -> 0."""
    x = x.float(); u = x * x; acc = torch.zeros_like(x)
    for c in _C8:
        acc = acc * u + c
    g = x * (0.5 + x * acc)
    return torch.where(x > 4, x, torch.where(x < -4, torch.zeros_like(x), g))
variants = {
    "poly8": poly8,
    "poly8-bf16": lambda x: poly8(x).bfloat16().float(),
    "tanh": lambda x: torch.nn.functional.gelu(x, approximate="tanh"),
    "bf16-round": lambda x: torch.nn.functional.gelu(x).bfloat16().float(),
    "tails 3.1": clamp_tails(-3.1, 3.1),
    "tails 4": clamp_tails(-4.0, 4.0),
    "tails -5.5/4": clamp_tails(-5.5, 4.0),
}
extra = sys.argv[1:]
for name, f in variants.items():
    if extra and name not in extra:
        continue
    set_act(f); out = run()
    p = [pcc(a, b) for a, b in zip(out, base)]
    print(f"{name:14s} mean {sum(p)/len(p):.8f} min {min(p):.8f}", flush=True)