"""Turns measured rigs (frames/<name>/rig.json) into an Emotecraft emote our loader reads.

    .venv/bin/python fit.py <name> [--out path.json] [--angle-tol 1] [--offset-tol 0.1]

Every channel is keyed independently, keeping only the ticks linear interpolation can't reproduce
within the tolerances. Values go through the inverse of EmotecraftLoader's conversions.
"""
import argparse
import json
from pathlib import Path

import numpy as np

from model import INDEX, Model, rest_rig, rig_to_pose

PARTS = {"root": "body", "head": "head", "body": "torso", "right_arm": "rightArm", "left_arm": "leftArm",
         "right_leg": "rightLeg", "left_leg": "leftLeg"}
# EmotecraftLoader.REST_PIVOTS: limb positions in the file are absolute pivots.
REST = {"right_arm": (-5, 2, 0), "left_arm": (5, 2, 0), "right_leg": (-1.9, 12, 0.1), "left_leg": (1.9, 12, 0.1)}
CHANNELS = ["x", "y", "z", "pitch", "yaw", "roll", "bend", "axis"]


def to_file(bone, channel, value):
    """The inverse of EmotecraftLoader.convert, writing degrees."""
    if channel in ("x", "y", "z"):
        if bone == "root":
            return value / 16 if channel == "z" else -value / 16
        return value + REST.get(bone, (0, 0, 0))["xyz".index(channel)]
    degrees = np.degrees(value)
    if channel == "bend":
        # Limbs bend the other way in Emotecraft; the torso agrees with ours.
        return degrees if bone == "body" else -degrees
    if bone == "root" and channel in ("pitch", "yaw"):
        return -degrees
    return degrees


def keep(values, tol):
    """Indices to key so linear interpolation stays within tol of every sample (Douglas-Peucker)."""
    keys = {0, len(values) - 1}

    def split(a, b):
        if b - a < 2:
            return
        t = np.arange(a, b + 1)
        line = values[a] + (values[b] - values[a]) * (t - a) / (b - a)
        err = np.abs(values[a:b + 1] - line)
        i = int(np.argmax(err))
        if err[i] > tol:
            keys.add(a + i)
            split(a, a + i)
            split(a + i, b)

    split(0, len(values) - 1)
    return sorted(keys)


def ground(rigs):
    """Shifts every tick vertically so the lowest foot over the whole emote is where a standing
    player's feet are. The viewer offset was a guess (measure.py keeps it fixed), so the measured
    height is only right up to a constant."""
    model = Model()
    lowest = lambda rig: max(model.pose(rig_to_pose(rig))[leg][:, 1].max() for leg in ("right_leg", "left_leg"))
    shift = lowest(rest_rig()) - max(lowest(r) for r in rigs)  # model Y points down
    for r in rigs:
        r[INDEX["root_y"]] += shift
    return shift


def periodic_smooth(values, harmonics, angle):
    """A loop's channel kept to its first few Fourier harmonics: removes the fit's tick-to-tick noise
    and joins the loop's seam. Returns the smoothed samples plus the loop's next sample (its start
    again, a whole number of turns on for an angle that spins)."""
    n = len(values)
    turns = 0.0
    if angle:
        # How far the angle travels in one loop: nothing for a sway, whole turns for a spin.
        step = (values[0] - values[-1] + np.pi) % (2 * np.pi) - np.pi
        turns = values[-1] + step - values[0]
    ramp = turns * np.arange(n + 1) / n
    spectrum = np.fft.rfft(values - ramp[:n])
    spectrum[harmonics + 1:] = 0
    smooth = np.fft.irfft(spectrum, n)
    return np.append(smooth, smooth[0]) + ramp


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("name")
    parser.add_argument("--out")
    parser.add_argument("--angle-tol", type=float, default=1.0, help="degrees")
    parser.add_argument("--offset-tol", type=float, default=0.1, help="pixels")
    parser.add_argument("--no-ground", action="store_true", help="keep the measured height")
    parser.add_argument("--harmonics", type=int, default=4,
                        help="for a whole loop, how many Fourier harmonics each channel keeps (0 = no smoothing)")
    parser.add_argument("--rig", help="read this rig file instead of frames/<name>/rig.json")
    parser.add_argument("--symmetric", action="store_true",
                        help="for a whole loop that sways once per loop, remove its left/right lopsidedness "
                             "(see symmetry.py)")
    parser.add_argument("--title", help="the emote's name (default: from <name>)")
    parser.add_argument("--description", help="default: where it was measured from")
    parser.add_argument("--author", default="Saturn capture")
    args = parser.parse_args()

    base = Path("frames") / args.name
    data = json.loads(Path(args.rig).read_text() if args.rig else (base / "rig.json").read_text())
    meta = json.loads((base / "views.json").read_text())
    ticks = data["ticks"]
    loop = meta.get("loop_ticks")
    rigs = [np.array(t["rig"]) for t in ticks]
    if args.symmetric:
        if not (loop and len(rigs) == loop):
            parser.error("--symmetric needs a whole loop")
        from symmetry import symmetrize
        rigs = list(symmetrize(rigs, sway=True))
        print("made symmetric: each tick averaged with the mirror of the pose half a loop on")
    if not args.no_ground:
        print(f"grounded: moved {ground(rigs):+.2f} px down")
    poses = [rig_to_pose(r) for r in rigs]
    first = ticks[0]["tick"]
    whole_loop = bool(loop) and len(poses) == loop
    if whole_loop and args.harmonics:
        print(f"smoothed each channel to {args.harmonics} harmonics of the {loop}-tick loop")

    moves = {}
    total = 0
    for bone, part in PARTS.items():
        for channel in CHANNELS:
            raw = np.array([getattr(p[bone], channel) for p in poses], float)
            angle = channel not in ("x", "y", "z")
            if angle:
                raw = np.unwrap(raw)
            if whole_loop:
                # A loop's last tick blends back into its first.
                raw = periodic_smooth(raw, args.harmonics, angle) if args.harmonics else np.append(raw, raw[0])
            values = np.array([to_file(bone, channel, v) for v in raw])
            if np.all(np.abs(values - values[0]) < 1e-6) and abs(values[0]) < 1e-6 and channel != "x":
                continue
            tol = args.offset_tol if channel in ("x", "y", "z") else args.angle_tol
            if bone == "root" and channel in ("x", "y", "z"):
                tol /= 16
            for i in keep(values, tol):
                moves.setdefault(i, {}).setdefault(part, {})[channel] = round(float(values[i]), 4)
                total += 1

    emote = {
        "version": 3,
        "name": args.title or args.name.replace("_", " ").title(),
        "description": args.description or f"Measured from {meta['url']}",
        "author": args.author,
        "emote": {
            "degrees": True,
            "isLoop": bool(loop),
            "returnTick": 0,
            "endTick": len(poses) - 1,
            "stopTick": len(poses) if whole_loop else len(poses) - 1,
            "moves": [{"tick": i, "easing": "LINEAR", **parts} for i, parts in sorted(moves.items())],
        },
    }
    out = Path(args.out) if args.out else base / f"{args.name}.json"
    out.write_text(json.dumps(emote, indent=2))
    print(f"wrote {out}: {len(emote['emote']['moves'])} ticks keyed, {total} channel keyframes "
          f"(from {len(poses)} samples, starting at captured tick {first})")


if __name__ == "__main__":
    main()