"""
How many views per SKU does enrollment need?

The number that decides the capture spec: one enrolled view identifies 57 % of
facings, two lift it to 72 %, three to 75 %. The shipped matcher goes the other
way (33 % -> 5 %) because its raw inlier score is not calibrated to be aggregated
across views, and it cannot even express two views of one SKU.

    python code/bench_identity.py            # expects the eval crops described below
"""
import glob, os, sys, random
sys.path.insert(0, "code")
import cv2, numpy as np, torch, timm
from feature_matcher import ProductFeatureMatcher

# Crops extracted from add/image6.jpg and labelled by shelf position; the labels
# are read off the shelf, not off the matcher being measured.
CROPS = sorted(glob.glob(os.environ.get("DORAYAKI_EVAL_CROPS", "/tmp/dorayaki-verify/eval") + "/*.png"))
items = [(p, os.path.basename(p).replace(".png","").split("_",1)[1]) for p in CROPS]
by_class = {}
for path, label in items:
    by_class.setdefault(label, []).append(path)

matcher = ProductFeatureMatcher(stock_dir="stock-lopia")
model = timm.create_model("vit_small_patch14_dinov2.lvd142m", pretrained=True, num_classes=0); model.eval()
MEAN = np.array([0.485,0.456,0.406], np.float32); STD = np.array([0.229,0.224,0.225], np.float32)

def embed(paths, size=518):
    batch = []
    for path in paths:
        image = cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB)
        h, w = image.shape[:2]; side = max(h, w)
        canvas = np.zeros((side, side, 3), np.uint8)
        canvas[(side-h)//2:(side-h)//2+h, (side-w)//2:(side-w)//2+w] = image
        image = cv2.resize(canvas, (size,size), interpolation=cv2.INTER_AREA).astype(np.float32)/255.0
        batch.append(((image-MEAN)/STD).transpose(2,0,1))
    with torch.no_grad():
        features = model(torch.from_numpy(np.stack(batch)))
    return torch.nn.functional.normalize(features, dim=1).numpy()

def matcher_best(crop):
    """Scores grouped by SKU, so several views of one SKU count as one answer."""
    scores, _ = matcher.score_crop(crop)
    grouped = {}
    for name, value in scores.items():
        sku = name.split("_")[0] + "_" + name.split("_")[1] if name.startswith("sku_l") else name
        grouped[sku] = max(grouped.get(sku, 0.0), float(value))
    return max(grouped.items(), key=lambda kv: kv[1])[0] if grouped else None

random.seed(7)
print(f"{len(items)} crops, {len(by_class)} classes")
print(f"  {'views/SKU':>9}  {'embedding top-1':>15}  {'matcher top-1':>14}")
for views in (1, 2, 3, 4):
    gallery, test = [], []
    for label, paths in by_class.items():
        shuffled = paths[:]; random.shuffle(shuffled)
        gallery += [(p, label) for p in shuffled[:views]]
        test += [(p, label) for p in shuffled[views:]]
    if not test:
        continue
    features_gallery = embed([p for p, _ in gallery])
    features_test = embed([p for p, _ in test])
    gallery_labels = [l for _, l in gallery]
    similarity = features_test @ features_gallery.T
    correct = sum(1 for i, (_, truth) in enumerate(test) if gallery_labels[int(similarity[i].argmax())] == truth)
    matcher_correct = sum(1 for p, truth in test if matcher_best(cv2.imread(p)) == truth)
    print(f"  {views:>9}  {correct:>6}/{len(test):<6} {100*correct//len(test):>3}%  "
          f"{matcher_correct:>6}/{len(test):<6} {100*matcher_correct//len(test):>3}%")