"""Webcam test: detect+align with buffalo_l detector, embed with TinyFaceNet, match to centroids.
Low-light handling: auto-gamma + CLAHE on the frame before detection/crop; per-track majority vote over last VOTE frames.
usage: python live_test.py [ckpt] [--raw]   (--raw = no normalization, for A/B)
q = quit, s = save snapshot to analysis/live_N.jpg, n = toggle normalization"""
import sys, time, collections, numpy as np, cv2, torch, torch.nn.functional as F
from insightface.app import FaceAnalysis
from insightface.utils.face_align import norm_crop
from model import TinyFaceNet
args=[a for a in sys.argv[1:] if not a.startswith("--")]
CK=args[0] if args else "dataset/best.pt"; THR=0.55; VOTE=7; NORM="--raw" not in sys.argv
ck=torch.load(CK,weights_only=False); names=ck["names"]; C=ck["centroids"]
m=TinyFaceNet(n_cls=len(names)); m.load_state_dict(ck["model"]); m.eval()
app=FaceAnalysis(name="buffalo_l",allowed_modules=["detection"],providers=["CPUExecutionProvider"]); app.prepare(ctx_id=-1,det_size=(320,320))
clahe=cv2.createCLAHE(clipLimit=2.0,tileGridSize=(8,8))
def normalize(img):
    """auto-gamma to lift dark frames + CLAHE on luminance. Returns (img, mean_brightness_before)."""
    mean=img.mean()
    if mean<110:                       # dark frame -> gamma lift towards mean ~110
        gamma=np.log(110/255)/np.log(max(mean,1)/255); lut=(np.linspace(0,1,256)**gamma*255).astype(np.uint8); img=cv2.LUT(img,lut)
    lab=cv2.cvtColor(img,cv2.COLOR_BGR2LAB); lab[...,0]=clahe.apply(lab[...,0])
    return cv2.cvtColor(lab,cv2.COLOR_LAB2BGR),mean
def iou(a,b):
    ix=max(0,min(a[2],b[2])-max(a[0],b[0])); iy=max(0,min(a[3],b[3])-max(a[1],b[1])); i=ix*iy
    return i/((a[2]-a[0])*(a[3]-a[1])+(b[2]-b[0])*(b[3]-b[1])-i+1e-6)
tracks=[]   # [bbox, deque(labels), deque(scores), last_seen_frame]
cap=cv2.VideoCapture(0,cv2.CAP_DSHOW); cap.set(3,640); cap.set(4,480)
if not cap.isOpened(): sys.exit("webcam open nahi hua")
n=0; t0=time.time(); fr=0
while True:
    ok,raw=cap.read()
    if not ok: break
    img,mean=normalize(raw) if NORM else (raw,raw.mean())
    for f in app.get(img):
        crop=norm_crop(img,f.kps,112); x=(torch.from_numpy(crop).permute(2,0,1)[None].float()/255.-0.5)/0.5
        with torch.no_grad(): e=F.normalize(m(x),dim=1)[0]
        sim=(C@e); j=int(sim.argmax()); s=float(sim[j]); inst=names[j] if s>=THR else "unknown"
        # --- track + vote ---
        bb=tuple(map(int,f.bbox)); tr=max(tracks,key=lambda t:iou(t[0],bb),default=None)
        if tr is None or iou(tr[0],bb)<0.3: tr=[bb,collections.deque(maxlen=VOTE),collections.deque(maxlen=VOTE),fr]; tracks.append(tr)
        tr[0]=bb; tr[1].append(inst); tr[2].append(s); tr[3]=fr
        lab,cnt=collections.Counter(tr[1]).most_common(1)[0]; conf=np.mean(tr[2])
        col=(0,255,0) if lab!="unknown" else (0,0,255)
        x1,y1,x2,y2=bb; cv2.rectangle(img,(x1,y1),(x2,y2),col,2)
        top3=" ".join(f"{names[k]}:{sim[k]:.2f}" for k in sim.argsort(descending=True)[:3])
        cv2.putText(img,f"{lab} {conf:.2f} ({cnt}/{len(tr[1])})",(x1,y1-8),0,0.7,col,2); cv2.putText(img,top3,(x1,y2+18),0,0.45,(255,255,0),1)
    tracks=[t for t in tracks if fr-t[3]<15]
    fr+=1; cv2.putText(img,f"{fr/(time.time()-t0):.1f}fps norm={'ON' if NORM else 'off'} light={mean:.0f}  q=quit s=save n=norm",(8,20),0,0.5,(255,255,255),1)
    cv2.imshow("live",img); k=cv2.waitKey(1)&0xFF
    if k==ord("q"): break
    if k==ord("n"): NORM=not NORM; tracks=[]
    if k==ord("s"): n+=1; cv2.imwrite(f"analysis/live_{n}.jpg",img); print("saved",n,flush=True)
cap.release(); cv2.destroyAllWindows()