"""Matcher variants. Each variant returns per-SKU evidence dicts so ranking and
open-set features can be analysed separately."""
import sys, os, json, time, glob, pickle, argparse
import cv2, numpy as np
sys.path.insert(0, os.path.dirname(__file__)); from diag import crops, sku_of

class Matcher:
    def __init__(self, stock, tmpl_h=None, root=False, geom='homog', nfeat=2000, scales=(None,)):
        self.sift = cv2.SIFT_create(nfeatures=nfeat); self.root = root; self.geom = geom
        self.T = []
        for p in sorted(glob.glob(f'{stock}/*.jp*g')+glob.glob(f'{stock}/*.png')):
            img = cv2.imread(p); name = os.path.splitext(os.path.basename(p))[0]
            for h in scales:
                im = img if h is None else cv2.resize(img, (max(1,int(img.shape[1]*h/img.shape[0])), h), interpolation=cv2.INTER_AREA)
                g = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY); kp, des = self.sift.detectAndCompute(g, None)
                des = self._root(des)
                hsv = cv2.cvtColor(im, cv2.COLOR_BGR2HSV)
                self.T.append(dict(name=name, sku=sku_of(name), kp=kp, des=des, shape=g.shape, hist=self._hist(hsv)))
        self.bf = cv2.BFMatcher(cv2.NORM_L2)
    def _root(self, des):
        if des is None or not self.root: return des
        des = des / (np.abs(des).sum(1, keepdims=True) + 1e-7); return np.sqrt(des).astype(np.float32)
    @staticmethod
    def _hist(hsv):
        m = (hsv[...,1] > 60) & (hsv[...,2] > 50)  # chromatic pixels only
        h = cv2.calcHist([hsv], [0], m.astype(np.uint8), [18], [0, 180]).ravel()
        return h / (h.sum() + 1e-6)
    def score(self, crop):
        g = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY); kq, dq = self.sift.detectAndCompute(g, None); dq = self._root(dq)
        hq = self._hist(cv2.cvtColor(crop, cv2.COLOR_BGR2HSV))
        ev = {}
        for t in self.T:
            e = dict(inliers=0, cover=0.0, color=float(np.minimum(hq, t['hist']).sum()), nq=len(kq))
            if dq is not None and t['des'] is not None and len(dq) >= 2 and len(t['des']) >= 2:
                good = [m for m, n in (p for p in self.bf.knnMatch(dq, t['des'], k=2) if len(p) == 2) if m.distance < 0.75*n.distance]
                if len(good) >= 4:
                    src = np.float32([kq[m.queryIdx].pt for m in good]); dst = np.float32([t['kp'][m.trainIdx].pt for m in good])
                    if self.geom == 'homog':
                        H, mask = cv2.findHomography(src, dst, cv2.RANSAC, 0.02*max(t['shape']))
                    else:
                        H, mask = cv2.estimateAffinePartial2D(src, dst, method=cv2.RANSAC, ransacReprojThreshold=0.02*max(t['shape']))
                    if mask is not None:
                        inl = mask.ravel().astype(bool); n = int(inl.sum())
                        if n >= 3:
                            pts = dst[inl]; hull = cv2.convexHull(pts)
                            e['inliers'] = n; e['cover'] = float(cv2.contourArea(hull) / (t['shape'][0]*t['shape'][1]))
                            # crop-side coverage too: does the match span the query?
                            qh = cv2.convexHull(src[inl]); e['qcover'] = float(cv2.contourArea(qh)/(g.shape[0]*g.shape[1]))
            k = t['sku']
            if k not in ev or e['inliers'] > ev[k]['inliers']: ev[k] = e
        return ev

if __name__ == '__main__':
    ap = argparse.ArgumentParser(); ap.add_argument('--stock'); ap.add_argument('--labels'); ap.add_argument('--live', action='store_true')
    ap.add_argument('--tmpl-h', type=int, default=None); ap.add_argument('--root', action='store_true'); ap.add_argument('--geom', default='homog')
    ap.add_argument('--scales', default=''); ap.add_argument('--out')
    a = ap.parse_args(); L = json.load(open(a.labels))
    scales = tuple(int(s) for s in a.scales.split(',')) if a.scales else ((a.tmpl_h,) if a.tmpl_h else (None,))
    t0 = time.perf_counter(); m = Matcher(a.stock, root=a.root, geom=a.geom, scales=scales); init = time.perf_counter()-t0
    rows = []
    for fc, c in crops('add/image6.jpg', L, a.live):
        if c.size == 0 or min(c.shape[:2]) < 8: rows.append(dict(index=fc['index'], ev={}, ms=0)); continue
        t = time.perf_counter(); ev = m.score(c); rows.append(dict(index=fc['index'], ev=ev, ms=(time.perf_counter()-t)*1000, shape=c.shape[:2]))
    pickle.dump(dict(rows=rows, enrolled=sorted({t['sku'] for t in m.T}), init_s=init, args=vars(a)), open(a.out, 'wb'))
    print(a.out, 'init', round(init,2), 'median ms', round(float(np.median([r['ms'] for r in rows if r['ms']])),1))