"""LFW verification benchmark for TinyFaceNet (standard 10-fold protocol, 6000 pairs).
Images: funneled LFW (sklearn download, pairs.txt parsed directly, images read one at a time). Faces are re-detected + 5-pt aligned with buffalo_l -> 112x112,
so the input matches training. Metric: 10-fold cross-validated accuracy with threshold picked on the other 9 folds (the LFW standard),
plus best-threshold accuracy, AUC, and TAR@FAR=1e-2/1e-3.
Stranger test (--strangers, needs a .pt ckpt with centroids + dataset/faces_v2.npz): treats all 12k LFW faces as people outside the family and reports,
per open-set threshold, how many get accepted as a family member vs how many held-out family faces are still recognised.
usage: python eval_lfw.py [ckpt=dataset/model_v1.pt | *.onnx] [--teacher] [--strangers]   (--teacher = also score buffalo_l itself as reference)
Cache: analysis/lfw_aligned.npz (aligned crops, built once, ~10 min CPU). Output: analysis/lfw_<tag>.json"""
import sys, os, json, time, numpy as np, cv2, torch, torch.nn.functional as F
import onnxruntime as ort
_IS=ort.InferenceSession
class _CappedSession(_IS):   # cap ORT threads BEFORE insightface imports; default = all cores spin-waiting on 5 tiny models -> thrash
    def __init__(s,path,sess_options=None,**kw):
        so=ort.SessionOptions(); so.intra_op_num_threads=int(os.environ.get("ORT_THREADS","4")); so.inter_op_num_threads=1
        super().__init__(path,sess_options=so,**kw)
ort.InferenceSession=_CappedSession
torch.set_num_threads(int(os.environ.get("THREADS","10")))
args=[a for a in sys.argv[1:] if not a.startswith("--")]; CK=args[0] if args else "dataset/model_v1.pt"; TAG=os.path.splitext(os.path.basename(CK))[0]
CACHE="analysis/lfw_aligned.npz"
# ---------------- 1. aligned crops (cached) ----------------
if not os.path.exists(CACHE):
    from sklearn.datasets import get_data_home
    from insightface.app import FaceAnalysis
    from insightface.utils.face_align import norm_crop
    home=os.path.join(get_data_home(),"lfw_home")
    if not os.path.exists(os.path.join(home,"pairs.txt")):   # first time: let sklearn fetch the tar + pairs.txt (does not load arrays)
        from sklearn.datasets._lfw import _check_fetch_lfw; _check_fetch_lfw(data_home=get_data_home(),funneled=True,download_if_missing=True)
    IMG=os.path.join(home,"lfw_funneled"); pairs=[]
    for l in open(os.path.join(home,"pairs.txt")).read().splitlines()[1:]:
        s=l.split("	")
        if len(s)==3: pairs.append(((s[0],int(s[1])),(s[0],int(s[2])),1))
        elif len(s)==4: pairs.append(((s[0],int(s[1])),(s[2],int(s[3])),0))
    N=len(pairs); Y=np.array([p[2] for p in pairs],np.int8); print("pairs",N,"pos",int(Y.sum()),flush=True)
    def load(name,idx): return cv2.imread(os.path.join(IMG,name,f"{name}_{idx:04d}.jpg"))   # BGR 250x250
    P=None
    app=FaceAnalysis(name="buffalo_l",providers=["CPUExecutionProvider"],allowed_modules=["detection","recognition"]); app.prepare(ctx_id=-1,det_size=(320,320))
    crops=np.zeros((N,2,112,112,3),np.uint8); ok=np.ones((N,2),bool); temb=np.zeros((N,2,512),np.float32); t0=time.time()
    for i in range(N):
        for k in range(2):
            img=load(*pairs[i][k]); faces=app.get(img)
            if not faces:   # fallback: center crop (LFW faces are centered)
                ok[i,k]=False; crops[i,k]=cv2.cvtColor(cv2.resize(img[70:180,70:180],(112,112)),cv2.COLOR_BGR2RGB); continue
            c=np.array([125,125]); f=min(faces,key=lambda f:np.linalg.norm((f.bbox[:2]+f.bbox[2:])/2-c))   # most central face
            crops[i,k]=cv2.cvtColor(norm_crop(img,f.kps,112),cv2.COLOR_BGR2RGB); temb[i,k]=f.normed_embedding
        if i%200==0: print(f"  aligned {i}/{N}  {time.time()-t0:.0f}s  det_fail={int((~ok[:i+1]).sum())}",flush=True)
    np.savez_compressed(CACHE,crops=crops,y=Y,ok=ok,teacher=temb); print("cached",CACHE,"det_fail",int((~ok).sum()))
d=np.load(CACHE); X=d["crops"]; Y=d["y"].astype(bool); N=len(Y)
# ---------------- 2. embed ----------------
def prep(x): return ((x.transpose(0,3,1,2).astype(np.float32)/255.)-0.5)/0.5
def embed_all(fn):
    flat=X.reshape(-1,112,112,3); out=[]
    for i in range(0,len(flat),128): out.append(fn(prep(flat[i:i+128])))
    E=np.concatenate(out); E/=np.linalg.norm(E,axis=1,keepdims=True); return E.reshape(N,2,-1)
t0=time.time()
if CK.endswith(".onnx"):
    import onnxruntime as ort; s=ort.InferenceSession(CK,providers=["CPUExecutionProvider"]); n=s.get_inputs()[0].name
    E=embed_all(lambda x:s.run(None,{n:x})[0])
else:
    from model import TinyFaceNet
    ck=torch.load(CK,weights_only=False); m=TinyFaceNet(); m.load_state_dict(ck["model"],strict=False); m.eval()
    def run(x):
        with torch.no_grad(): return m(torch.from_numpy(x)).numpy()
    E=embed_all(run)
lat=(time.time()-t0)/(2*N)*1000; print(f"embedded {2*N} faces in {time.time()-t0:.0f}s ({lat:.1f} ms/face)")
# ---------------- 3. metrics ----------------
def kfold_acc(sim,y,k=10):
    folds=np.array_split(np.arange(len(y)),k); ths=np.arange(-1,1,0.002); accs=[]
    for f in folds:
        tr=np.setdiff1d(np.arange(len(y)),f); a=[((sim[tr]>=t)==y[tr]).mean() for t in ths]; t=ths[int(np.argmax(a))]
        accs.append(((sim[f]>=t)==y[f]).mean())
    return float(np.mean(accs)),float(np.std(accs))
def tar_at_far(sim,y,far):
    neg=np.sort(sim[~y])[::-1]; t=neg[int(far*len(neg))]; return float((sim[y]>=t).mean()),float(t)
def report(E,name):
    sim=(E[:,0]*E[:,1]).sum(1); acc,sd=kfold_acc(sim,Y)
    ths=np.arange(-1,1,0.002); best=max(((sim>=t)==Y).mean() for t in ths)
    from sklearn.metrics import roc_auc_score
    r=dict(acc_10fold=round(acc,4),std=round(sd,4),acc_best_thr=round(float(best),4),auc=round(float(roc_auc_score(Y,sim)),4),
           tar_far1e2=round(tar_at_far(sim,Y,1e-2)[0],4),tar_far1e3=round(tar_at_far(sim,Y,1e-3)[0],4),
           mean_sim_pos=round(float(sim[Y].mean()),3),mean_sim_neg=round(float(sim[~Y].mean()),3))
    print(f"{name:12s} LFW 10-fold acc={acc:.4f}+-{sd:.4f}  best={best:.4f}  AUC={r['auc']:.4f}  TAR@FAR1e-2={r['tar_far1e2']:.3f}  TAR@FAR1e-3={r['tar_far1e3']:.3f}  pos/neg sim={r['mean_sim_pos']}/{r['mean_sim_neg']}")
    return r
res={"model":CK,"n_pairs":int(N),"det_fail":int((~d["ok"]).sum()),"lat_ms":round(lat,2),"student":report(E,TAG)}
if "--strangers" in sys.argv and not CK.endswith(".onnx") and "centroids" in ck:
    names=list(map(str,ck["names"])); C=ck["centroids"].numpy(); fam_idx=[i for i in range(len(names)) if names[i]!="other"]; Cf=C[fam_idx]
    f=np.load("dataset/faces_v2.npz"); lab=np.array(list(map(str,f["labels"]))); DT=f["date"]
    fam=np.array([(x[:6] in ("202511","202605","202606")) for x in DT])&(lab!="other")
    with torch.no_grad(): EF=F.normalize(m(torch.from_numpy(prep(f["imgs"][fam]))),dim=1).numpy()
    yf=np.array([names.index(l) for l in lab[fam]]); EL=E.reshape(-1,E.shape[-1])
    sL=(EL@Cf.T).max(1); sF=EF@Cf.T; predF=np.array(fam_idx)[sF.argmax(1)]; bestF=sF.max(1)
    print(f"strangers n={len(EL)} max-sim to family mean={sL.mean():.3f} median={np.median(sL):.3f}; family test faces n={len(EF)}")
    st={}
    for thr in [0.45,0.50,0.55,0.60,0.65,0.70,0.75,0.80]:
        fa=float((sL>=thr).mean()); rec=float(((bestF>=thr)&(predF==yf)).mean()); st[str(thr)]=dict(stranger_accept=round(fa,4),family_recall=round(rec,4))
        print(f"  THR={thr:.2f}  stranger accepted as family {fa*100:5.1f}%   family recognised {rec*100:5.1f}%")
    res["stranger_test"]=dict(n_strangers=int(len(EL)),n_family=int(len(EF)),stranger_max_sim_mean=round(float(sL.mean()),4),by_thr=st)
if "--teacher" in sys.argv: res["teacher_buffalo_l"]=report(d["teacher"]/np.maximum(np.linalg.norm(d["teacher"],axis=2,keepdims=True),1e-8),"buffalo_l")
json.dump(res,open(f"analysis/lfw_{TAG}.json","w"),indent=1); print("wrote",f"analysis/lfw_{TAG}.json")