"""Puts captured frames next to our render of the fitted pose, with the same skin and cameras.

    .venv/bin/python compare.py <name> [tick] [views]

Writes frames/<name>/compare.png: for each view, theirs on the left and ours on the right.
"""
import json
import sys
from pathlib import Path

import numpy as np
from PIL import Image

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

name = sys.argv[1]
tick = int(sys.argv[2]) if len(sys.argv) > 2 else 0
views = sys.argv[3].split(",") if len(sys.argv) > 3 else ["y000", "y090", "y180", "y270", "y045"]
base = Path("frames") / name
meta = json.loads((base / "views.json").read_text())
data = json.loads((base / "rig.json").read_text())
rig = np.array(next(t["rig"] for t in data["ticks"] if t["tick"] == tick))
skin = np.asarray(Image.open(Path("skins") / f"{meta.get('skin', 'Steve')}.png").convert("RGBA"))
model = Model(spacing=0.12, skin=skin)
scale = 2
cams = {v: Camera(meta["views"][v], meta["size"], meta["aspect"], scale) for v in views}
ours = render_colors(np.concatenate(list(model.pose(rig_to_pose(rig)).values())), model.colors, data["calib"], cams)
bg = np.array([40, 40, 52], np.uint8)
tiles = []
for v in views:
    theirs = np.asarray(Image.open(base / v / f"{tick:04d}.png").convert("RGBA").reduce(scale))
    a = np.where(theirs[..., 3:] > 127, theirs[..., :3], bg)
    image, mask = ours[v]
    b = np.where(mask[..., None], (image * 255).astype(np.uint8), bg)
    tiles.append(np.concatenate([a, b], axis=1))
sheet = np.concatenate(tiles, axis=0)
ys, xs = np.nonzero((sheet != bg).any(axis=2))
Image.fromarray(sheet[:, max(xs.min() - 8, 0):xs.max() + 8]).save(base / "compare.png")
print("wrote", base / "compare.png")