"""Quantize a TinyFaceNet checkpoint and measure accuracy drop.
Variants: fp32 | int8_ort (ONNX Runtime static, QOperator, per-channel -- the real deployable path)
          | int8_w (weight-only fake-quant, per-channel symmetric) | int4_w (weight-only fake-quant, per-channel) | int4_g32 (per-group of 32)
Metrics (date-held-out test set): family acc (centroid argmax), open-set acc with THR, mean cosine agreement with fp32 embeddings over the FULL dataset,
and argmax agreement with fp32 over the full dataset (statistically more stable than the 59-face test set).
usage: python quantize_eval.py <ckpt> <tag>   -> analysis/quant_<tag>.json, dataset/<tag>_fp32.onnx, dataset/<tag>_int8.onnx"""
import sys, os, json, time, copy, numpy as np, torch, torch.nn.functional as F, onnx, onnxruntime as ort
from onnxruntime.quantization import quantize_static, CalibrationDataReader, QuantType, QuantFormat
from model import TinyFaceNet
torch.set_num_threads(4); CK,TAG=sys.argv[1],sys.argv[2]; THR=0.55
ck=torch.load(CK,weights_only=False); names=list(map(str,ck["names"])); C=ck["centroids"].numpy(); oi=names.index("other")
d=np.load("dataset/faces_v2.npz"); X=d["imgs"]; L=np.array(list(map(str,d["labels"]))); DT=d["date"]; S=d["src"]
Y=np.array([names.index(l) for l in L]); test=np.array([(x[:6] in ("202511","202605","202606")) for x in DT])
full=np.linspace(0,len(X)-1,1500).astype(int)                      # agreement subset (1500 faces, spread over sources)
def prep(x): return ((x.transpose(0,3,1,2).astype(np.float32)/255.)-0.5)/0.5
def base(): m=TinyFaceNet(n_cls=len(names)); m.load_state_dict(ck["model"]); m.eval(); return m
def run_torch(m,ix):
    out=[]
    with torch.no_grad():
        for i in range(0,len(ix),256): out.append(F.normalize(m(torch.from_numpy(prep(X[ix[i:i+256]]))),dim=1).numpy())
    return np.concatenate(out)
def run_ort(path,ix):
    s=ort.InferenceSession(path,providers=["CPUExecutionProvider"]); n=s.get_inputs()[0].name; out=[]
    for i in range(0,len(ix),64): e=s.run(None,{n:prep(X[ix[i:i+64]])})[0]; out.append(e/np.linalg.norm(e,axis=1,keepdims=True))
    return np.concatenate(out)
def metrics(E,Eref,ix):
    sim=E[np.isin(ix,np.nonzero(test)[0])]@C.T if False else None
    te=np.nonzero(test)[0]; Et=E[:len(te)]; yt=Y[te]; sim=Et@C.T; pred=sim.argmax(1); best=sim.max(1)
    fam=yt!=oi; acc=float((pred[fam]==yt[fam]).mean())
    pred_os=np.where(best>=THR,pred,oi); acc_os=float((pred_os==yt).mean())                   # open-set: unknown->other
    per={names[c]:round(float((pred[yt==c]==c).mean()),3) for c in range(len(names)) if (yt==c).any()}
    Ef=E[len(te):]; Rf=Eref[len(te):]; cos=float((Ef*Rf).sum(1).mean()); agree=float(((Ef@C.T).argmax(1)==(Rf@C.T).argmax(1)).mean())
    return dict(fam_acc=acc,openset_acc=acc_os,per_class=per,cos_vs_fp32=cos,argmax_agree_vs_fp32=agree)
def fake_quant(m,bits,group=None):
    m=copy.deepcopy(m); qmax=2**(bits-1)-1
    for name,p in m.named_parameters():
        if p.dim()<2 or name.startswith(("distill","cls")): continue
        w=p.data; o=w.shape[0]; wf=w.reshape(o,-1)
        if group and wf.shape[1]%group==0: wg=wf.reshape(o,-1,group); sc=wg.abs().amax(2,keepdim=True).clamp_min(1e-8)/qmax; wf=((wg/sc).round().clamp(-qmax,qmax)*sc).reshape(o,-1)
        else: sc=wf.abs().amax(1,keepdim=True).clamp_min(1e-8)/qmax; wf=(wf/sc).round().clamp(-qmax,qmax)*sc
        p.data=wf.reshape(w.shape)
    return m
def lat(fn,n=50):
    x=np.zeros((1,112,112,3),np.uint8); fn(np.array([0])); t=time.perf_counter()
    for _ in range(n): fn(np.array([0]))
    return (time.perf_counter()-t)/n*1000
ix=np.concatenate([np.nonzero(test)[0],full]); res={}; m=base()
# ---- fp32 torch ----
E0=run_torch(m,ix); res["fp32"]=dict(**metrics(E0,E0,ix),size_kb=round(sum(p.numel() for n,p in m.named_parameters() if not n.startswith(("distill","cls")))*4/1024),lat_ms=round(lat(lambda i:run_torch(m,i)),2),note="PyTorch fp32 (deploy params only)")
# ---- ONNX export + ORT fp32 ----
fp=f"dataset/{TAG}_fp32.onnx"; q8=f"dataset/{TAG}_int8.onnx"
torch.onnx.export(m,torch.zeros(1,3,112,112),fp,input_names=["x"],output_names=["emb"],dynamic_axes={"x":{0:"b"},"emb":{0:"b"}},opset_version=17,dynamo=False)
E=run_ort(fp,ix); res["fp32_onnx"]=dict(**metrics(E,E0,ix),size_kb=round(os.path.getsize(fp)/1024),lat_ms=round(lat(lambda i:run_ort(fp,i)),2),note="ONNX Runtime fp32")
# ---- ORT static int8 ----
class CR(CalibrationDataReader):
    def __init__(s): tr=np.nonzero(~test)[0]; s.ix=np.linspace(0,len(tr)-1,200).astype(int); s.tr=tr; s.i=0
    def get_next(s):
        if s.i>=len(s.ix): return None
        b=s.tr[s.ix[s.i:s.i+8]]; s.i+=8; return {"x":prep(X[b])}
quantize_static(fp,q8,CR(),quant_format=QuantFormat.QOperator,per_channel=True,activation_type=QuantType.QUInt8,weight_type=QuantType.QInt8)
E=run_ort(q8,ix); res["int8_ort"]=dict(**metrics(E,E0,ix),size_kb=round(os.path.getsize(q8)/1024),lat_ms=round(lat(lambda i:run_ort(q8,i)),2),note="ONNX Runtime static int8, per-channel, 200 calib imgs (deployable)")
# ---- weight-only fake quant ----
for tag,bits,g,note in [("int8_w",8,None,"simulated: int8 weights per-channel, fp32 activations"),
                        ("int4_w",4,None,"simulated: int4 weights per-channel (coarse)"),
                        ("int4_g32",4,32,"simulated: int4 weights, group=32 scales (llama.cpp-style)")]:
    mq=fake_quant(m,bits,g); E=run_torch(mq,ix); core=sum(p.numel() for n,p in m.named_parameters() if not n.startswith(("distill","cls")))
    kb=core*bits/8/1024*(1+ (2/g if g else 0.002))
    res[tag]=dict(**metrics(E,E0,ix),size_kb=round(kb),lat_ms=None,note=note)
json.dump(dict(ckpt=CK,tag=TAG,names=names,n_test=int(test.sum()),n_test_family=int((test&(Y!=oi)).sum()),n_agree=len(full),thr=THR,results=res),open(f"analysis/quant_{TAG}.json","w"),indent=1)
for k,v in res.items(): print(f"{k:10s} fam_acc={v['fam_acc']:.3f} openset={v['openset_acc']:.3f} cos={v['cos_vs_fp32']:.4f} agree={v['argmax_agree_vs_fp32']:.3f} size={v['size_kb']}KB lat={v['lat_ms']}ms")