"""Draws the fitted pose's outline over the captured frames, to see where the fit is off.

    .venv/bin/python overlay.py <name> [--ticks 0:13:3] [--views y000,y090,y045,top000]

Writes frames/<name>/overlay.png: rows are views, columns ticks. Red is where only our model is, blue
where only theirs is, and grey where both are.
"""
import argparse
import json
from pathlib import Path

import numpy as np
from PIL import Image

from model import Model, rig_to_pose
from render import Camera, render_masks


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("name")
    parser.add_argument("--ticks", default="0:13:3")
    parser.add_argument("--views", default="y000,y090,y045,top000")
    args = parser.parse_args()

    base = Path("frames") / args.name
    meta = json.loads((base / "views.json").read_text())
    data = json.loads((base / "rig.json").read_text())
    by_tick = {t["tick"]: np.array(t["rig"]) for t in data["ticks"]}
    ticks = [t for t in range(*map(int, args.ticks.split(":"))) if t in by_tick]
    views = args.views.split(",")
    scale = 2
    cams = {v: Camera(meta["views"][v], meta["size"], meta["aspect"], scale) for v in views}
    model = Model(spacing=0.25)

    tiles = []
    for v in views:
        row = []
        for t in ticks:
            pts = np.concatenate(list(model.pose(rig_to_pose(by_tick[t])).values()))
            ours = render_masks(pts, data["calib"], {v: cams[v]})[v]
            frame = Image.open(base / v / f"{t:04d}.png").reduce(scale)
            theirs = np.asarray(frame)[:, :, 3] > 127
            img = np.full(ours.shape + (3,), 30, np.uint8)
            img[theirs & ours] = (150, 150, 150)
            img[ours & ~theirs] = (230, 60, 60)
            img[theirs & ~ours] = (60, 110, 240)
            row.append(img)
        tiles.append(row)

    both = np.zeros(tiles[0][0].shape[:2], bool)
    for row in tiles:
        for img in row:
            both |= img.sum(axis=2) > 90
    ys, xs = np.nonzero(both)
    y0, y1, x0, x1 = max(ys.min() - 4, 0), ys.max() + 4, max(xs.min() - 4, 0), xs.max() + 4
    sheet = np.concatenate([np.concatenate([img[y0:y1, x0:x1] for img in row], axis=1) for row in tiles], axis=0)
    Image.fromarray(sheet).save(base / "overlay.png")
    print("wrote", base / "overlay.png", "ticks", ticks)


if __name__ == "__main__":
    main()