additional_scorers.py

"""Additional physics-plausibility scorers: timing, direction, and temporal profile.

These capture failure modes that the primary trajectory/energy scorers miss:
  - contact_timing_score  — WHEN does the physics event happen?
  - direction_constraint_score — does motion go the RIGHT WAY?
  - temporal_symmetry_score — is the acceleration profile physically monotonic?

Each function takes (video_path, seed) and returns a float in [0, 1].
"""

from __future__ import annotations

from pathlib import Path

import numpy as np

from score import _read_frames, _compute_flows, _subtract_camera_motion
from seeds import SeedTask

# Seed-id → expected peak timing region (min_frac, max_frac)
_PEAK_TIMING: dict[str, tuple[float, float]] = {
    "gravity_freefall":    (0.60, 1.00),
    "gravity_projectile":  (0.60, 1.00),
    "gravity_buoyancy":    (0.00, 0.40),
    "incline_roll":        (0.55, 1.00),
    "collision_elastic":   (0.25, 0.75),
    "collision_inelastic": (0.25, 0.75),
}

# Seed-id → expected dominant direction sign: (dx_sign, dy_sign)
#   +1 = positive expected, -1 = negative expected, 0 = don't care
_DIRECTION_EXPECTED: dict[str, tuple[int, int]] = {
    "gravity_freefall":    ( 0, +1),   # downward
    "gravity_projectile":  (+1, +1),   # rightward + downward
    "gravity_buoyancy":    ( 0, -1),   # upward
    "incline_roll":        (+1, +1),   # rightward + downward
    "collision_elastic":   ( 0,  0),   # varies
    "collision_inelastic": ( 0,  0),   # varies
}

# Which seeds should have monotonically increasing velocity
_MONOTONIC_SEEDS = {"gravity_freefall", "gravity_projectile", "incline_roll"}
# Which seeds should show peak-then-decrease (collision profile)
_COLLISION_SEEDS = {"collision_elastic", "collision_inelastic"}


def _per_frame_magnitude(flows: list[np.ndarray]) -> np.ndarray:
    """Mean magnitude of camera-corrected flow per frame."""
    mags: list[float] = []
    for raw_flow in flows:
        flow = _subtract_camera_motion(raw_flow)
        mag = np.sqrt(flow[..., 0] ** 2 + flow[..., 1] ** 2)
        thresh = np.percentile(mag, 80)
        mask = mag > thresh
        mags.append(float(mag[mask].mean()) if mask.sum() > 10 else 0.0)
    return np.array(mags)


def _mean_flow_vector(flows: list[np.ndarray]) -> tuple[float, float]:
    """Mean (dx, dy) across all frames, after camera subtraction."""
    dx_all, dy_all = [], []
    for raw_flow in flows:
        flow = _subtract_camera_motion(raw_flow)
        mag = np.sqrt(flow[..., 0] ** 2 + flow[..., 1] ** 2)
        thresh = np.percentile(mag, 80)
        mask = mag > thresh
        if mask.sum() > 10:
            dx_all.append(float(flow[mask, 0].mean()))
            dy_all.append(float(flow[mask, 1].mean()))
    if not dx_all:
        return 0.0, 0.0
    return float(np.mean(dx_all)), float(np.mean(dy_all))


# ═════════════════════════════════════════════════════════════════════════
#  1) Contact timing
# ═════════════════════════════════════════════════════════════════════════

def contact_timing_score(video_path: str | Path, seed: SeedTask) -> float:
    """Score whether the peak-motion event occurs at the expected time.

    Returns 1.0 when the peak falls inside the expected window,
    decaying linearly to 0.0 as it moves away.
    """
    frames = _read_frames(video_path)
    if len(frames) < 6:
        return 0.0

    flows = _compute_flows(frames)
    if not flows:
        return 0.0

    mag = _per_frame_magnitude(flows)
    T = len(mag)
    t_peak = int(np.argmax(mag))
    t_frac = t_peak / T

    lo, hi = _PEAK_TIMING.get(seed.id, (0.0, 1.0))
    if lo <= t_frac <= hi:
        return 1.0
    dist = min(abs(t_frac - lo), abs(t_frac - hi))
    return float(np.clip(1.0 - dist * 3.0, 0.0, 1.0))


# ═════════════════════════════════════════════════════════════════════════
#  2) Direction constraint
# ═════════════════════════════════════════════════════════════════════════

def direction_constraint_score(video_path: str | Path, seed: SeedTask) -> float:
    """Score whether the dominant motion direction matches the expected physics.

    Returns 1.0 for perfect directional alignment, 0.0 for opposite.
    """
    frames = _read_frames(video_path)
    if len(frames) < 6:
        return 0.0

    flows = _compute_flows(frames)
    if not flows:
        return 0.0

    mean_dx, mean_dy = _mean_flow_vector(flows)
    expected = _DIRECTION_EXPECTED.get(seed.id, (0, 0))
    dx_sign, dy_sign = expected

    if dx_sign == 0 and dy_sign == 0:
        return 1.0

    scores: list[float] = []
    if dx_sign != 0:
        s = dx_sign * mean_dx / (abs(mean_dx) + 0.1)
        scores.append(float(np.clip(s, 0.0, 1.0)))
    if dy_sign != 0:
        s = dy_sign * mean_dy / (abs(mean_dy) + 0.1)
        scores.append(float(np.clip(s, 0.0, 1.0)))

    return float(np.mean(scores)) if scores else 1.0


# ═════════════════════════════════════════════════════════════════════════
#  3) Temporal symmetry (monotonicity)
# ═════════════════════════════════════════════════════════════════════════

def temporal_symmetry_score(video_path: str | Path, seed: SeedTask) -> float:
    """Score whether the velocity profile has the expected monotonic shape.

    For freefall/roll: velocity should increase monotonically.
    For collisions: velocity should rise to a peak then fall.
    """
    frames = _read_frames(video_path)
    if len(frames) < 6:
        return 0.0

    flows = _compute_flows(frames)
    if not flows:
        return 0.0

    vel = _per_frame_magnitude(flows)
    T = len(vel)
    if T < 4:
        return 0.0

    if seed.id in _MONOTONIC_SEEDS:
        increasing = sum(1 for t in range(T - 1) if vel[t + 1] > vel[t])
        return float(increasing / (T - 1))

    if seed.id in _COLLISION_SEEDS:
        peak_t = int(np.argmax(vel))
        if peak_t < 1 or peak_t >= T - 1:
            return 0.3

        before_inc = sum(1 for t in range(peak_t) if vel[t + 1] > vel[t])
        after_dec = sum(1 for t in range(peak_t, T - 1) if vel[t + 1] < vel[t])

        before_score = before_inc / peak_t if peak_t > 0 else 0.0
        after_score = after_dec / (T - 1 - peak_t) if (T - 1 - peak_t) > 0 else 0.0
        return float(0.5 * before_score + 0.5 * after_score)

    return 0.5


if __name__ == "__main__":
    from seeds import get_seed

    test_cases = [
        ("outputs/incline_roll__0000_kling.mp4", "incline_roll"),
        ("outputs/gravity_freefall__0000_kling.mp4", "gravity_freefall"),
        ("outputs/collision_elastic__0006_kling.mp4", "collision_elastic"),
    ]

    for path, seed_id in test_cases:
        if not Path(path).exists():
            continue
        seed = get_seed(seed_id)
        print(f"\n{seed_id} ({path}):")
        print(f"  contact_timing:  {contact_timing_score(path, seed):.4f}")
        print(f"  direction:       {direction_constraint_score(path, seed):.4f}")
        print(f"  temporal:        {temporal_symmetry_score(path, seed):.4f}")