main.py

#!/usr/bin/env python3
"""
CFB40 - College Football Play Detection Pipeline

End-to-end pipeline for detecting and extracting plays from football game videos.
This script guides the user through:
1. Video selection (or uses test video in --testing mode)
2. Interactive scorebug region selection (or load from saved config)
3. Interactive play clock region selection (or load from saved config)
4. Play detection using the state machine
5. Clip extraction and export

Usage:
    # Testing mode: uses test video at 38:40-48:40 (10 min), still asks for region selection
    python main.py --testing

    # Full mode: prompts for video file and processes entire video
    python main.py

    # Full mode with specific video
    python main.py --video "full_videos/OSU vs ND 01.20.25.mp4"

    # Use pre-saved regions (auto-detects config based on video name)
    python main.py --video "full_videos/OSU vs ND 01.20.25.mp4" --use-saved-regions

    # Use pre-saved regions from a specific config file
    python main.py --video "full_videos/OSU vs ND 01.20.25.mp4" --use-saved-regions --regions-config "output/my_config.json"

    # Testing mode with saved regions
    python main.py --testing --use-saved-regions
"""

import argparse
import logging
import sys
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

from config import (
    SessionConfig,
    save_session_config,
    load_session_config,
    get_video_basename,
    PROJECT_ROOT,
    OUTPUT_DIR,
    DEFAULT_VIDEO_PATH,
    TESTING_START_TIME,
    TESTING_END_TIME,
    EXPECTED_PLAYS_TESTING,
)
from config.session import parse_time_string
from ui import (
    print_banner,
    select_flag_region,
    select_scorebug_region,
    select_playclock_region,
    select_timeout_region,
    get_video_path_from_user,
)
from ui.api import extract_sample_frames_for_selection
from video import generate_clips
from pipeline import run_extraction, print_results_summary

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)


# =============================================================================
# Argument Parsing
# =============================================================================


def _parse_args() -> argparse.Namespace:
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description="CFB40 Play Detection Pipeline")
    parser.add_argument("--testing", action="store_true", help="Testing mode: use test video at 38:40-48:40 (10 min)")
    parser.add_argument("--video", type=str, help="Path to video file")
    parser.add_argument("--start", type=str, help="Start time (MM:SS or seconds)")
    parser.add_argument("--end", type=str, help="End time (MM:SS or seconds)")
    parser.add_argument("--skip-clips", action="store_true", help="Skip clip generation")
    parser.add_argument(
        "--all-plays-debug",
        action="store_true",
        help="Generate individual play clips in addition to combined video (for debugging)",
    )
    parser.add_argument(
        "--clip-method",
        type=str,
        choices=["stream_copy", "reencode", "ultrafast"],
        default="stream_copy",
        help="Clip extraction method: stream_copy (fastest), reencode (best compression), ultrafast (faster encoding)",
    )
    parser.add_argument(
        "--use-saved-regions",
        action="store_true",
        help="Use pre-existing regions from saved config instead of interactive selection (default: prompt user)",
    )
    parser.add_argument(
        "--regions-config",
        type=str,
        help="Path to specific config file for regions (only used with --use-saved-regions). If not provided, auto-detects based on video name.",
    )
    parser.add_argument(
        "--parallel",
        type=int,
        default=2,
        metavar="N",
        help="Number of parallel workers for detection (default: 2, use 1 for sequential)",
    )

    return parser.parse_args()


# =============================================================================
# Phase 1: Configuration
# =============================================================================


def _phase1_configuration(args: argparse.Namespace) -> Optional[Tuple[str, float, Optional[float], str]]:
    """
    Phase 1: Parse configuration from args and determine video/time settings.

    Args:
        args: Parsed command line arguments.

    Returns:
        Tuple of (video_path, start_time, end_time, video_basename), or None on error.
    """
    print("[Phase 1] Configuration")
    print("-" * 50)

    if args.testing:
        # Testing mode: use default video and segment
        video_path = str(DEFAULT_VIDEO_PATH)
        start_time = TESTING_START_TIME
        end_time = TESTING_END_TIME
        video_basename = get_video_basename(video_path, testing_mode=True)

        print("Mode: Testing (10-minute segment)")
        print(f"Video: {video_path}")
        print(f"Segment: {start_time // 60:.0f}:{start_time % 60:02.0f} - {end_time // 60:.0f}:{end_time % 60:02.0f}")
    else:
        # Full mode: get video from user or args
        if args.video:
            video_path = args.video
        else:
            video_path = get_video_path_from_user(PROJECT_ROOT)

        if not video_path or not Path(video_path).exists():
            logger.error("No valid video path provided")
            return None

        video_basename = get_video_basename(video_path, testing_mode=False)

        # Parse start/end times
        start_time = 0.0
        end_time = None

        if args.start:
            start_time = parse_time_string(args.start)
        if args.end:
            end_time = parse_time_string(args.end)

        print("Mode: Full processing")
        print(f"Video: {video_path}")
        if end_time:
            print(f"Segment: {start_time // 60:.0f}:{int(start_time) % 60:02d} - {end_time // 60:.0f}:{int(end_time) % 60:02d}")
        else:
            print(f"Segment: {start_time // 60:.0f}:{int(start_time) % 60:02d} - end")

    print(f"Output basename: {video_basename}")

    # Verify video exists
    if not Path(video_path).exists():
        logger.error("Video not found: %s", video_path)
        return None

    return video_path, start_time, end_time, video_basename


# =============================================================================
# Phase 2: Region Setup
# =============================================================================


def _phase2_load_saved_regions(
    args: argparse.Namespace,
    video_path: str,
    start_time: float,
    end_time: Optional[float],
    video_basename: str,
) -> Optional[SessionConfig]:
    """
    Phase 2 (saved mode): Load regions from a saved config file.

    Args:
        args: Parsed command line arguments.
        video_path: Path to video file.
        start_time: Start time in seconds.
        end_time: End time in seconds (or None for full video).
        video_basename: Base name for output files.

    Returns:
        SessionConfig with loaded regions, or None on error.
    """
    print("[Phase 2] Loading Saved Regions")
    print("-" * 50)

    # Determine config path
    if args.regions_config:
        config_path = Path(args.regions_config)
    else:
        config_path = OUTPUT_DIR / f"{video_basename}_config.json"

    if not config_path.exists():
        logger.error("Config file not found: %s", config_path)
        logger.error("Run without --use-saved-regions first to create regions, or specify --regions-config")
        return None

    print(f"Loading config from: {config_path}")
    session_config = load_session_config(str(config_path))

    # Update video path and time bounds from command line args (regions stay from config)
    session_config.video_path = video_path
    session_config.start_time = start_time
    session_config.end_time = end_time
    session_config.video_basename = video_basename
    session_config.template_path = str(OUTPUT_DIR / f"{video_basename}_template.png")
    session_config.config_path = str(OUTPUT_DIR / f"{video_basename}_config.json")

    # Print loaded regions
    _print_region_summary(session_config)

    return session_config


def _phase2_interactive_selection(
    video_path: str,
    start_time: float,
    end_time: Optional[float],
    video_basename: str,
) -> Optional[SessionConfig]:
    """
    Phase 2 (interactive mode): Guide user through region selection.

    Args:
        video_path: Path to video file.
        start_time: Start time in seconds.
        end_time: End time in seconds (or None for full video).
        video_basename: Base name for output files.

    Returns:
        SessionConfig with selected regions, or None on error/cancel.
    """
    print("[Phase 2] Interactive Region Selection")
    print("-" * 50)

    # Extract sample frames for region selection
    print("\nExtracting sample frames for region selection...")
    frames = extract_sample_frames_for_selection(video_path, start_time, num_frames=5, interval=2.0)
    if not frames:
        logger.error("Could not extract frames from video")
        return None
    print(f"  Extracted {len(frames)} frames")

    # Select scorebug region
    scorebug_bbox, selected_frame = select_scorebug_region(frames, video_path=video_path)
    if scorebug_bbox is None:
        logger.error("Scorebug region selection cancelled")
        return None

    sb_x, sb_y, sb_w, sb_h = scorebug_bbox
    print(f"Scorebug region: x={sb_x}, y={sb_y}, w={sb_w}, h={sb_h}")

    # Select play clock region (from same frame)
    playclock_region = select_playclock_region(selected_frame, scorebug_bbox)
    if playclock_region is None:
        logger.error("Play clock region selection cancelled")
        return None

    pc_x, pc_y, pc_w, pc_h = playclock_region
    print(f"Play clock region: x_offset={pc_x}, y_offset={pc_y}, w={pc_w}, h={pc_h}")

    # Select timeout regions
    away_timeout_region, home_timeout_region = _select_timeout_regions(selected_frame, scorebug_bbox)

    # Select FLAG region (where "1st & 10" / "FLAG" appears)
    flag_region = _do_flag_region_selection(selected_frame, scorebug_bbox)

    # Create session config
    session_config = SessionConfig(
        video_path=video_path,
        start_time=start_time,
        end_time=end_time,
        scorebug_x=sb_x,
        scorebug_y=sb_y,
        scorebug_width=sb_w,
        scorebug_height=sb_h,
        playclock_x_offset=pc_x,
        playclock_y_offset=pc_y,
        playclock_width=pc_w,
        playclock_height=pc_h,
        home_timeout_x=home_timeout_region[0],
        home_timeout_y=home_timeout_region[1],
        home_timeout_width=home_timeout_region[2],
        home_timeout_height=home_timeout_region[3],
        away_timeout_x=away_timeout_region[0],
        away_timeout_y=away_timeout_region[1],
        away_timeout_width=away_timeout_region[2],
        away_timeout_height=away_timeout_region[3],
        flag_x_offset=flag_region[0],
        flag_y_offset=flag_region[1],
        flag_width=flag_region[2],
        flag_height=flag_region[3],
        template_path=str(OUTPUT_DIR / f"{video_basename}_template.png"),
        config_path=str(OUTPUT_DIR / f"{video_basename}_config.json"),
        video_basename=video_basename,
    )

    # Save config and generate template
    save_session_config(session_config, OUTPUT_DIR, selected_frame=selected_frame)

    return session_config


def _select_timeout_regions(selected_frame, scorebug_bbox) -> Tuple[Tuple[int, int, int, int], Tuple[int, int, int, int]]:
    """
    Select timeout indicator regions for both teams.

    Args:
        selected_frame: Frame image for selection.
        scorebug_bbox: Scorebug bounding box for reference.

    Returns:
        Tuple of (away_timeout_region, home_timeout_region).
    """
    print("\n" + "=" * 60)
    print("Timeout Tracker Setup")
    print("Select the timeout indicator regions (3 stacked ovals per team)")

    # Select away team timeouts (usually on left)
    away_timeout_region = select_timeout_region(selected_frame, scorebug_bbox, "away")
    if away_timeout_region is None:
        logger.warning("Away timeout region selection cancelled - clock reset classification will be limited")
        away_timeout_region = (0, 0, 0, 0)
    else:
        at_x, at_y, at_w, at_h = away_timeout_region
        print(f"Away timeout region: x={at_x}, y={at_y}, w={at_w}, h={at_h}")

    # Select home team timeouts (usually on right)
    home_timeout_region = select_timeout_region(selected_frame, scorebug_bbox, "home")
    if home_timeout_region is None:
        logger.warning("Home timeout region selection cancelled - clock reset classification will be limited")
        home_timeout_region = (0, 0, 0, 0)
    else:
        ht_x, ht_y, ht_w, ht_h = home_timeout_region
        print(f"Home timeout region: x={ht_x}, y={ht_y}, w={ht_w}, h={ht_h}")

    return away_timeout_region, home_timeout_region


def _do_flag_region_selection(selected_frame, scorebug_bbox) -> Tuple[int, int, int, int]:
    """
    Select FLAG indicator region (where '1st & 10' / 'FLAG' appears).

    Args:
        selected_frame: Frame image for selection.
        scorebug_bbox: Scorebug bounding box for reference.

    Returns:
        FLAG region as (x_offset, y_offset, width, height) relative to scorebug.
    """
    print("\n" + "=" * 60)
    print("FLAG Region Setup")
    print("Select where '1st & 10' / 'FLAG' appears on the scorebug")

    flag_region = select_flag_region(selected_frame, scorebug_bbox)
    if flag_region is None:
        logger.warning("FLAG region selection cancelled - FLAG detection will be disabled")
        return (0, 0, 0, 0)

    fl_x, fl_y, fl_w, fl_h = flag_region
    print(f"FLAG region: x_offset={fl_x}, y_offset={fl_y}, w={fl_w}, h={fl_h}")
    return flag_region


def _print_region_summary(config: SessionConfig) -> None:
    """Print a summary of the configured regions."""
    print("Loaded regions:")
    print(f"  Scorebug: x={config.scorebug_x}, y={config.scorebug_y}, w={config.scorebug_width}, h={config.scorebug_height}")
    print(f"  Play clock: x_offset={config.playclock_x_offset}, y_offset={config.playclock_y_offset}, w={config.playclock_width}, h={config.playclock_height}")

    if config.home_timeout_width > 0:
        print(f"  Home timeout: x={config.home_timeout_x}, y={config.home_timeout_y}, w={config.home_timeout_width}, h={config.home_timeout_height}")
    if config.away_timeout_width > 0:
        print(f"  Away timeout: x={config.away_timeout_x}, y={config.away_timeout_y}, w={config.away_timeout_width}, h={config.away_timeout_height}")
    if config.flag_width > 0:
        print(f"  FLAG region: x_offset={config.flag_x_offset}, y_offset={config.flag_y_offset}, w={config.flag_width}, h={config.flag_height}")


# =============================================================================
# Phase 3: Detection
# =============================================================================


def _phase3_extraction(session_config: SessionConfig, num_workers: int) -> Dict[str, Any]:
    """
    Phase 3: Run play extraction on the video.

    Args:
        session_config: Configuration with video path and regions.
        num_workers: Number of parallel workers.

    Returns:
        Extraction results dictionary.
    """
    print("\n" + "=" * 60)
    return run_extraction(session_config, OUTPUT_DIR, num_workers=num_workers)


# =============================================================================
# Phase 4: Clip Generation
# =============================================================================


def _phase4_clips(
    results: Dict[str, Any],
    video_path: str,
    video_basename: str,
    clip_method: str,
    generate_individual: bool,
) -> Dict[str, Any]:
    """
    Phase 4: Generate video clips for detected plays.

    Args:
        results: Detection results with plays.
        video_path: Path to source video.
        video_basename: Base name for output files.
        clip_method: Method for clip extraction.
        generate_individual: Whether to generate individual play clips.

    Returns:
        Clip timing information.
    """
    print("\n" + "=" * 60)
    print("\n[Phase 4] Generating Clips...")
    print("-" * 50)

    return generate_clips(
        plays=results.get("plays", []),
        video_path=video_path,
        output_dir=OUTPUT_DIR,
        video_basename=video_basename,
        clip_method=clip_method,
        generate_individual=generate_individual,
    )


# =============================================================================
# Phase 5: Results Summary
# =============================================================================


def _phase5_summary(
    results: Dict[str, Any],
    testing_mode: bool,
    clip_timing: Dict[str, Any],
    video_basename: str,
    generate_individual: bool,
) -> None:
    """
    Phase 5: Print final results summary.

    Args:
        results: Detection results.
        testing_mode: Whether running in testing mode.
        clip_timing: Timing info from clip generation.
        video_basename: Base name for output files.
        generate_individual: Whether individual clips were generated.
    """
    print_results_summary(
        results=results,
        testing_mode=testing_mode,
        clip_timing=clip_timing,
        video_basename=video_basename,
        generate_individual=generate_individual,
        expected_plays=EXPECTED_PLAYS_TESTING,
    )


# =============================================================================
# Main Entry Point
# =============================================================================


def main() -> int:
    """
    Main entry point for CFB40 play detection pipeline.

    Orchestrates the 5 phases:
    1. Configuration - Parse args, determine video/time settings
    2. Region Setup - Load saved regions or interactive selection
    3. Detection - Run play detection
    4. Clips - Generate video clips
    5. Summary - Print results

    Returns:
        Exit code (0 for success, 1 for error).
    """
    args = _parse_args()
    print_banner()

    # Phase 1: Configuration
    config_result = _phase1_configuration(args)
    if config_result is None:
        return 1
    video_path, start_time, end_time, video_basename = config_result

    # Phase 2: Region Setup
    print("\n" + "=" * 60)
    if args.use_saved_regions:
        session_config = _phase2_load_saved_regions(args, video_path, start_time, end_time, video_basename)
    else:
        session_config = _phase2_interactive_selection(video_path, start_time, end_time, video_basename)

    if session_config is None:
        return 1

    # Phase 3: Extraction
    results = _phase3_extraction(session_config, args.parallel)

    # Phase 4: Clip Generation
    clip_timing = {}
    if not args.skip_clips:
        clip_timing = _phase4_clips(results, video_path, video_basename, args.clip_method, args.all_plays_debug)

    # Phase 5: Results Summary
    _phase5_summary(results, args.testing, clip_timing, video_basename, args.all_plays_debug)

    return 0


if __name__ == "__main__":
    sys.exit(main())