"""Build a rebake4 plan by pinning the print to the consultant's neckline.
usage: mkplan.py dir out.json '<json list of shots>'
shot: {a0,a1,seed, t:[x0,y0,x1,y1] neck+collar template at the seed, at:[x,y] print centre at the seed, w: print width at the seed,
       box:[x0,x1,y0,y1] erase box round the print (units of scale), plus any per-frame extras: {"extras":{"a-b":{...}}}}"""
import cv2, numpy as np, json, sys
d, out, shots = sys.argv[1], sys.argv[2], json.loads(sys.argv[3])
G = lambda k: cv2.cvtColor(cv2.imread(f'{d}/f{k:03d}.png'), cv2.COLOR_BGR2GRAY)
P = {}
for sh in shots:
    x0, y0, x1, y1 = sh['t']; T = G(sh['seed'])[y0:y1, x0:x1]; tw, th = x1 - x0, y1 - y0
    tr = {sh['seed']: (1.0, x0, y0, 1.0)}
    for rng in (range(sh['seed'] + 1, sh['a1'] + 1), range(sh['seed'] - 1, sh['a0'] - 1, -1)):
        _, px, py, ps = tr[sh['seed']]
        for k in rng:
            g = G(k); best = (-1,)
            for sc in ps * np.linspace(0.95, 1.05, 5):
                w, h = int(round(tw * sc)), int(round(th * sc)); R = sh.get('search', 45)
                sx0, sy0 = max(0, int(px - R)), max(0, int(py - R)); sub = g[sy0:int(py + R) + h, sx0:int(px + R) + w]
                if sub.shape[0] <= h or sub.shape[1] <= w: continue
                res = cv2.matchTemplate(sub, cv2.resize(T, (w, h), interpolation=cv2.INTER_AREA), cv2.TM_CCOEFF_NORMED)
                _, mv, _, ml = cv2.minMaxLoc(res)
                if mv > best[0]: best = (mv, sx0 + ml[0], sy0 + ml[1], sc)
            if best[0] < sh.get('min', 0.3): best = (best[0], px, py, ps)        # lost (light sweep): hold
            tr[k] = best; _, px, py, ps = best
    ks = sorted(tr); arr = np.array([tr[k] for k in ks], np.float64)
    for c in (1, 2, 3):                                                          # light smoothing
        arr[:, c] = cv2.GaussianBlur(arr[:, c].reshape(-1, 1), (1, 0), sh.get('smooth', 1.5), borderType=cv2.BORDER_REPLICATE).ravel()
    ax, ay = sh['at'][0] - x0, sh['at'][1] - y0
    for k, (sc_, tx, ty, s) in zip(ks, arr):
        px_, py_ = tx + ax * s, ty + ay * s
        e = {'prints': [[px_, py_, s, 0.0, sh['w'], 1.0]], 'print_at': [px_, py_], 'print_w': sh['w'] * s, 'box': sh['box'],
             'sigma': float(np.clip(0.5 + (0.9 - sc_) * 3, 0.5, 1.8)), 'score': round(float(sc_), 2)}
        for rg, ex in sh.get('extras', {}).items():
            a, b = map(int, rg.split('-'))
            if a <= k <= b: e.update(ex)
        P[k] = e
    print(sh['seed'], 'scores', ' '.join(f'{k}:{tr[k][0]:.2f}' for k in ks if k % 6 == 0))
json.dump(P, open(out, 'w'))