import sys,os,glob,json,time,numpy as np,cv2,random
sys.path.insert(0,'code'); sys.path.insert(0,os.path.dirname(__file__)); os.environ['YOLO_OFFLINE']='1'
from gallery import ProductGallery
from detector import ShelfProductDetector
from diag import crops, sku_of
from PIL import Image
S=os.path.dirname(__file__); random.seed(0)
g=ProductGallery('stock-lopia-clean',['sku_l01','sku_l02','sku_l03','sku_l05','sku_l06','sku_l09','sku_l11']); assert g.load()
refs=[(sku_of(os.path.basename(p)[:-4]),g.embed(cv2.imread(p))) for p in sorted(glob.glob('stock-lopia-clean/*.jpg'))]
det=ShelfProductDetector(); dis=[]
for p in sorted(glob.glob('add/*')):
    if 'image6' in p: continue
    im=cv2.cvtColor(np.array(Image.open(p).convert('RGB')),cv2.COLOR_RGB2BGR)
    for j,b in enumerate(det.detect(im).get('boxes',[])):
        x1,y1,x2,y2=[int(v) for v in b['box']]; c=im[max(0,y1):y2,max(0,x1):x2]
        if min(c.shape[:2])>=40: dis.append((f'd{j}_{os.path.basename(p)}',c))
random.shuffle(dis)
D=[(k,g.embed(c)) for k,c in dis]
L=json.load(open(f'{S}/../../../../eval/labels/image6.json')); E={k for k,_ in refs}
Q=[(fc['label'],g.embed(c)) for fc,c in crops('add/image6.jpg',L,True) if fc['label'] in E]
t=time.perf_counter(); _=g.embed(Q and dis[0][1]); print('embed ms',round((time.perf_counter()-t)*1000,1))
for N in [0,50,200,len(D)]:
    R=refs+D[:N]; M=np.stack([e for _,e in R]); names=[k for k,_ in R]
    rec={1:0,5:0,10:0,20:0}
    for truth,q in Q:
        s=M@q; order=[]
        for i in np.argsort(-s):
            if names[i] not in order: order.append(names[i])
        for K in rec: rec[K]+= truth in order[:K]
    print(f'skus={len(set(names)):4d} DINO recall@1/5/10/20 '+'/'.join(str(rec[K]) for K in rec)+f' of {len(Q)}')