tests/test_visualizer.py

"""Tests for PoseVisualizer — no GPU, no model downloads."""
import numpy as np
import pytest
from formscout.types import IngestResult, Pose2DResult


def _make_ingest(n=5, h=480, w=640, fps=30.0):
    frames = [np.zeros((h, w, 3), dtype=np.uint8) for _ in range(n)]
    return IngestResult(frames=frames, fps=fps, duration=n / fps, n_people=1, width=w, height=h)


def _make_pose(n=5, w=640, h=480):
    """Synthetic Pose2DResult: 17 joints at fixed pixel positions, conf=0.9."""
    kps_per_frame = []
    for i in range(n):
        frame_kps = {}
        for j in range(17):
            frame_kps[j] = {
                "x": float(50 + j * 30 + i * 2),
                "y": float(100 + j * 20),
                "conf": 0.9,
            }
        kps_per_frame.append(frame_kps)
    return Pose2DResult(keypoints=kps_per_frame, fps=30.0, confidence=0.9, notes="")


class TestComputeJointVelocity:
    def test_returns_17_joints(self):
        from formscout.agents.visualizer import compute_joint_velocity
        pose = _make_pose(n=5)
        result = compute_joint_velocity(pose.keypoints, fps=30.0)
        assert len(result) == 17

    def test_each_list_has_n_frames(self):
        from formscout.agents.visualizer import compute_joint_velocity
        pose = _make_pose(n=5)
        result = compute_joint_velocity(pose.keypoints, fps=30.0)
        for joint_idx, speeds in result.items():
            assert len(speeds) == 5, f"joint {joint_idx} has {len(speeds)} speeds, expected 5"

    def test_speeds_are_non_negative(self):
        from formscout.agents.visualizer import compute_joint_velocity
        pose = _make_pose(n=5)
        result = compute_joint_velocity(pose.keypoints, fps=30.0)
        for speeds in result.values():
            assert all(s >= 0.0 for s in speeds)

    def test_missing_keypoints_give_zero_speed(self):
        from formscout.agents.visualizer import compute_joint_velocity
        empty_kps = [{} for _ in range(5)]
        result = compute_joint_velocity(empty_kps, fps=30.0)
        for speeds in result.values():
            assert all(s == 0.0 for s in speeds)


class TestDrawSkeleton:
    def test_skeleton_draws_without_error(self):
        from formscout.agents.visualizer import PoseVisualizer
        vis = PoseVisualizer()
        frame = np.zeros((480, 640, 3), dtype=np.uint8)
        kps = {j: {"x": float(50 + j * 30), "y": float(100 + j * 20), "conf": 0.9}
               for j in range(17)}
        result = vis._draw_skeleton(frame.copy(), kps)
        assert result.shape == frame.shape
        assert not np.array_equal(result, frame)

    def test_low_confidence_keypoints_not_drawn(self):
        from formscout.agents.visualizer import PoseVisualizer
        vis = PoseVisualizer()
        frame = np.zeros((480, 640, 3), dtype=np.uint8)
        kps = {j: {"x": float(50 + j * 30), "y": 100.0, "conf": 0.1} for j in range(17)}
        result = vis._draw_skeleton(frame.copy(), kps)
        assert np.array_equal(result, frame)


class TestDrawTrails:
    def test_trails_draw_without_error(self):
        from formscout.agents.visualizer import PoseVisualizer, TRAIL_LENGTH
        from collections import deque
        vis = PoseVisualizer()
        frame = np.zeros((480, 640, 3), dtype=np.uint8)
        trail_history = {
            0: deque([(100 + i * 5, 200 + i * 3) for i in range(5)], maxlen=TRAIL_LENGTH)
        }
        result = vis._draw_trails(frame.copy(), trail_history)
        assert result.shape == frame.shape
        assert not np.array_equal(result, frame)

    def test_short_trail_no_crash(self):
        from formscout.agents.visualizer import PoseVisualizer, TRAIL_LENGTH
        from collections import deque
        vis = PoseVisualizer()
        frame = np.zeros((480, 640, 3), dtype=np.uint8)
        trail_history = {0: deque([(100, 200)], maxlen=TRAIL_LENGTH)}
        result = vis._draw_trails(frame.copy(), trail_history)
        assert np.array_equal(result, frame)


class TestDrawVelocityArrows:
    def test_arrows_draw_without_error(self):
        from formscout.agents.visualizer import PoseVisualizer
        vis = PoseVisualizer()
        frame = np.zeros((480, 640, 3), dtype=np.uint8)
        kps = {j: {"x": float(50 + j * 30), "y": float(100 + j * 20), "conf": 0.9}
               for j in range(17)}
        prev_kps = {j: {"x": float(48 + j * 30), "y": float(98 + j * 20), "conf": 0.9}
                    for j in range(17)}
        velocities = {j: [0.0] * 5 for j in range(17)}
        velocities[5] = [0.0, 10.0, 50.0, 80.0, 120.0]
        result = vis._draw_velocity_arrows(frame.copy(), kps, prev_kps, velocities, frame_idx=4)
        assert result.shape == frame.shape

    def test_no_prev_kps_no_crash(self):
        from formscout.agents.visualizer import PoseVisualizer
        vis = PoseVisualizer()
        frame = np.zeros((480, 640, 3), dtype=np.uint8)
        kps = {j: {"x": float(50 + j * 30), "y": 100.0, "conf": 0.9} for j in range(17)}
        velocities = {j: [50.0] * 5 for j in range(17)}
        result = vis._draw_velocity_arrows(frame.copy(), kps, None, velocities, frame_idx=0)
        assert result.shape == frame.shape


class TestRenderVideo:
    def test_creates_mp4_file(self, tmp_path):
        from formscout.agents.visualizer import PoseVisualizer
        vis = PoseVisualizer()
        ingest = _make_ingest(n=5)
        pose = _make_pose(n=5)
        out = str(tmp_path / "out.mp4")
        result = vis.render_video(ingest, pose, {"skeleton"}, out)
        assert result is not None
        import os
        assert os.path.exists(result)
        assert os.path.getsize(result) > 0

    def test_empty_layers_returns_none(self, tmp_path):
        from formscout.agents.visualizer import PoseVisualizer
        vis = PoseVisualizer()
        out = str(tmp_path / "out.mp4")
        result = vis.render_video(_make_ingest(), _make_pose(), set(), out)
        assert result is None

    def test_no_detections_returns_none(self, tmp_path):
        from formscout.agents.visualizer import PoseVisualizer
        vis = PoseVisualizer()
        ingest = _make_ingest(n=5)
        empty_pose = Pose2DResult(
            keypoints=[{} for _ in range(5)], fps=30.0, confidence=0.0, notes=""
        )
        out = str(tmp_path / "out.mp4")
        result = vis.render_video(ingest, empty_pose, {"skeleton"}, out)
        assert result is None

    def test_last_velocities_set_after_render(self, tmp_path):
        from formscout.agents.visualizer import PoseVisualizer
        vis = PoseVisualizer()
        out = str(tmp_path / "out.mp4")
        vis.render_video(_make_ingest(n=5), _make_pose(n=5), {"skeleton"}, out)
        assert len(vis.last_velocities) == 17


class TestBuildVelocitySummary:
    def test_returns_markdown_table(self):
        from formscout.agents.visualizer import build_velocity_summary, compute_joint_velocity
        pose = _make_pose(n=10)
        vels = compute_joint_velocity(pose.keypoints, fps=30.0)
        result = build_velocity_summary(pose.keypoints, vels)
        assert "|" in result
        assert any(name in result for name in ["knee", "shoulder", "hip", "ankle"])

    def test_empty_keypoints_returns_empty_string(self):
        from formscout.agents.visualizer import build_velocity_summary
        empty_kps = [{} for _ in range(5)]
        vels = {j: [0.0] * 5 for j in range(17)}
        result = build_velocity_summary(empty_kps, vels)
        assert result == ""