src/detection/scorebug.py

"""
Scorebug detector module.

This module provides functions to detect the presence and location of the scorebug
(score overlay) in video frames.
"""

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

import cv2
import numpy as np

from .models import ScorebugDetection

logger = logging.getLogger(__name__)


class DetectScoreBug:
    """
    Detects the scorebug in video frames.

    The detector supports two modes:
    1. Full-frame search: Template matching across entire frame (slower, use for initial detection)
    2. Fixed-region check: Only check known location for presence (much faster)

    Additionally, split detection can be enabled to handle partial scorebug overlays:
    - When enabled, left and right halves of the scorebug are matched independently
    - Detection passes if either half exceeds the split threshold OR full template exceeds full threshold
    - This helps detect scorebugs when player stats graphics appear on one side

    For optimal performance, use fixed_region mode after determining scorebug location once.
    """

    # Detection thresholds
    FULL_THRESHOLD = 0.5  # Threshold for full-template matching when split detection is enabled
    SPLIT_THRESHOLD = 0.7  # Threshold for half-template matching (left or right)
    LEGACY_THRESHOLD = 0.6  # Original threshold when split detection is disabled

    def __init__(
        self,
        template_path: Optional[str] = None,
        fixed_region: Optional[Tuple[int, int, int, int]] = None,
        fixed_region_config_path: Optional[str] = None,
        use_split_detection: bool = True,
    ):
        """
        Initialize the scorebug detector.

        Args:
            template_path: Path to a template image of the scorebug (optional)
            fixed_region: Fixed region where scorebug appears (x, y, w, h) - enables fast mode
            fixed_region_config_path: Path to JSON config with fixed region (alternative to fixed_region)
            use_split_detection: Enable split-half detection for robustness to partial overlays (default: True)
        """
        self.template: Optional[np.ndarray[Any, Any]] = None
        self.template_path = template_path
        self.fixed_region = fixed_region
        self._use_fixed_region = fixed_region is not None
        self.use_split_detection = use_split_detection

        # Pre-computed template halves for split detection (populated when template is loaded)
        self._template_left: Optional[np.ndarray[Any, Any]] = None
        self._template_right: Optional[np.ndarray[Any, Any]] = None

        if template_path:
            self.load_template(template_path)

        # Load fixed region from config file if provided
        if fixed_region_config_path and not fixed_region:
            self._load_fixed_region_config(fixed_region_config_path)

        mode = "fixed_region" if self._use_fixed_region else "full_search"
        split_mode = "split_detection" if use_split_detection else "full_only"
        logger.info("DetectScoreBug initialized (template: %s, mode: %s, split: %s)", template_path is not None, mode, split_mode)
        if self._use_fixed_region:
            logger.info("  Fixed region: %s", self.fixed_region)

    @property
    def is_fixed_region_mode(self) -> bool:
        """Check if detector is using fixed region mode for faster detection."""
        return self._use_fixed_region

    def _load_fixed_region_config(self, config_path: str) -> None:
        """Load fixed region from a JSON config file."""
        path = Path(config_path)
        if not path.exists():
            logger.warning("Fixed region config not found: %s", config_path)
            return

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

        if "scorebug_region" in data:
            region = data["scorebug_region"]
            self.fixed_region = (region["x"], region["y"], region["width"], region["height"])
            self._use_fixed_region = True
            logger.info("Loaded fixed region from config: %s", self.fixed_region)

    def load_template(self, template_path: str) -> None:
        """
        Load a template image for matching.

        Args:
            template_path: Path to the template image
        """
        self.template = cv2.imread(template_path)
        if self.template is None:
            raise ValueError(f"Could not load template image: {template_path}")

        self.template_path = template_path
        logger.info("Loaded template: %s (size: %dx%d)", template_path, self.template.shape[1], self.template.shape[0])

        # Pre-compute template halves for split detection
        if self.use_split_detection:
            half_width = self.template.shape[1] // 2
            self._template_left = self.template[:, :half_width].copy()
            self._template_right = self.template[:, half_width:].copy()
            logger.info(
                "  Split detection enabled: left half=%dx%d, right half=%dx%d",
                self._template_left.shape[1],
                self._template_left.shape[0],
                self._template_right.shape[1],
                self._template_right.shape[0],
            )

    def detect(self, frame: np.ndarray[Any, Any]) -> ScorebugDetection:
        """
        Detect scorebug in a frame.

        Uses fixed-region mode if configured (much faster), otherwise searches entire frame.
        If a fixed region is set but no template is loaded, assumes scorebug is present
        at the fixed location (useful when coordinates come from user selection).

        In fixed-region mode with template:
        - detected=True (always assume present for play tracking, since we know the location)
        - template_matched=True/False (actual visibility for special play end detection)

        Args:
            frame: Input frame (BGR format)

        Returns:
            ScorebugDetection object with detection results
        """
        # Use fixed-region mode if configured (much faster - only checks known location)
        if self._use_fixed_region and self.fixed_region is not None:
            # If no template loaded, assume scorebug is present at fixed location
            # This is used when coordinates are provided via fixed config (no verification needed)
            if self.template is None:
                x, y, w, h = self.fixed_region
                logger.debug("Fixed region mode without template - assuming scorebug present at %s", self.fixed_region)
                return ScorebugDetection(detected=True, confidence=1.0, bbox=(x, y, w, h), method="fixed_region_assumed")

            detection = self._detect_in_fixed_region(frame)
        else:
            # Need template for full-frame search
            if self.template is None:
                logger.debug("No template loaded, cannot detect scorebug")
                return ScorebugDetection(detected=False, confidence=0.0, method="none")

            # Full-frame template matching (slower, searches entire frame)
            detection = self._detect_by_template_fullsearch(frame)

        if detection.detected:
            logger.debug("Scorebug detected with confidence %.2f using %s", detection.confidence, detection.method)
        else:
            logger.debug("No scorebug detected (confidence: %.2f)", detection.confidence)

        return detection

    # pylint: disable=too-many-locals
    def _detect_in_fixed_region(self, frame: np.ndarray[Any, Any]) -> ScorebugDetection:
        """
        Detect scorebug by checking only the fixed known location.

        This is MUCH faster than full-frame search since we only compare
        the template against a single position.

        In fixed-region mode:
        - detected=True always (we know where the scorebug is, so assume present for play tracking)
        - template_matched=True/False (actual template match result for special play end detection)

        When split detection is enabled:
        - Matches full template AND left/right halves independently
        - Template match passes if: max(left_conf, right_conf) >= SPLIT_THRESHOLD OR full_conf >= FULL_THRESHOLD
        - This handles cases where player stats graphics replace one side of the scorebug

        Args:
            frame: Input frame

        Returns:
            Detection result with detected=True (assumed) and template_matched=actual result
        """
        # Asserts: this method should only be called when fixed_region and template are set
        assert self.fixed_region is not None
        assert self.template is not None

        x, y, _, _ = self.fixed_region
        th, tw = self.template.shape[:2]

        # Validate region bounds
        frame_h, frame_w = frame.shape[:2]
        if x < 0 or y < 0 or x + tw > frame_w or y + th > frame_h:
            logger.warning("Fixed region out of frame bounds")
            # Out of bounds - can't verify template, but still assume present for play tracking
            return ScorebugDetection(detected=True, confidence=0.0, bbox=self.fixed_region, method="fixed_region", template_matched=False)

        # Extract the region where scorebug should be
        region = frame[y : y + th, x : x + tw]

        # Compare full template to region using normalized cross-correlation
        result = cv2.matchTemplate(region, self.template, cv2.TM_CCOEFF_NORMED)
        full_confidence = float(result[0, 0])

        # Use split detection if enabled
        if self.use_split_detection and self._template_left is not None and self._template_right is not None:
            half_width = tw // 2

            # Extract left and right halves of the region
            region_left = region[:, :half_width]
            region_right = region[:, half_width:]

            # Match left half
            left_result = cv2.matchTemplate(region_left, self._template_left, cv2.TM_CCOEFF_NORMED)
            left_confidence = float(left_result[0, 0])

            # Match right half
            right_result = cv2.matchTemplate(region_right, self._template_right, cv2.TM_CCOEFF_NORMED)
            right_confidence = float(right_result[0, 0])

            # Template match passes if:
            # - Either half exceeds SPLIT_THRESHOLD (handles partial overlays)
            # - OR full template exceeds FULL_THRESHOLD
            max_half = max(left_confidence, right_confidence)
            template_matched = max_half >= self.SPLIT_THRESHOLD or full_confidence >= self.FULL_THRESHOLD

            # Use max_half as primary confidence when it's higher (indicates partial overlay)
            effective_confidence = max(full_confidence, max_half)

            logger.debug(
                "Split detection: full=%.3f, left=%.3f, right=%.3f, max_half=%.3f, template_matched=%s",
                full_confidence,
                left_confidence,
                right_confidence,
                max_half,
                template_matched,
            )

            # detected=True (assume present for play tracking), template_matched=actual result
            return ScorebugDetection(
                detected=True,
                confidence=effective_confidence,
                bbox=(x, y, tw, th),
                method="fixed_region_split",
                left_confidence=left_confidence,
                right_confidence=right_confidence,
                template_matched=template_matched,
            )

        # Legacy mode: single threshold on full template
        threshold = self.LEGACY_THRESHOLD
        template_matched = full_confidence >= threshold
        # detected=True (assume present for play tracking), template_matched=actual result
        return ScorebugDetection(detected=True, confidence=full_confidence, bbox=(x, y, tw, th), method="fixed_region", template_matched=template_matched)

    # pylint: disable=too-many-locals
    def _detect_by_template_fullsearch(self, frame: np.ndarray[Any, Any]) -> ScorebugDetection:
        """
        Detect scorebug using full-frame template matching.

        This searches the entire frame for the template - slower but works
        when scorebug position is unknown.

        When split detection is enabled:
        - First finds best full-template match location
        - Then also checks left/right half confidences at that location
        - Detection passes if: max(left_conf, right_conf) >= SPLIT_THRESHOLD OR full_conf >= FULL_THRESHOLD

        Args:
            frame: Input frame

        Returns:
            Detection result
        """
        if self.template is None:
            return ScorebugDetection(detected=False, confidence=0.0, method="full_search")

        # Perform template matching across entire frame
        result = cv2.matchTemplate(frame, self.template, cv2.TM_CCOEFF_NORMED)
        _, max_val, _, max_loc = cv2.minMaxLoc(result)

        # Get bounding box dimensions
        h, w = self.template.shape[:2]
        bbox = (max_loc[0], max_loc[1], w, h)

        # Use split detection if enabled
        if self.use_split_detection and self._template_left is not None and self._template_right is not None:
            x, y = max_loc
            half_width = w // 2

            # Extract region at best match location
            region = frame[y : y + h, x : x + w]

            # Handle edge case where region might be partially out of bounds
            if region.shape[0] != h or region.shape[1] != w:
                # Fall back to full-template match only
                threshold = self.LEGACY_THRESHOLD
                if max_val >= threshold:
                    return ScorebugDetection(detected=True, confidence=float(max_val), bbox=bbox, method="full_search")
                return ScorebugDetection(detected=False, confidence=float(max_val), method="full_search")

            # Extract left and right halves
            region_left = region[:, :half_width]
            region_right = region[:, half_width:]

            # Match left half
            left_result = cv2.matchTemplate(region_left, self._template_left, cv2.TM_CCOEFF_NORMED)
            left_confidence = float(left_result[0, 0])

            # Match right half
            right_result = cv2.matchTemplate(region_right, self._template_right, cv2.TM_CCOEFF_NORMED)
            right_confidence = float(right_result[0, 0])

            # Detection passes if either half exceeds split threshold OR full exceeds full threshold
            max_half = max(left_confidence, right_confidence)
            detected = max_half >= self.SPLIT_THRESHOLD or max_val >= self.FULL_THRESHOLD

            # Use max_half as primary confidence when it's higher
            effective_confidence = max(float(max_val), max_half)

            logger.debug(
                "Split detection (fullsearch): full=%.3f, left=%.3f, right=%.3f, max_half=%.3f, detected=%s", max_val, left_confidence, right_confidence, max_half, detected
            )

            return ScorebugDetection(
                detected=detected, confidence=effective_confidence, bbox=bbox, method="full_search_split", left_confidence=left_confidence, right_confidence=right_confidence
            )

        # Legacy mode: single threshold on full template
        threshold = self.LEGACY_THRESHOLD
        if max_val >= threshold:
            return ScorebugDetection(detected=True, confidence=float(max_val), bbox=bbox, method="full_search")
        return ScorebugDetection(detected=False, confidence=float(max_val), method="full_search")

    def set_fixed_region(self, region: Tuple[int, int, int, int]) -> None:
        """
        Set a fixed region for fast detection mode.

        Call this after discovering the scorebug location to switch to fast mode.

        Args:
            region: (x, y, width, height) of the scorebug location
        """
        self.fixed_region = region
        self._use_fixed_region = True
        logger.info("Fixed region set: %s - now using fast detection mode", region)

    def save_fixed_region_config(self, config_path: str) -> None:
        """Save the fixed region to a config file for reuse."""
        if self.fixed_region is None:
            logger.warning("No fixed region to save")
            return

        x, y, w, h = self.fixed_region
        data = {"scorebug_region": {"x": x, "y": y, "width": w, "height": h}}

        path = Path(config_path)
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)

        logger.info("Saved fixed region config to: %s", config_path)

    def discover_and_lock_region(self, frame: np.ndarray[Any, Any]) -> bool:
        """
        Discover scorebug location using full search, then lock to fixed region mode.

        This is useful for the first frame - find the scorebug once, then use
        fast fixed-region mode for all subsequent frames.

        Args:
            frame: Frame to search

        Returns:
            True if scorebug was found and region was locked, False otherwise
        """
        # Temporarily disable fixed region to do full search
        old_use_fixed = self._use_fixed_region
        self._use_fixed_region = False

        detection = self._detect_by_template_fullsearch(frame)

        if detection.detected and detection.bbox:
            self.set_fixed_region(detection.bbox)
            return True
        self._use_fixed_region = old_use_fixed
        return False

    def visualize_detection(self, frame: np.ndarray[Any, Any], detection: ScorebugDetection) -> np.ndarray[Any, Any]:
        """
        Draw detection results on frame for visualization.

        Args:
            frame: Input frame
            detection: Detection result

        Returns:
            Frame with visualization overlay
        """
        vis_frame = frame.copy()

        if detection.detected and detection.bbox:
            x, y, w, h = detection.bbox

            # Draw bounding box
            color = (0, 255, 0)  # Green for detected
            cv2.rectangle(vis_frame, (x, y), (x + w, y + h), color, 2)

            # Add confidence text
            text = f"{detection.method}: {detection.confidence:.2f}"
            cv2.putText(vis_frame, text, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

        return vis_frame