formscout/agents/visualizer.py

"""
PoseVisualizer — annotated overlay video with skeleton, trails, velocity arrows.

Input:  IngestResult + Pose2DResult
Output: .mp4 path (or None on failure/empty layers)
Failure: returns None, never raises.
"""
from __future__ import annotations

import colorsys
import logging
import math
import tempfile
from collections import deque

import cv2
import numpy as np

logger = logging.getLogger(__name__)

# ── COCO constants ────────────────────────────────────────────────────────────

COCO_KEYPOINTS = [
    "nose", "left_eye", "right_eye", "left_ear", "right_ear",
    "left_shoulder", "right_shoulder", "left_elbow", "right_elbow",
    "left_wrist", "right_wrist", "left_hip", "right_hip",
    "left_knee", "right_knee", "left_ankle", "right_ankle",
]

COCO_SKELETON = [
    (0, 1), (0, 2), (1, 3), (2, 4),          # face
    (5, 6), (5, 7), (7, 9), (6, 8), (8, 10), # arms
    (5, 11), (6, 12), (11, 12),               # torso
    (11, 13), (13, 15), (12, 14), (14, 16),  # legs
]

TRAIL_LENGTH = 10
MAX_ARROW_PX = 40
CONF_THRESHOLD = 0.3


# ── Kalman filter ─────────────────────────────────────────────────────────────

class SimpleKalmanFilter:
    """4-state Kalman filter (x, y, vx, vy) for joint tracking."""

    def __init__(self, process_noise: float = 0.01, measurement_noise: float = 0.1):
        self.is_initialized = False
        self.state = np.zeros(4)
        self.cov = np.eye(4) * 0.1
        self.Q = np.eye(4) * process_noise
        self.R = np.eye(2) * measurement_noise
        self.H = np.array([[1, 0, 0, 0], [0, 1, 0, 0]], dtype=float)

    def predict(self, dt: float = 1.0):
        F = np.array([[1, 0, dt, 0], [0, 1, 0, dt], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=float)
        self.state = F @ self.state
        self.cov = F @ self.cov @ F.T + self.Q

    def update(self, x: float, y: float):
        z = np.array([x, y])
        if not self.is_initialized:
            self.state[:2] = z
            self.is_initialized = True
            return
        S = self.H @ self.cov @ self.H.T + self.R
        K = self.cov @ self.H.T @ np.linalg.inv(S)
        self.state = self.state + K @ (z - self.H @ self.state)
        self.cov = (np.eye(4) - K @ self.H) @ self.cov

    def velocity_magnitude(self) -> float:
        vx, vy = self.state[2], self.state[3]
        return math.sqrt(vx * vx + vy * vy)

    def velocity_vector(self) -> tuple[float, float]:
        return float(self.state[2]), float(self.state[3])


# ── Velocity computation ──────────────────────────────────────────────────────

def compute_joint_velocity(
    keypoints_per_frame: list[dict],
    fps: float,
) -> dict[int, list[float]]:
    """
    Compute Kalman-filtered per-joint speed (px/s) for each frame.

    Returns dict[joint_idx, [speed_frame0, ...]] for all 17 COCO joints.
    Missing/low-confidence keypoints yield speed=0.0 for that frame.
    """
    dt = 1.0 / fps if fps > 0 else 1.0
    filters: dict[int, SimpleKalmanFilter] = {j: SimpleKalmanFilter() for j in range(17)}
    result: dict[int, list[float]] = {j: [] for j in range(17)}

    for frame_kps in keypoints_per_frame:
        for j in range(17):
            kf = filters[j]
            kp = frame_kps.get(j)
            kf.predict(dt)
            if kp and kp.get("conf", 0.0) >= CONF_THRESHOLD:
                kf.update(kp["x"], kp["y"])
                speed = kf.velocity_magnitude()
            else:
                speed = 0.0
            result[j].append(speed)

    return result


# ── Helpers ───────────────────────────────────────────────────────────────────

def _conf_to_bgr(conf: float) -> tuple[int, int, int]:
    """Map confidence 0→1 to BGR color red→green via HSV."""
    hue = conf * 120.0 / 360.0
    r, g, b = colorsys.hsv_to_rgb(hue, 1.0, 1.0)
    return (int(b * 255), int(g * 255), int(r * 255))


# ── PoseVisualizer ────────────────────────────────────────────────────────────

class PoseVisualizer:
    """Renders skeleton, trails, and velocity arrows onto video frames."""

    def __init__(self):
        self.last_velocities: dict[int, list[float]] = {}

    # ── Skeleton ──────────────────────────────────────────────────────────────

    def _draw_skeleton(self, frame: np.ndarray, kps: dict) -> np.ndarray:
        """Draw COCO-17 bones (white) and joints (confidence-colored) onto frame."""
        visible = {j: kp for j, kp in kps.items() if kp.get("conf", 0.0) >= CONF_THRESHOLD}

        # Bones
        for j1, j2 in COCO_SKELETON:
            if j1 in visible and j2 in visible:
                p1 = (int(visible[j1]["x"]), int(visible[j1]["y"]))
                p2 = (int(visible[j2]["x"]), int(visible[j2]["y"]))
                cv2.line(frame, p1, p2, (255, 255, 255), 2)

        # Joints
        for j, kp in visible.items():
            pt = (int(kp["x"]), int(kp["y"]))
            color = _conf_to_bgr(kp["conf"])
            cv2.circle(frame, pt, 4, color, -1)
            cv2.circle(frame, pt, 5, (255, 255, 255), 1)

        return frame

    # ── Trails ───────────────────────────────────────────────────────────────

    def _draw_trails(self, frame: np.ndarray, trail_history: dict) -> np.ndarray:
        """Draw fading motion trails for each joint."""
        for joint_idx, trail in trail_history.items():
            pts = list(trail)
            if len(pts) < 2:
                continue
            for i in range(1, len(pts)):
                alpha = i / len(pts)
                brightness = int(255 * alpha)
                color = (brightness, brightness, brightness)
                thickness = max(1, int(3 * alpha))
                p1 = (int(pts[i - 1][0]), int(pts[i - 1][1]))
                p2 = (int(pts[i][0]), int(pts[i][1]))
                cv2.line(frame, p1, p2, color, thickness)
        return frame

    # ── Velocity arrows ───────────────────────────────────────────────────────

    def _draw_velocity_arrows(
        self,
        frame: np.ndarray,
        kps: dict,
        prev_kps: dict | None,
        velocities: dict[int, list[float]],
        frame_idx: int,
    ) -> np.ndarray:
        """Draw per-joint velocity arrows scaled by speed."""
        if prev_kps is None:
            return frame

        all_speeds = [velocities[j][frame_idx] for j in range(17) if frame_idx < len(velocities.get(j, []))]
        peak = max(all_speeds) if all_speeds else 1.0
        if peak == 0.0:
            return frame

        for j in range(17):
            kp = kps.get(j)
            pk = prev_kps.get(j)
            if not kp or not pk:
                continue
            if kp.get("conf", 0.0) < CONF_THRESHOLD:
                continue
            speeds = velocities.get(j, [])
            if frame_idx >= len(speeds):
                continue
            speed = speeds[frame_idx]
            if speed == 0.0:
                continue

            dx = kp["x"] - pk["x"]
            dy = kp["y"] - pk["y"]
            mag = math.sqrt(dx * dx + dy * dy)
            if mag < 1e-6:
                continue

            length = min(speed / peak * MAX_ARROW_PX, MAX_ARROW_PX)
            nx, ny = dx / mag, dy / mag
            start = (int(kp["x"]), int(kp["y"]))
            end = (int(kp["x"] + nx * length), int(kp["y"] + ny * length))

            ratio = speed / peak
            if ratio < 0.33:
                color = (0, 200, 0)     # green
            elif ratio < 0.66:
                color = (0, 140, 255)   # orange
            else:
                color = (0, 0, 255)     # red

            cv2.arrowedLine(frame, start, end, color, 2, tipLength=0.35)

        return frame

    # ── Public ────────────────────────────────────────────────────────────────

    def render_video(
        self,
        ingest,
        pose2d,
        layers: set[str],
        output_path: str,
    ) -> str | None:
        """
        Render annotated video. Returns output_path on success, None otherwise.
        layers: subset of {"skeleton", "trails", "velocity_arrows"}
        """
        if not layers:
            return None

        if not any(pose2d.keypoints):
            return None

        try:
            velocities = compute_joint_velocity(pose2d.keypoints, ingest.fps)
            self.last_velocities = velocities

            frames = ingest.frames
            orig_h, orig_w = frames[0].shape[:2]
            fps = ingest.fps or 30.0

            # Cap at 1280px wide — big frames are slow and don't need to be HQ
            max_w = 1280
            if orig_w > max_w:
                scale = max_w / orig_w
                out_w = max_w
                out_h = int(orig_h * scale)
            else:
                scale = 1.0
                out_w, out_h = orig_w, orig_h

            # Scale keypoint coordinates to match resized frames
            def _scale_kps(kps: dict) -> dict:
                if scale == 1.0:
                    return kps
                return {
                    j: {**kp, "x": kp["x"] * scale, "y": kp["y"] * scale}
                    for j, kp in kps.items()
                }

            scaled_keypoints = [_scale_kps(k) for k in pose2d.keypoints]

            # Prefer H.264 (avc1): browser-playable via OpenCV's bundled FFmpeg, no
            # external ffmpeg needed. Fall back to mp4v only if avc1 can't open
            # (older OpenCV builds) — that case is transcoded after the write.
            writer = None
            used_fourcc = None
            for _tag in ("avc1", "mp4v"):
                candidate = cv2.VideoWriter(
                    output_path, cv2.VideoWriter_fourcc(*_tag), fps, (out_w, out_h)
                )
                if candidate.isOpened():
                    writer, used_fourcc = candidate, _tag
                    break
                candidate.release()
            if writer is None:
                logger.warning("VideoWriter failed to open: %s", output_path)
                return None

            trail_history: dict[int, deque] = {j: deque(maxlen=TRAIL_LENGTH) for j in range(17)}
            prev_kps: dict | None = None

            for frame_idx, (frame, kps) in enumerate(zip(frames, scaled_keypoints)):
                if scale != 1.0:
                    out_frame = cv2.resize(frame, (out_w, out_h), interpolation=cv2.INTER_AREA)
                else:
                    out_frame = frame.copy()

                if "trails" in layers:
                    for j, kp in kps.items():
                        if kp.get("conf", 0.0) >= CONF_THRESHOLD:
                            trail_history[j].append((kp["x"], kp["y"]))
                    out_frame = self._draw_trails(out_frame, trail_history)

                if "skeleton" in layers:
                    out_frame = self._draw_skeleton(out_frame, kps)

                if "velocity_arrows" in layers:
                    out_frame = self._draw_velocity_arrows(
                        out_frame, kps, prev_kps, velocities, frame_idx
                    )

                writer.write(out_frame)
                prev_kps = kps

            writer.release()

            # mp4v is NOT browser-playable; if we had to fall back to it, transcode
            # to H.264 with ffmpeg when available (best effort).
            if used_fourcc != "avc1":
                import os
                import shutil
                import subprocess
                if shutil.which("ffmpeg"):
                    raw = output_path + ".raw.mp4"
                    try:
                        os.replace(output_path, raw)
                        subprocess.run(
                            ["ffmpeg", "-y", "-i", raw, "-c:v", "libx264",
                             "-pix_fmt", "yuv420p", "-movflags", "+faststart", output_path],
                            check=True, capture_output=True,
                        )
                        os.unlink(raw)
                    except Exception as e:
                        logger.warning("ffmpeg transcode failed (%s) — leaving mp4v", e)
                        if os.path.exists(raw):
                            os.replace(raw, output_path)
                else:
                    logger.warning("wrote mp4v but no avc1/ffmpeg — overlay may not "
                                   "play in-browser")

            return output_path

        except Exception as e:
            logger.warning("render_video failed: %s", e)
            return None

    def render_frame(
        self,
        ingest,
        pose2d,
        frame_idx: int,
        layers: set[str],
        caption: str = "",
        out_png: str | None = None,
    ) -> str | None:
        """Render a single annotated still (skeleton + optional trails + caption).

        frame_idx is typically the governing frame from BiomechFeatures.timing.
        Returns the PNG path on success, None on any failure. Never raises.
        """
        try:
            if not (0 <= frame_idx < len(ingest.frames)) or frame_idx >= len(pose2d.keypoints):
                return None

            frame = ingest.frames[frame_idx].copy()
            kps = pose2d.keypoints[frame_idx]

            if "trails" in layers:
                trail: dict[int, deque] = {j: deque(maxlen=TRAIL_LENGTH) for j in range(17)}
                start = max(0, frame_idx - TRAIL_LENGTH)
                for fi in range(start, frame_idx + 1):
                    for j, kp in pose2d.keypoints[fi].items():
                        if kp.get("conf", 0.0) >= CONF_THRESHOLD:
                            trail[j].append((kp["x"], kp["y"]))
                frame = self._draw_trails(frame, trail)

            if "skeleton" in layers:
                frame = self._draw_skeleton(frame, kps)

            if caption:
                cv2.rectangle(frame, (0, 0), (frame.shape[1], 28), (0, 0, 0), -1)
                cv2.putText(frame, caption[:80], (8, 20), cv2.FONT_HERSHEY_SIMPLEX,
                            0.55, (255, 255, 255), 1, cv2.LINE_AA)

            if out_png is None:
                out_png = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name

            ok = cv2.imwrite(out_png, frame)
            return out_png if ok else None
        except Exception as e:
            logger.warning("render_frame failed: %s", e)
            return None


# ── Velocity summary ──────────────────────────────────────────────────────────

def build_velocity_summary(
    keypoints_per_frame: list[dict],
    velocities: dict[int, list[float]],
) -> str:
    """Return markdown table of per-joint avg/peak velocity. Empty string if no valid joints."""
    n_frames = len(keypoints_per_frame)
    if n_frames == 0:
        return ""

    rows = []
    for j in range(17):
        detected = sum(
            1 for kps in keypoints_per_frame
            if kps.get(j, {}).get("conf", 0.0) >= CONF_THRESHOLD
        )
        if detected < n_frames * 0.5:
            continue

        speeds = velocities.get(j, [])
        if not speeds:
            continue

        avg_speed = sum(speeds) / len(speeds)
        peak_speed = max(speeds)
        rows.append((COCO_KEYPOINTS[j], avg_speed, peak_speed))

    if not rows:
        return ""

    rows.sort(key=lambda r: r[2], reverse=True)
    lines = [
        "| Joint | Avg (px/s) | Peak (px/s) |",
        "|---|---|---|",
    ]
    for name, avg, peak in rows:
        lines.append(f"| {name} | {avg:.1f} | {peak:.1f} |")
    return "\n".join(lines)