"""End-to-end check of measure.py on frames with a known answer: our own twerk test animation, rendered
by model.py as if it were a capture, then measured like one.

    .venv/bin/python synth_test.py make <capture to copy cameras from> [ticks]   # writes frames/synth
    .venv/bin/python measure.py synth --ticks 0:4
    .venv/bin/python synth_test.py check
"""
import json
import sys
from pathlib import Path

import numpy as np
from PIL import Image
from scipy.spatial import cKDTree

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

# Where the fake viewer puts the model: Minecraft's player scale, feet 25 units below the origin.
CALIB = (0.9375, 0, -1.5, 0)
SKIN = "Steve"
OUT = Path("frames/synth")


def twerk(t):
    """TestAnimations.twerk, written in rig terms."""
    up = 0.5 * (1 - np.cos(t / 8 * 2 * np.pi))
    thigh, knee, splay, tilt, arch = -1.1 + 0.15 * up, 1.75 - 0.3 * up, 0.55, 1.4, 0.4 + 0.3 * up
    drop = 12 - 6 * (np.cos(thigh) + np.cos(thigh + knee)) * np.cos(splay)
    back = -6 * (np.sin(thigh) + np.sin(thigh + knee))
    chest = tilt - arch
    v = rest_rig()
    s = lambda name, value: v.__setitem__(INDEX[name], value)
    s("hip_y", drop), s("hip_z", back)
    s("body_pitch", tilt), s("body_bend", arch), s("body_axis", np.pi)
    s("head_pitch", tilt - chest * 0.5)
    for arm, roll in (("right_arm", 0.3), ("left_arm", -0.3)):
        s(f"{arm}_pitch", tilt - chest * 1.1), s(f"{arm}_roll", roll), s(f"{arm}_bend", 0.5)
    for leg, sign in (("right_leg", 1), ("left_leg", -1)):
        s(f"{leg}_pitch", thigh), s(f"{leg}_roll", sign * splay), s(f"{leg}_bend", knee)
        s(f"{leg}_axis", np.pi - sign * 0.35)
    return v


def skin():
    return np.asarray(Image.open(f"skins/{SKIN}.png").convert("RGBA"))


def make(source, ticks):
    meta = json.loads((Path("frames") / source / "views.json").read_text())
    model = Model(spacing=0.08, skin=skin())
    cams = {n: Camera(v, meta["size"], meta["aspect"], 1) for n, v in meta["views"].items()}
    for tick in range(ticks):
        pts = np.concatenate(list(model.pose(rig_to_pose(twerk(tick))).values()))
        for name, (image, mask) in render_colors(pts, model.colors, CALIB, cams).items():
            (OUT / name).mkdir(parents=True, exist_ok=True)
            rgba = np.concatenate([image * 255, mask[..., None] * 255], axis=-1).astype(np.uint8)
            Image.fromarray(rgba).save(OUT / name / f"{tick:04d}.png")
    meta |= {"url": "synthetic twerk", "skin": SKIN, "ticks": ticks, "loop_ticks": None}
    (OUT / "views.json").write_text(json.dumps(meta, indent=2))
    print(f"wrote {ticks} ticks to {OUT}")


def check():
    data = json.loads((OUT / "rig.json").read_text())
    print("calibration:", np.round(data["calib"], 3), "true:", CALIB)
    model = Model(skin=skin())
    for entry in data["ticks"]:
        truth, got = model.pose(rig_to_pose(twerk(entry["tick"]))), model.pose(rig_to_pose(np.array(entry["rig"])))
        # The fitted pose sits in the fitted calibration's frame; compare where the viewer shows them.
        to = lambda p, c: c[0] * p + np.asarray(c[1:]) * [1, -1, -1]
        parts = []
        for part in truth:
            a, b = to(truth[part], CALIB), to(got[part], data["calib"])
            parts.append((part, np.linalg.norm(a - b, axis=1).mean(), cKDTree(b).query(a)[0].mean()))
        print(f"tick {entry['tick']}: score {entry['score']:.3f}  "
              + "  ".join(f"{p} {d:.2f}/{s:.2f}" for p, d, s in parts))
    print("(per part: mean point-for-point / shape-only distance, in model pixels)")


if __name__ == "__main__":
    if sys.argv[1] == "make":
        make(sys.argv[2], int(sys.argv[3]) if len(sys.argv) > 3 else 4)
    else:
        check()