# Grouper (the Critique's cr_iou.py, outputs in build/checks and previews/compare): recompute the silhouette overlap of the model's lookalike renders with the Designer's masks
import numpy as np, os, json
from PIL import Image
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
BOX = {"side": (430, 255, 1010, 520), "top": (440, 20, 1000, 250), "bottom": (440, 520, 1000, 745),
       "front": (95, 215, 385, 550), "back": (1075, 235, 1345, 550)}
out = {}
for v, b in BOX.items():
    m = np.array(Image.open(os.path.join(ROOT, "reference", "mask_%s.png" % v)).convert("L").crop(b)) > 127
    im = Image.open(os.path.join(ROOT, "previews", "compare", "lk_%s.png" % v)).convert("RGBA")
    w, h = b[2]-b[0], b[3]-b[1]
    a = np.array(im.resize((w, h), Image.LANCZOS))[:, :, 3] > 127
    inter = (a & m).sum(); uni = (a | m).sum()
    # extra (model outside mask) and missing (mask not covered)
    extra = (a & ~m).sum(); miss = (m & ~a).sum()
    # per-column extents for side/top to find where they differ
    out[v] = dict(iou=round(inter/uni, 4), model_px=int(a.sum()), mask_px=int(m.sum()), extra_px=int(extra), missing_px=int(miss),
                  extra_pct_of_mask=round(100*extra/m.sum(), 2), missing_pct_of_mask=round(100*miss/m.sum(), 2))
    # save diff image at 3x: green = model only, red = sheet only, grey = both
    d = np.zeros((h, w, 3), np.uint8) + 255
    d[a & m] = (170, 170, 170); d[a & ~m] = (0, 170, 0); d[m & ~a] = (220, 0, 0)
    Image.fromarray(d).resize((w*3, h*3), Image.NEAREST).save(os.path.join(ROOT, "previews", "compare", "iou_diff_%s.png" % v))
    # column-wise top/bottom difference (sheet px)
    if v in ("side", "top", "bottom"):
        rows = []
        for x in range(w):
            ma = np.where(a[:, x])[0]; mm = np.where(m[:, x])[0]
            if len(ma) and len(mm):
                rows.append((x + b[0], int(ma.min()) - int(mm.min()), int(ma.max()) - int(mm.max())))
        big = [r for r in rows if abs(r[1]) >= 3 or abs(r[2]) >= 3]
        out[v]["cols_off_3px"] = len(big)
        out[v]["cols_total"] = len(rows)
        # summarize runs
        runs = []
        for r in big:
            if runs and r[0] - runs[-1][1] <= 2:
                runs[-1][1] = r[0]; runs[-1][2] = max(runs[-1][2], abs(r[1])); runs[-1][3] = max(runs[-1][3], abs(r[2]))
            else:
                runs.append([r[0], r[0], abs(r[1]), abs(r[2])])
        out[v]["runs(x0,x1,max_top_off,max_bot_off)"] = runs
print("IOU", json.dumps({k: (v["iou"], v["missing_px"], v["extra_px"]) for k, v in out.items()}))
json.dump(out, open(os.path.join(ROOT, "build", "checks", "iou.json"), "w"), indent=1)