src/tracking/models.py

"""
Pydantic models for play tracking.

These models represent detected plays, their temporal boundaries,
and the state machine configuration/state for play detection.

Architecture:
- PlayTracker (parent): Coordinates sub-trackers, holds all plays
- NormalPlayTracker: Handles standard plays (clock reset to 40, countdown)
- SpecialPlayTracker: Handles 40→25 transitions (timeouts, punts, FGs, XPs)
"""

from enum import Enum
from typing import Optional, List, Tuple

from pydantic import BaseModel, Field

# =============================================================================
# Utility Functions
# =============================================================================


def determine_timeout_team(home_change: int, away_change: int) -> Optional[str]:
    """
    Determine which team called a timeout based on timeout count changes.

    Validates that EXACTLY one team's count decreased by 1 while the other
    team's count stayed the same.

    Args:
        home_change: Change in home team timeouts (before - after, positive = decrease)
        away_change: Change in away team timeouts (before - after, positive = decrease)

    Returns:
        "home" or "away" if a valid timeout pattern detected, None otherwise
    """
    if home_change == 1 and away_change == 0:
        return "home"
    if away_change == 1 and home_change == 0:
        return "away"
    return None


# =============================================================================
# Enums
# =============================================================================


class TrackerMode(Enum):
    """Which sub-tracker is currently active."""

    NORMAL = "normal"  # Normal play tracker is active
    SPECIAL = "special"  # Special play tracker is handling 40→25 transition


class PlayState(Enum):
    """Current state of play detection (used by NormalPlayTracker)."""

    IDLE = "idle"  # No scorebug detected, waiting
    PRE_SNAP = "pre_snap"  # Scorebug visible, clock ticking down before snap
    PLAY_IN_PROGRESS = "play_in_progress"  # Ball snapped, play is live
    POST_PLAY = "post_play"  # Play ended, waiting for next play setup
    NO_SCOREBUG = "no_scorebug"  # Scorebug lost during/after play (e.g., replay)


class SpecialPlayPhase(Enum):
    """
    State machine phases for SpecialPlayTracker when handling 40→25 transitions.

    Flow:
    1. CHECKING_COUNTDOWN: Check if 25 immediately counts down (weird clock)
    2. MONITORING_TIMEOUT: Wait 4-7s for timeout indicator change
    3. WAITING_FOR_END: Track as special play until 7s or scorebug disappears
    4. RESOLVED: Classification complete, ready to hand back control
    """

    CHECKING_COUNTDOWN = "checking_countdown"  # First 2s: checking if 25 counts down
    MONITORING_TIMEOUT = "monitoring_timeout"  # 4-7s: watching for timeout indicator change
    WAITING_FOR_END = "waiting_for_end"  # Tracking special play until end condition
    RESOLVED = "resolved"  # Classification complete


class PlayEvent(BaseModel):
    """Represents a detected play with start and end times."""

    play_number: int = Field(..., description="Sequential play number")
    start_time: float = Field(..., description="Video timestamp (seconds) when play started")
    end_time: float = Field(..., description="Video timestamp (seconds) when play ended - from backward counting")
    confidence: float = Field(..., description="Overall confidence score")
    start_method: str = Field(..., description="How start was detected: 'clock_reset', 'clock_reset_25', 'clock_freeze'")
    end_method: str = Field(..., description="How end was detected: 'backward_calc' (primary), 'direct_detect' (secondary)")
    direct_end_time: Optional[float] = Field(None, description="End time from direct detection (for comparison)")
    start_clock_value: Optional[int] = Field(None, description="Clock value at start detection")
    end_clock_value: Optional[int] = Field(None, description="Clock value used for backward calculation")
    play_type: str = Field("normal", description="Type of play: 'normal', 'special' (punt/fg/xp after 25-second reset)")
    has_flag: bool = Field(False, description="Whether a penalty FLAG was detected during this play")


class TrackPlayStateConfig(BaseModel):
    """Configuration settings for play state tracking.

    These values control the detection thresholds and timing parameters
    used by the state machine to identify play boundaries.
    """

    clock_stable_frames: int = Field(3, description="Frames with same clock value to consider it 'stable'")
    max_play_duration: float = Field(15.0, description="Maximum expected play duration in seconds")
    scorebug_lost_timeout: float = Field(30.0, description="Seconds before resetting state when scorebug lost")
    required_countdown_ticks: int = Field(3, description="Number of consecutive descending ticks required to confirm play end")
    min_clock_jump_for_reset: int = Field(5, description="Minimum jump in clock value to consider it a valid reset (40 from X where X <= 40 - this value)")

    # FLAG detection settings (penalty flag indicator)
    flag_extension_timeout: float = Field(15.0, description="Max seconds to extend play capture after FLAG first detected")
    capture_flag_plays: bool = Field(True, description="Whether to capture plays with FLAGS (bypasses quiet time filter)")

    # Opening kickoff detection settings
    opening_kickoff_min_consecutive_frames: int = Field(3, description="Required consecutive frames with valid clock readings to confirm kickoff start")
    max_opening_kickoff_duration: float = Field(90.0, description="Maximum valid kickoff duration (seconds). If exceeded, reject and keep searching.")


class TimeoutInfo(BaseModel):
    """Timeout information for a frame."""

    home_timeouts: Optional[int] = Field(None, description="Number of home team timeouts remaining")
    away_timeouts: Optional[int] = Field(None, description="Number of away team timeouts remaining")
    confidence: float = Field(0.0, description="Confidence of the timeout reading")


class FlagInfo(BaseModel):
    """FLAG indicator information for a frame."""

    detected: bool = Field(False, description="Whether FLAG is detected (yellow present AND valid hue)")
    yellow_ratio: float = Field(0.0, description="Ratio of yellow pixels in FLAG region")
    mean_hue: float = Field(0.0, description="Mean hue of yellow pixels (helps distinguish orange)")
    scorebug_verified: bool = Field(True, description="Whether scorebug was verified present via template matching")
    is_valid_yellow: bool = Field(False, description="True if mean_hue >= threshold (not orange)")


class ClockResetStats(BaseModel):
    """Statistics about clock reset classifications."""

    total: int = Field(0, description="Total 40→25 resets detected")
    weird_clock: int = Field(0, description="Class A: 25 counts down immediately (rejected)")
    timeout: int = Field(0, description="Class B: Timeout indicator changed")
    special: int = Field(0, description="Class C: Special play (injury/punt/FG/XP)")


# =============================================================================
# State models for restructured tracker architecture
# =============================================================================


class SpecialPlayHandoff(BaseModel):
    """
    Data passed from NormalPlayTracker to SpecialPlayTracker when a 40→25 transition is detected.

    Contains all context needed for the SpecialPlayTracker to classify and handle the transition.
    """

    transition_timestamp: float = Field(..., description="When the 40→25 transition occurred")
    home_timeouts_at_40: Optional[int] = Field(None, description="Home team timeouts before transition")
    away_timeouts_at_40: Optional[int] = Field(None, description="Away team timeouts before transition")
    timeout_confidence_at_40: float = Field(0.0, description="Confidence of timeout reading at 40")
    was_in_play: bool = Field(False, description="Whether a play was in progress when transition detected")
    play_start_time: Optional[float] = Field(None, description="Start time of play in progress (if any)")
    time_at_40: float = Field(0.0, description="How long clock was at 40 before transitioning to 25")
    from_freeze_detection: bool = Field(False, description="Whether this handoff came from clock freeze→25 detection (always treat as special play)")


class NormalTrackerState(BaseModel):
    """
    State for NormalPlayTracker.

    Tracks clock behavior and play boundaries for normal plays.
    """

    # Current detection state
    state: PlayState = Field(PlayState.IDLE, description="Current state of the normal play tracker")

    # Clock tracking
    last_clock_value: Optional[int] = Field(None, description="Last observed play clock value")
    last_clock_timestamp: Optional[float] = Field(None, description="Timestamp of last clock reading")
    clock_stable_count: int = Field(0, description="Number of consecutive frames with same clock value")
    last_scorebug_timestamp: Optional[float] = Field(None, description="Timestamp of last scorebug detection")

    # Current play tracking
    current_play_start_time: Optional[float] = Field(None, description="Start time of the current play")
    current_play_start_method: Optional[str] = Field(None, description="Method used to detect play start")
    current_play_start_clock: Optional[int] = Field(None, description="Clock value when play started")
    current_play_clock_base: int = Field(40, description="Clock base for current play (40 or 25)")
    current_play_type: str = Field("normal", description="Type of current play: 'normal' or 'special'")
    first_40_timestamp: Optional[float] = Field(None, description="When we first saw 40 in current play")
    countdown_history: List[Tuple[float, int]] = Field(default_factory=list, description="(timestamp, clock_value) pairs")
    direct_end_time: Optional[float] = Field(None, description="Direct end time observation")

    # Timeout tracking
    last_home_timeouts: Optional[int] = Field(None, description="Last observed home team timeouts")
    last_away_timeouts: Optional[int] = Field(None, description="Last observed away team timeouts")
    last_timeout_confidence: float = Field(0.0, description="Confidence of last timeout reading")

    # FLAG tracking (penalty flag indicator)
    flag_detected_at: Optional[float] = Field(None, description="When FLAG first appeared in current FLAG event")
    flag_last_seen_at: Optional[float] = Field(None, description="Last time FLAG was visible")
    flag_active: bool = Field(False, description="Whether FLAG is currently active (extends play capture)")
    current_play_has_flag: bool = Field(False, description="Whether FLAG was detected during the current play")

    # Clock freeze tracking (for special play detection)
    clock_freeze_start_timestamp: Optional[float] = Field(None, description="When clock first reached current frozen value")
    clock_freeze_value: Optional[int] = Field(None, description="The clock value that has been frozen/stable")

    # Last play info (for continuation detection)
    last_play_end_time: Optional[float] = Field(None, description="End time of the last completed play")

    # Handoff signaling
    request_special_handoff: bool = Field(False, description="Flag to request handoff to SpecialPlayTracker")
    pending_handoff: Optional[SpecialPlayHandoff] = Field(None, description="Data for pending handoff to special tracker")

    # Opening kickoff tracking
    # The opening kickoff is special because the scorebug appears mid-play
    # We detect it when scorebug first appears and end it on first clock reset to 40
    # Requires k consecutive valid clock readings to confirm kickoff (filters out pre-game noise)
    opening_kickoff_active: bool = Field(False, description="Whether we're currently tracking the opening kickoff")
    opening_kickoff_complete: bool = Field(False, description="Whether opening kickoff has been recorded (only happens once)")
    first_scorebug_timestamp: Optional[float] = Field(None, description="When scorebug first appeared in the video (verified by template matching)")
    first_clock_reading_timestamp: Optional[float] = Field(None, description="When we first got a valid clock reading")
    opening_kickoff_consecutive_readings: int = Field(0, description="Count of consecutive frames with valid clock readings for kickoff detection")
    opening_kickoff_candidate_timestamp: Optional[float] = Field(None, description="Timestamp when consecutive clock readings started (candidate kickoff start)")


class SpecialTrackerState(BaseModel):
    """
    State for SpecialPlayTracker.

    Manages the 40→25 transition classification state machine.
    """

    # Current phase of special play handling
    phase: SpecialPlayPhase = Field(SpecialPlayPhase.RESOLVED, description="Current phase of special play handling")

    # Transition context (from handoff)
    transition_timestamp: Optional[float] = Field(None, description="When the 40→25 transition occurred")
    home_timeouts_at_40: Optional[int] = Field(None, description="Home team timeouts before transition")
    away_timeouts_at_40: Optional[int] = Field(None, description="Away team timeouts before transition")
    timeout_confidence_at_40: float = Field(0.0, description="Confidence of timeout reading at 40")
    was_in_play: bool = Field(False, description="Whether a play was in progress when transition detected")
    play_start_time: Optional[float] = Field(None, description="Start time of play that led to this transition")
    time_at_40: float = Field(0.0, description="How long clock was at 40 before transitioning")
    from_freeze_detection: bool = Field(False, description="Whether this came from clock freeze→25 detection (always special play)")

    # Classification tracking
    lowest_clock_seen: int = Field(25, description="Lowest clock value seen since 25 (for countdown detection)")
    last_scorebug_timestamp: Optional[float] = Field(None, description="Last time scorebug was visible")

    # Scorebug disappearance tracking (debounce for noisy detection)
    consecutive_scorebug_absent: int = Field(0, description="Number of consecutive frames without scorebug detection")

    # Timeout confirmation tracking (require k consecutive confirmations)
    pending_timeout_team: Optional[str] = Field(None, description="Team with pending timeout detection")
    consecutive_timeout_confirmations: int = Field(0, description="Number of consecutive frames confirming timeout")
    saw_timeout_increase: bool = Field(False, description="Whether we saw any timeout count INCREASE (invalidates timeout detection)")

    # Resolution
    resolution_complete: bool = Field(False, description="Whether classification is complete")
    classification: Optional[str] = Field(None, description="Final classification: 'weird_clock', 'timeout', 'special'")
    timeout_team: Optional[str] = Field(None, description="Which team called timeout: 'home' or 'away'")