"""Build aligned 112x112 face dataset with teacher (buffalo_l) embeddings.
Sources: labeled jpg faces (from clustering) + video frames (auto-labeled by centroid)."""
import os, sys, json, re, numpy as np, cv2, collections
from PIL import Image
from insightface.app import FaceAnalysis
from insightface.utils.face_align import norm_crop
D=os.environ.get("FACE_DATA","data")  # flat folder of jpg/mp4, filenames YYYYMMDD_*; OUT="dataset"
labels=json.load(open("analysis/labels.json")); names=sorted(set(labels.values()))+["other"]
rows=json.load(open("analysis/faces_clustered.json"))
app=FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"]); app.prepare(ctx_id=-1, det_size=(320,320))
def date_of(f):
    m=re.match(r"(\d{8})_",f); return m.group(1) if m else "00000000"
imgs,embs,labs,srcs,dates=[],[],[],[],[]
def add(face,img,lab,src,date):
    crop=norm_crop(img,face.kps,112); imgs.append(crop); embs.append(face.normed_embedding.astype(np.float32)); labs.append(lab); srcs.append(src); dates.append(date)
# ---- jpg faces ----
byfile=collections.defaultdict(list)
for r in rows: byfile[r["file"]].append(r)
other_cap=400; other_n=0
for i,(f,rs) in enumerate(sorted(byfile.items())):
    im=Image.open(os.path.join(D,f)); im.draft("RGB",(1280,1280)); im=im.convert("RGB")
    W,H=rs[0]["w"],rs[0]["h"]; s0=1280/max(W,H); s0=min(s0,1)        # scale used in scan
    s1=im.width/W                                                      # scale after draft
    img=cv2.cvtColor(np.array(im),cv2.COLOR_RGB2BGR)
    for r in rs:
        lab=labels.get(str(r["cluster"]),"other")
        if lab=="other":
            if other_n>=other_cap: continue
        x1,y1,x2,y2=[v/s0*s1 for v in r["box"]]; m=0.6*max(x2-x1,y2-y1)
        X1,Y1,X2,Y2=int(max(0,x1-m)),int(max(0,y1-m)),int(min(img.shape[1],x2+m)),int(min(img.shape[0],y2+m))
        sub=img[Y1:Y2,X1:X2]
        if sub.size==0: continue
        fs=app.get(sub)
        if not fs: continue
        fc=max(fs,key=lambda q:(q.bbox[2]-q.bbox[0])*(q.bbox[3]-q.bbox[1]))
        add(fc,sub,lab,"jpg:"+f,date_of(f))
        if lab=="other": other_n+=1
    if i%50==0: print("jpg",i,len(byfile),len(imgs),flush=True)
print("jpg done",len(imgs),collections.Counter(labs),flush=True)
# ---- centroids for auto-label ----
E=np.array(embs); L=np.array(labs)
cent={n:E[L==n].mean(0) for n in names if n!="other" and (L==n).any()}
for n in cent: cent[n]/=np.linalg.norm(cent[n])
C=np.stack([cent[n] for n in cent]); cn=list(cent)
# ---- video frames @1fps ----
vid_counts=collections.Counter()
for f in sorted(x for x in os.listdir(D) if x.lower().endswith(".mp4")):
    cap=cv2.VideoCapture(os.path.join(D,f)); fps=cap.get(5) or 30; n=int(cap.get(7)); step=int(round(fps))
    for fi in range(0,n,step):
        cap.set(1,fi); ok,img=cap.read()
        if not ok: continue
        h,w=img.shape[:2]; s=960/max(h,w)
        if s<1: img=cv2.resize(img,None,fx=s,fy=s)
        for fc in app.get(img):
            if max(fc.bbox[2]-fc.bbox[0],fc.bbox[3]-fc.bbox[1])<40: continue
            sim=C@fc.normed_embedding; j=int(sim.argmax())
            if sim[j]>=0.45: lab=cn[j]
            elif sim[j]<0.30: lab="other"
            else: continue   # ambiguous, skip
            add(fc,img,lab,f"vid:{f}:{fi}",date_of(f)); vid_counts[lab]+=1
    print("vid",f,dict(vid_counts),flush=True)
np.savez_compressed(os.path.join(OUT,"faces.npz"),imgs=np.array(imgs,np.uint8),embs=np.array(embs,np.float32),labels=np.array(labs),src=np.array(srcs),date=np.array(dates))
print("DONE",len(imgs),collections.Counter(labs))