formscout/types.py

"""
FormScout typed agent contracts.
Every agent accepts and returns frozen dataclasses defined here.
Validate at every boundary — never accept raw dicts across agent boundaries.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any


@dataclass(frozen=True)
class IngestResult:
    """Output of IngestAgent — decoded video frames + metadata."""
    frames: list  # list of np.ndarray HWC BGR
    fps: float
    duration: float
    n_people: int
    width: int
    height: int
    confidence: float = 1.0
    notes: str = ""


@dataclass(frozen=True)
class SegmentResult:
    """Output of SegmentationAgent — per-frame athlete masks."""
    athlete_track_id: int
    masks: list  # list of np.ndarray bool HW per frame
    confidence: float = 1.0
    notes: str = ""


@dataclass(frozen=True)
class Pose2DResult:
    """Output of Pose2DAgent — per-frame 2D keypoints (COCO 17-joint)."""
    keypoints: list  # list[dict[int, dict]] frame→joint→{x,y,conf}
    fps: float
    confidence: float = 0.0
    notes: str = ""


@dataclass(frozen=True)
class Body3DResult:
    """Output of Body3DAgent — optional 3D joint positions."""
    used: bool
    joints_3d: list  # list[dict] frame→joint→{x,y,z} — empty if used=False
    confidence: float = 0.0
    notes: str = ""


@dataclass(frozen=True)
class MovementResult:
    """Output of MovementClassifierAgent — which FMS test is being performed."""
    test_name: str  # "deep_squat"|"hurdle_step"|...|"unknown"
    side: str       # "left"|"right"|"na"
    confidence: float = 0.0
    notes: str = ""

    def __post_init__(self):
        valid_tests = {
            "deep_squat", "hurdle_step", "inline_lunge",
            "shoulder_mobility", "active_slr",
            "trunk_stability_pushup", "rotary_stability", "unknown",
        }
        if self.test_name not in valid_tests:
            raise ValueError(f"test_name must be one of {valid_tests}, got '{self.test_name}'")
        valid_sides = {"left", "right", "na"}
        if self.side not in valid_sides:
            raise ValueError(f"side must be one of {valid_sides}, got '{self.side}'")


@dataclass(frozen=True)
class BiomechFeatures:
    """Output of BiomechanicsAgent — measured angles, alignments, timing."""
    test_name: str
    view: str           # "2d" | "3d"
    side: str           # "left"|"right"|"na"
    angles: dict        # named angle → degrees
    alignments: dict    # named alignment → value
    symmetry_delta: float | None  # |left - right| or None for non-bilateral
    timing: dict        # event name → frame index
    confidence: float = 0.0
    notes: str = ""

    def __post_init__(self):
        if self.view not in ("2d", "3d"):
            raise ValueError(f"view must be '2d' or '3d', got '{self.view}'")


@dataclass(frozen=True)
class ScoreResult:
    """Output of ScoringAgent (ST-GCN head) — provisional numeric score."""
    score: int  # 0–3
    rationale: str
    confidence: float
    needs_human: bool = False
    notes: str = ""

    def __post_init__(self):
        if not self.needs_human and not (0 <= self.score <= 3):
            raise ValueError(f"score must be 0–3, got {self.score}")


@dataclass(frozen=True)
class RetrievalResult:
    """Output of RetrievalAgent — similar scored exemplars from the index."""
    exemplars: list  # list of {clip_id, score, similarity, rationale}
    confidence: float = 1.0
    notes: str = ""


@dataclass(frozen=True)
class JudgeResult:
    """Output of JudgeAgent — final VLM-scored result with rationale."""
    score: int | None   # 0–3 or None if needs_human=True
    rationale: str
    compensation_tags: list
    corrective_hint: str
    confidence: float
    needs_human: bool = False
    notes: str = ""

    def __post_init__(self):
        if not self.needs_human and self.score is not None:
            if not (0 <= self.score <= 3):
                raise ValueError(f"score must be 0–3 when needs_human=False, got {self.score}")
        if self.needs_human and self.score is not None:
            raise ValueError("score must be None when needs_human=True")


@dataclass(frozen=True)
class ReportResult:
    """Output of ReportAgent — assembled scorecard."""
    per_test: list  # list of dicts with test_name, score, judge_result, features
    composite: int | None  # None if any test unscored
    asymmetries: list  # list of {test, left_score, right_score, delta}
    overlay_video_path: str | None
    pdf_path: str | None
    low_confidence_flags: list
    disagreement_flags: list
    notes: str = ""


@dataclass
class PipelineState:
    """Mutable state threaded through the Director."""
    video_path: str
    ingest: IngestResult | None = None
    segment: SegmentResult | None = None
    pose2d: Pose2DResult | None = None
    body3d: Body3DResult | None = None
    movement: MovementResult | None = None
    features: BiomechFeatures | None = None
    stgcn_score: ScoreResult | None = None
    retrieval: RetrievalResult | None = None
    judge: JudgeResult | None = None
    report: ReportResult | None = None
    errors: list = field(default_factory=list)
    warnings: list = field(default_factory=list)