scripts/test_flag_region_selection.py

#!/usr/bin/env python
"""
Test script for FLAG region selection (Step 1 of FLAG detection implementation).

This script allows interactive selection of the FLAG indicator region on the scorebug.
The FLAG region is where "1st & 10" / "FLAG" text appears on the scorebug.

Usage:
    # Use default Tennessee video
    python scripts/test_flag_region_selection.py

    # Use a specific video
    python scripts/test_flag_region_selection.py --video "full_videos/OSU vs Texas 01.10.25.mkv"

The script will:
1. Load sample frames from the specified video
2. Display the frame with the existing scorebug region highlighted
3. Allow user to click/drag to select the FLAG region
4. Save the selected region to output/{video_basename}_flag_config.json
5. Optionally update the session config if it exists
6. Display a cropped preview of the selected region
"""

import json
import logging
import sys
from pathlib import Path
from typing import Any, Optional, Tuple

import cv2
import numpy as np

# Add src to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

from ui.models import BBox
from ui.selector import RegionSelector
from video.frame_extractor import extract_sample_frames

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

# Project paths
PROJECT_ROOT = Path(__file__).parent.parent
DATA_CONFIG_DIR = PROJECT_ROOT / "data" / "config"
OUTPUT_DIR = PROJECT_ROOT / "output"
FULL_VIDEOS_DIR = PROJECT_ROOT / "full_videos"

# Default test video
DEFAULT_VIDEO = FULL_VIDEOS_DIR / "OSU vs Tenn 12.21.24.mkv"
DEFAULT_START_TIME = 38 * 60 + 40  # 38:40 - where test segment starts


class FlagRegionSelectionSession:
    """Interactive session for selecting the FLAG indicator region with padding around scorebug."""

    def __init__(
        self,
        frame: np.ndarray[Any, Any],
        scorebug_bbox: BBox,
        scale_factor: int = 3,
        padding: int = 100,
    ):
        """
        Initialize FLAG region selection session.

        Args:
            frame: Full video frame (BGR)
            scorebug_bbox: Bounding box of the scorebug in the frame
            scale_factor: How much to scale up the region for easier selection
            padding: Pixels of padding around the scorebug to include in the view
        """
        self.full_frame = frame
        self.scorebug_bbox = scorebug_bbox
        self.scale_factor = scale_factor
        self.padding = padding
        self.window_name = "Select FLAG Region (1st & 10 location)"

        # Calculate padded region bounds (with boundary checks)
        sb = scorebug_bbox
        self.pad_x1 = max(0, sb.x - padding)
        self.pad_y1 = max(0, sb.y - padding)
        self.pad_x2 = min(frame.shape[1], sb.x2 + padding)
        self.pad_y2 = min(frame.shape[0], sb.y2 + padding)

        # Scorebug position within the padded region (for drawing reference outline)
        self.scorebug_in_padded = BBox(
            x=sb.x - self.pad_x1,
            y=sb.y - self.pad_y1,
            width=sb.width,
            height=sb.height,
        )

        # Extract the padded region from the frame
        self.padded_region = frame[self.pad_y1 : self.pad_y2, self.pad_x1 : self.pad_x2].copy()
        padded_w, padded_h = self.padded_region.shape[1], self.padded_region.shape[0]

        # Scale up the padded region for easier selection
        self.scaled_region = cv2.resize(
            self.padded_region,
            (padded_w * scale_factor, padded_h * scale_factor),
            interpolation=cv2.INTER_LINEAR,
        )

        # Setup selector
        self.selector = RegionSelector(self.window_name, mode="two_click")

        # State
        self.should_quit = False
        self.should_confirm = False

    def _render_frame(self, display_frame: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]:
        """Add overlays and instructions to the display frame."""
        # Draw scorebug boundary for reference (gray dashed appearance)
        sb_scaled = self.scorebug_in_padded.scaled(self.scale_factor)
        cv2.rectangle(
            display_frame,
            (sb_scaled.x, sb_scaled.y),
            (sb_scaled.x2, sb_scaled.y2),
            (128, 128, 128),
            2,
        )
        # Label the scorebug region
        cv2.putText(
            display_frame,
            "Scorebug",
            (sb_scaled.x, sb_scaled.y - 5),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (128, 128, 128),
            1,
        )

        # Draw clicked points
        for i, point in enumerate(self.selector.points):
            color = (0, 255, 0) if i == 0 else (0, 0, 255)  # Green for first, red for second
            cv2.circle(display_frame, point, 5, color, -1)
            label = "Top-Left" if i == 0 else "Bottom-Right"
            cv2.putText(display_frame, label, (point[0] + 10, point[1]), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

        # Draw preview rectangle when dragging
        if len(self.selector.points) == 1 and self.selector.current_point:
            cv2.rectangle(display_frame, self.selector.points[0], self.selector.current_point, (0, 255, 255), 2)

        # Draw final selection rectangle
        if self.selector.selection_complete:
            bbox = self.selector.get_bbox()
            if bbox:
                # Yellow highlight for FLAG region
                cv2.rectangle(display_frame, (bbox.x, bbox.y), (bbox.x2, bbox.y2), (0, 255, 255), 3)

        # Instructions at top
        if len(self.selector.points) == 0:
            instruction = "Step 1: Click TOP-LEFT corner of FLAG/down-distance area"
            color = (0, 255, 255)
        elif len(self.selector.points) == 1:
            instruction = "Step 2: Click BOTTOM-RIGHT corner of FLAG/down-distance area"
            color = (0, 255, 255)
        else:
            instruction = "Press ENTER to confirm, 'r' to reset, 'q' to quit"
            color = (0, 255, 0)

        cv2.putText(display_frame, instruction, (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

        return display_frame

    def run(self) -> Optional[BBox]:
        """
        Run the interactive selection loop.

        Returns:
            BBox in scaled coordinates, or None if cancelled
        """
        cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
        cv2.resizeWindow(
            self.window_name,
            self.scaled_region.shape[1],
            self.scaled_region.shape[0] + 50,
        )
        cv2.setMouseCallback(self.window_name, self.selector.mouse_callback)

        while True:
            # Get display frame
            display_frame = self.scaled_region.copy()
            display_frame = self._render_frame(display_frame)

            cv2.imshow(self.window_name, display_frame)
            key = cv2.waitKey(30) & 0xFF

            # Handle key presses
            if key == ord("q"):
                self.should_quit = True
            elif key == ord("r"):
                self.selector.reset()
            elif key == 13:  # Enter key
                if self.selector.selection_complete:
                    self.should_confirm = True

            # Check exit conditions
            if self.should_quit:
                cv2.destroyAllWindows()
                return None

            if self.should_confirm:
                bbox = self.selector.get_bbox()
                cv2.destroyAllWindows()
                return bbox

        return None

    def get_scorebug_relative_bbox(self) -> Optional[BBox]:
        """
        Get the selection converted to coordinates relative to the scorebug.

        The selection is made in the scaled padded region. This method:
        1. Unscales to original pixel coordinates
        2. Converts from padded-region coordinates to scorebug-relative coordinates

        Returns:
            BBox with x_offset, y_offset relative to the scorebug top-left corner
        """
        bbox = self.selector.get_bbox()
        if bbox is None:
            return None

        # Unscale from the zoomed view
        unscaled = bbox.unscaled(self.scale_factor)

        # Convert from padded-region coords to scorebug-relative coords
        # The scorebug starts at self.scorebug_in_padded within the padded region
        scorebug_relative = BBox(
            x=unscaled.x - self.scorebug_in_padded.x,
            y=unscaled.y - self.scorebug_in_padded.y,
            width=unscaled.width,
            height=unscaled.height,
        )

        return scorebug_relative

    def get_absolute_bbox(self) -> Optional[BBox]:
        """
        Get the selection in absolute frame coordinates.

        Returns:
            BBox in absolute frame coordinates
        """
        bbox = self.selector.get_bbox()
        if bbox is None:
            return None

        # Unscale from the zoomed view
        unscaled = bbox.unscaled(self.scale_factor)

        # Add the padded region offset to get absolute coordinates
        absolute = BBox(
            x=unscaled.x + self.pad_x1,
            y=unscaled.y + self.pad_y1,
            width=unscaled.width,
            height=unscaled.height,
        )

        return absolute


def load_saved_scorebug_config() -> Optional[Tuple[BBox, str]]:
    """
    Load the scorebug region from an existing config file.

    Returns:
        Tuple of (scorebug_bbox, template_path) or None if not found
    """
    # Look for any existing config file in output/
    config_files = list(OUTPUT_DIR.glob("*_config.json"))

    # Filter out playclock and timeout configs
    main_configs = [f for f in config_files if "playclock" not in f.name and "timeout" not in f.name]

    if not main_configs:
        logger.warning("No existing config files found in output/")
        return None

    # Use the most recent one
    config_path = max(main_configs, key=lambda p: p.stat().st_mtime)
    logger.info("Loading scorebug config from: %s", config_path)

    with open(config_path, "r", encoding="utf-8") as f:
        config = json.load(f)

    scorebug_bbox = BBox(
        x=config["scorebug_x"],
        y=config["scorebug_y"],
        width=config["scorebug_width"],
        height=config["scorebug_height"],
    )

    template_path = config.get("template_path", "")

    return scorebug_bbox, template_path


def show_preview(frame: np.ndarray[Any, Any], flag_region_bbox: BBox, scorebug_bbox: BBox) -> None:
    """
    Display a preview of the selected FLAG region.

    Args:
        frame: Full video frame
        flag_region_bbox: Selected FLAG region (relative to scorebug)
        scorebug_bbox: Scorebug bounding box
    """
    # Calculate absolute coordinates
    abs_x = scorebug_bbox.x + flag_region_bbox.x
    abs_y = scorebug_bbox.y + flag_region_bbox.y
    abs_x2 = abs_x + flag_region_bbox.width
    abs_y2 = abs_y + flag_region_bbox.height

    # Extract the FLAG region
    flag_roi = frame[abs_y:abs_y2, abs_x:abs_x2].copy()

    # Scale up for visibility
    scale = 4
    preview = cv2.resize(flag_roi, (flag_roi.shape[1] * scale, flag_roi.shape[0] * scale), interpolation=cv2.INTER_LINEAR)

    # Also show the full frame with the region highlighted
    frame_with_highlight = frame.copy()
    # Draw scorebug outline (gray)
    cv2.rectangle(
        frame_with_highlight,
        (scorebug_bbox.x, scorebug_bbox.y),
        (scorebug_bbox.x2, scorebug_bbox.y2),
        (128, 128, 128),
        2,
    )
    # Draw FLAG region outline (yellow)
    cv2.rectangle(
        frame_with_highlight,
        (abs_x, abs_y),
        (abs_x2, abs_y2),
        (0, 255, 255),
        3,
    )

    # Add labels
    cv2.putText(frame_with_highlight, "Scorebug", (scorebug_bbox.x, scorebug_bbox.y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (128, 128, 128), 2)
    cv2.putText(frame_with_highlight, "FLAG Region", (abs_x, abs_y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)

    # Show both windows
    cv2.namedWindow("FLAG Region Preview (4x scaled)", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("FLAG Region Preview (4x scaled)", preview.shape[1], preview.shape[0])
    cv2.imshow("FLAG Region Preview (4x scaled)", preview)

    cv2.namedWindow("Full Frame with FLAG Region", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("Full Frame with FLAG Region", 1280, 720)
    cv2.imshow("Full Frame with FLAG Region", frame_with_highlight)

    print("\n" + "=" * 50)
    print("FLAG Region Preview")
    print("=" * 50)
    print("  Relative to scorebug:")
    print(f"    x_offset: {flag_region_bbox.x}")
    print(f"    y_offset: {flag_region_bbox.y}")
    print(f"    width:    {flag_region_bbox.width}")
    print(f"    height:   {flag_region_bbox.height}")
    print("  Absolute in frame:")
    print(f"    x: {abs_x}, y: {abs_y}")
    print(f"    width: {flag_region_bbox.width}, height: {flag_region_bbox.height}")
    print("=" * 50)
    print("Press any key to close preview...")

    cv2.waitKey(0)
    cv2.destroyAllWindows()


def get_video_basename(video_path: str) -> str:
    """Get a clean basename from video path for config naming."""
    basename = Path(video_path).stem
    for char in [" ", ".", "-"]:
        basename = basename.replace(char, "_")
    while "__" in basename:
        basename = basename.replace("__", "_")
    return basename.strip("_")


def save_flag_region(flag_bbox: BBox, source_video: str, scorebug_template: str, video_path: str) -> Path:
    """
    Save the FLAG region configuration to video-specific JSON file.

    Also updates the session config if it exists.

    Args:
        flag_bbox: FLAG region bounding box (relative to scorebug)
        source_video: Name of the source video
        scorebug_template: Name of the scorebug template file
        video_path: Full path to the video (used for naming)

    Returns:
        Path to the saved config file
    """
    video_basename = get_video_basename(video_path)

    config = {
        "flag_region": {
            "x_offset": flag_bbox.x,
            "y_offset": flag_bbox.y,
            "width": flag_bbox.width,
            "height": flag_bbox.height,
        },
        "source_video": source_video,
        "scorebug_template": scorebug_template,
    }

    # Save to video-specific file in output directory
    output_path = OUTPUT_DIR / f"{video_basename}_flag_config.json"
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(config, f, indent=2)

    logger.info("Saved FLAG region config to: %s", output_path)

    # Also update session config if it exists
    session_config_path = OUTPUT_DIR / f"{video_basename}_config.json"
    if session_config_path.exists():
        try:
            with open(session_config_path, "r", encoding="utf-8") as f:
                session_config = json.load(f)

            session_config["flag_x_offset"] = flag_bbox.x
            session_config["flag_y_offset"] = flag_bbox.y
            session_config["flag_width"] = flag_bbox.width
            session_config["flag_height"] = flag_bbox.height

            with open(session_config_path, "w", encoding="utf-8") as f:
                json.dump(session_config, f, indent=2)

            logger.info("Updated session config: %s", session_config_path)
        except Exception as e:
            logger.warning("Could not update session config: %s", e)

    return output_path


def print_banner() -> None:
    """Print the script banner and introduction."""
    print("\n" + "=" * 60)
    print("   FLAG Region Selection Test")
    print("=" * 60)
    print("\nThis script helps select the FLAG/down-distance region on the scorebug.")
    print("The FLAG indicator appears where '1st & 10' normally shows.")
    print("-" * 60)


def print_instructions() -> None:
    """Print selection instructions and wait for user."""
    print("\n" + "-" * 60)
    print("INSTRUCTIONS:")
    print("  1. The scorebug region will be shown (scaled 3x)")
    print("  2. Click the TOP-LEFT corner of the down/distance area")
    print("     (where '1st & 10' or 'FLAG' appears)")
    print("  3. Click the BOTTOM-RIGHT corner")
    print("  4. Press ENTER to confirm, 'r' to reset, 'q' to quit")
    print("-" * 60)
    input("\nPress Enter to start selection...")


def parse_args():
    """Parse command line arguments."""
    import argparse

    parser = argparse.ArgumentParser(description="Select FLAG region for a video")
    parser.add_argument(
        "--video",
        type=str,
        default=str(DEFAULT_VIDEO),
        help="Path to video file (default: OSU vs Tenn 12.21.24.mkv)",
    )
    parser.add_argument(
        "--start-time",
        type=float,
        default=DEFAULT_START_TIME,
        help="Start time in seconds for frame extraction (default: 38:40)",
    )
    return parser.parse_args()


def load_scorebug_config_for_video(video_path: str) -> Optional[Tuple[BBox, str]]:
    """Load scorebug config for a specific video."""
    video_basename = get_video_basename(video_path)
    session_config_path = OUTPUT_DIR / f"{video_basename}_config.json"

    if session_config_path.exists():
        with open(session_config_path, "r", encoding="utf-8") as f:
            config = json.load(f)

        scorebug_bbox = BBox(
            x=config["scorebug_x"],
            y=config["scorebug_y"],
            width=config["scorebug_width"],
            height=config["scorebug_height"],
        )
        template_path = config.get("template_path", "")
        logger.info("Loaded session config from: %s", session_config_path)
        return scorebug_bbox, template_path

    # Fall back to generic config
    logger.warning("No session config found for video, trying generic config...")
    return load_saved_scorebug_config()


def main() -> int:
    """Main entry point for FLAG region selection test."""
    args = parse_args()
    print_banner()

    # Determine video path
    video_path = Path(args.video)
    if not video_path.exists():
        # Try relative to project root
        video_path = PROJECT_ROOT / args.video
        if not video_path.exists():
            print(f"\nERROR: Video not found: {args.video}")
            return 1

    print(f"Using video: {video_path.name}")
    video_basename = get_video_basename(str(video_path))
    print(f"Config basename: {video_basename}")

    # Load scorebug config for this video
    result = load_scorebug_config_for_video(str(video_path))
    if result is None:
        print("\nERROR: No existing scorebug config found.")
        print("Please run main.py first to set up the scorebug region for this video.")
        return 1

    scorebug_bbox, template_path = result
    print(f"\nLoaded scorebug region: {scorebug_bbox.to_tuple()}")

    start_time = args.start_time
    print(f"Starting at: {int(start_time) // 60}:{int(start_time) % 60:02d}")

    # Extract sample frames
    print("\nExtracting sample frames...")
    frames = extract_sample_frames(str(video_path), start_time, num_frames=1, interval=0.0)

    if not frames:
        print("ERROR: Failed to extract frames from video")
        return 1

    timestamp, frame = frames[0]
    print(f"Loaded frame at {timestamp:.1f}s")

    # Show instructions and start selection
    print_instructions()

    # Run selection session (with 100px padding around scorebug)
    session = FlagRegionSelectionSession(frame, scorebug_bbox, scale_factor=3, padding=100)
    scaled_bbox = session.run()

    if scaled_bbox is None:
        print("\nSelection cancelled.")
        return 0

    # Get bbox relative to scorebug (for config) and absolute (for preview)
    flag_bbox = session.get_scorebug_relative_bbox()
    if flag_bbox is None:
        print("\nNo valid selection made.")
        return 0

    print(f"\nSelected FLAG region (relative to scorebug): {flag_bbox.to_tuple()}")

    # Show preview
    show_preview(frame, flag_bbox, scorebug_bbox)

    # Ask to save
    print("\nSave this FLAG region? (y/n)")
    response = input("> ").strip().lower()

    if response == "y":
        # Extract video name and template name for metadata
        video_name = video_path.name
        template_name = Path(template_path).name if template_path else "unknown"

        save_path = save_flag_region(flag_bbox, video_name, template_name, str(video_path))
        print(f"\n✓ FLAG region saved to: {save_path}")
    else:
        print("\nSelection not saved.")

    return 0


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