import os, numpy as np, torch, torch.nn as nn, torch.nn.functional as F, time, json, collections, sys
from model import TinyFaceNet
torch.manual_seed(0); np.random.seed(0); torch.set_num_threads(int(os.environ.get("THREADS","10")))  # leave cores for the desktop
import os; d=np.load(os.environ.get("DATA","dataset/faces.npz")); X=d["imgs"]; T=d["embs"]; L=d["labels"]; S=d["src"]; DT=d["date"]
names=sorted(set(L)); idx={n:i for i,n in enumerate(names)}; Y=np.array([idx[l] for l in L])
# date-based split: hold out Nov-2025 and May/Jun-2026
test=np.array([(x[:6] in ("202511","202605","202606")) for x in DT]); train=~test
print("classes",names); print("train",train.sum(),"test",test.sum())
print("train per class",dict(collections.Counter(L[train]))); print("test per class",dict(collections.Counter(L[test])))
EPOCHS=int(sys.argv[1]) if len(sys.argv)>1 else 40; BS=128
def aug(x):  # x: B,112,112,3 uint8 numpy -> float tensor B,3,112,112
    x=torch.from_numpy(x).permute(0,3,1,2).float()/255.
    B=x.shape[0]
    flip=torch.rand(B)<0.5; x[flip]=x[flip].flip(-1)
    # random shift/scale via affine
    th=torch.zeros(B,2,3); sc=1+0.15*(torch.rand(B)*2-1); th[:,0,0]=sc; th[:,1,1]=sc; th[:,:,2]=0.12*(torch.rand(B,2)*2-1)
    x=F.grid_sample(x,F.affine_grid(th,x.shape,align_corners=False),padding_mode="border",align_corners=False)
    # color jitter
    x=x*(1+0.4*(torch.rand(B,1,1,1)-0.5))+0.2*(torch.rand(B,1,1,1)-0.5)
    # low-light sim: gamma darken (p=0.4), contrast crush (p=0.3), sensor noise (p=0.4)
    g=torch.rand(B)<0.4
    if g.any(): x[g]=x[g].clamp(1e-3,1)**(1/(0.3+0.7*torch.rand(int(g.sum()),1,1,1)))*(0.5+0.5*torch.rand(int(g.sum()),1,1,1))
    cc=torch.rand(B)<0.3
    if cc.any(): mu=x[cc].mean((1,2,3),keepdim=True); x[cc]=(x[cc]-mu)*(0.4+0.6*torch.rand(int(cc.sum()),1,1,1))+mu
    nz=torch.rand(B)<0.4
    if nz.any(): x[nz]=x[nz]+torch.randn_like(x[nz])*(0.08*torch.rand(int(nz.sum()),1,1,1))
    gray=torch.rand(B)<0.15; x[gray]=x[gray].mean(1,keepdim=True).expand(-1,3,-1,-1)
    blur=torch.rand(B)<0.2
    if blur.any(): x[blur]=F.avg_pool2d(x[blur],3,1,1)
    # cutout
    for i in torch.nonzero(torch.rand(B)<0.3).flatten():
        cy,cx=np.random.randint(0,112,2); x[i,:,max(0,cy-16):cy+16,max(0,cx-16):cx+16]=0
    return (x.clamp(0,1)-0.5)/0.5
def prep(x): return (torch.from_numpy(x).permute(0,3,1,2).float()/255.-0.5)/0.5
m=TinyFaceNet(n_cls=len(names)); opt=torch.optim.AdamW(m.parameters(),2e-3,weight_decay=5e-4)
tr=np.nonzero(train)[0]; te=np.nonzero(test)[0]; Tt=torch.from_numpy(T)
# ---- cap "other" in train: prefer jpg crops, fill with video frames (spread out) ----
OTHER_CAP=600
oth=tr[L[tr]=="other"]; fam=tr[L[tr]!="other"]
oj=oth[np.char.startswith(S[oth],"jpg:")]; ov=oth[~np.char.startswith(S[oth],"jpg:")]
need=max(0,OTHER_CAP-len(oj)); ov=ov[np.linspace(0,len(ov)-1,min(need,len(ov))).astype(int)] if need else ov[:0]
tr=np.concatenate([fam,oj[:OTHER_CAP],ov])
print("train after other-cap",len(tr),dict(collections.Counter(L[tr])))
# ---- class-balanced sampling (sqrt inverse freq) ----
cnt=collections.Counter(Y[tr]); w=np.array([1/np.sqrt(cnt[y]) for y in Y[tr]]); w/=w.sum()
steps=EPOCHS*(len(tr)//BS+1); sched=torch.optim.lr_scheduler.OneCycleLR(opt,2e-3,total_steps=steps,pct_start=0.1)
def embed(ix):
    m.eval(); out=[]
    with torch.no_grad():
        for i in range(0,len(ix),256): out.append(F.normalize(m(prep(X[ix[i:i+256]])),dim=1))
    return torch.cat(out)
def evaluate():
    Etr=embed(tr); Ete=embed(te); ytr=torch.from_numpy(Y[tr]); yte=torch.from_numpy(Y[te])
    C=torch.stack([F.normalize(Etr[ytr==c].mean(0),dim=0) if (ytr==c).any() else torch.zeros(Etr.shape[1]) for c in range(len(names))])
    sim=Ete@C.T; pred=sim.argmax(1); acc=(pred==yte).float().mean().item()
    per={str(names[c]):round((pred[yte==c]==c).float().mean().item(),3) for c in range(len(names)) if (yte==c).any()}
    fam=yte!=idx["other"]; facc=(pred[fam]==yte[fam]).float().mean().item() if fam.any() else 0.
    return facc,per,C
best=0
for ep in range(EPOCHS):
    m.train(); ep_ix=np.random.choice(tr,size=len(tr),replace=True,p=w); t0=time.time(); tl=0
    for i in range(0,len(ep_ix),BS):
        b=ep_ix[i:i+BS]; x=aug(X[b]); y=torch.from_numpy(Y[b]); t=Tt[b]
        e=m(x); lc=F.cross_entropy(m.cls(e),y,label_smoothing=0.1); ld=1-F.cosine_similarity(m.distill(e),t).mean()
        loss=lc+2.0*ld; opt.zero_grad(); loss.backward(); opt.step(); sched.step(); tl+=loss.item()*len(b)
    acc,per,C=evaluate()
    print(f"ep{ep} loss={tl/len(tr):.3f} fam_acc={acc:.4f} {per} {time.time()-t0:.0f}s",flush=True)
    if acc>=best: best=acc; torch.save({"model":m.state_dict(),"names":names,"centroids":C},os.environ.get("OUT","dataset/best.pt"))
print("BEST",best)