# Copyright 2026 Viggle AI. Licensed under the Apache License, Version 2.0 (see LICENSE-CODE).
# SPDX-License-Identifier: Apache-2.0
"""Keyframe camera paths for the demo service -- pure math, no models.

A key is `{pos, look, src, t, ease?, focal?}`:
  pos    where the camera is, in the world frame W (the source camera of frame `start`: x right, y down,
         z forward), in units of the pivot depth `zm`
  look   the 3D point the camera looks at, same frame and units
  src    which source frame this key shows (absolute frame index of the clip)
  t      which output frame this key lands on: the first key is frame 0, the last is frame `frames - 1`,
         and they increase strictly in between
  ease   optional, eases the pose motion of the segment leaving this key (time stays linear)
  focal  optional focal multiplier, default 1

Between keys `pos` and `look` run on a Catmull-Rom curve (cubic Hermite with finite-difference tangents,
non-uniform in time), `src` and `focal` run linearly. The orientation is derived from `look - pos` with roll
locked to zero -- the training corpus has no roll (p99 2.1 deg), so a rolled camera is genuinely out of
distribution. Time never runs backwards: `src` must be non-decreasing. Bullet time is two keys with the same
`src` at different `t`; a plain move is keys whose `src` advance one frame per output frame.
"""
import math

import numpy as np

UP = np.array([0.0, -1.0, 0.0])   # y is down in the camera frame


def key_times(path, frames):
    """Output frame index of every key."""
    t = [int(k["t"]) for k in path]
    assert len(t) >= 2 and t[0] == 0 and t[-1] == frames - 1 and all(b > a for a, b in zip(t, t[1:])), \
        f"key output frames {t} must run from 0 to {frames - 1}, strictly increasing"
    return t


def hermite(tk, pk, t, ease):
    """Catmull-Rom through (tk, pk[k]) evaluated at t. pk: (K, D). ease[k]: ease segment k."""
    tk, pk = np.asarray(tk, float), np.asarray(pk, float)
    K = len(tk)
    d = np.diff(pk, axis=0) / np.diff(tk)[:, None]              # chord slope of every segment
    m = np.zeros_like(pk)
    m[0], m[-1] = d[0], d[-1]
    m[1:-1] = 0.5 * (d[:-1] + d[1:])
    m[1:-1][d[:-1] * d[1:] <= 0] = 0    # Fritsch-Carlson: no overshoot, and a hold between equal keys stays exactly still
    lim = 3 * np.minimum(np.abs(d[:-1]), np.abs(d[1:]))   # ... and no tangent steeper than 3x the gentler chord, or a slow-then-fast pair dips backwards first
    m[1:-1] = np.clip(m[1:-1], -lim, lim)
    out = np.zeros((len(t), pk.shape[1]))
    for i, x in enumerate(t):
        k = min(int(np.searchsorted(tk, x, side="right")) - 1, K - 2)
        h = tk[k + 1] - tk[k]
        s = (x - tk[k]) / h
        if ease[k]:
            s = (1 - math.cos(math.pi * s)) / 2
        h00, h10, h01, h11 = 2 * s**3 - 3 * s**2 + 1, s**3 - 2 * s**2 + s, -2 * s**3 + 3 * s**2, s**3 - s**2
        out[i] = h00 * pk[k] + h10 * h * m[k] + h01 * pk[k + 1] + h11 * h * m[k + 1]
    return out


def look_at(pos, look, prev=None):
    """c2w rotation (columns right, down, forward) looking from pos at look with zero roll."""
    f = look - pos
    n = np.linalg.norm(f)
    if n < 1e-6:
        return prev if prev is not None else np.eye(3)
    f = f / n
    r = np.cross(f, UP)
    if np.linalg.norm(r) < 1e-6:   # looking straight up or down: keep x as right
        r = np.array([1.0, 0.0, 0.0])
    r = r / np.linalg.norm(r)
    d = np.cross(f, r)
    return np.stack([r, d, f], 1)


def plan_path(path, frames, zm):
    """-> per-frame c2w in W (4x4, translation in scene units), per-frame source frame, per-frame focal,
    per-segment source speed (source frames per output frame: 0 = frozen, 1 = real time)."""
    tk = key_times(path, frames)
    src = [int(k["src"]) for k in path]
    assert all(b >= a for a, b in zip(src, src[1:])), f"source frames {src} run backwards"
    t = np.arange(frames)
    ease = [bool(k.get("ease", False)) for k in path]
    pos = hermite(tk, [k["pos"] for k in path], t, ease) * zm
    look = hermite(tk, [k["look"] for k in path], t, ease) * zm
    tmap = np.rint(np.interp(t, tk, src)).astype(int).tolist()
    focal = np.interp(t, tk, [float(k.get("focal", 1.0)) for k in path]).tolist()
    c2w = np.tile(np.eye(4), (frames, 1, 1))
    R = None
    for i in range(frames):
        R = look_at(pos[i], look[i], R)
        c2w[i, :3, :3], c2w[i, :3, 3] = R, pos[i]
    speed = [(src[k + 1] - src[k]) / (tk[k + 1] - tk[k]) for k in range(len(tk) - 1)]
    return c2w, tmap, focal, speed


def describe(c2w_W, piv_W, zm):
    """The inverse: per-frame (pos, look) in path units from c2w in W. `look` is the point on the optical
    axis at the pivot's depth, so a camera aimed at the pivot reports the pivot itself."""
    out = []
    for M in np.asarray(c2w_W):
        p, f, r = M[:3, 3], M[:3, 2], M[:3, 0]
        d = max(float((piv_W - p) @ f), 0.05 * zm)
        roll = math.degrees(math.atan2(-float(r @ UP), math.hypot(r[0], r[2])))   # 0 for a level camera
        out.append(dict(pos=(p / zm).round(4).tolist(), look=((p + f * d) / zm).round(4).tolist(), roll=round(roll, 2)))
    return out