scripts/test_texas_yellow_digits.py

"""
Test if yellow-tinted play clock digits are causing misreads in Texas video.

Compares template matching results against OCR ground truth to identify
any systematic misreads due to the yellow color not normalizing correctly.

Usage:
    python scripts/test_texas_yellow_digits.py
"""

import json
import logging
import sys
from pathlib import Path
from typing import Any, Dict, List, Tuple

import cv2
import numpy as np

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

from readers import ReadPlayClock
from setup import DigitTemplateLibrary
from detection.scorebug import DetectScoreBug
from utils.regions import preprocess_playclock_region
from utils.color import normalize_to_grayscale, detect_red_digits

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

# Lazy import easyocr to avoid slow startup
_ocr_reader = None


def get_ocr_reader():
    """Lazy-load EasyOCR reader."""
    global _ocr_reader
    if _ocr_reader is None:
        import easyocr
        _ocr_reader = easyocr.Reader(["en"], gpu=False, verbose=False)
    return _ocr_reader


def seconds_to_timestamp(seconds: float) -> str:
    """Convert seconds to timestamp string (H:MM:SS)."""
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = int(seconds % 60)
    if hours > 0:
        return f"{hours}:{minutes:02d}:{secs:02d}"
    return f"{minutes}:{secs:02d}"


def load_texas_config() -> Dict[str, Any]:
    """Load the saved config for Texas video."""
    config_path = Path("output/OSU_vs_Texas_01_10_25_config.json")
    with open(config_path, "r") as f:
        return json.load(f)


def read_clock_with_ocr(region: np.ndarray) -> Tuple[bool, int, float, str]:
    """
    Read play clock value using OCR.
    
    Returns:
        Tuple of (detected, value, confidence, raw_text)
    """
    reader = get_ocr_reader()
    
    # Scale up for better OCR
    scaled = cv2.resize(region, None, fx=4, fy=4, interpolation=cv2.INTER_LINEAR)
    
    # Try OCR on the scaled image
    results = reader.readtext(scaled, allowlist="0123456789")
    
    if not results:
        return False, -1, 0.0, ""
    
    # Get the best result
    best_result = max(results, key=lambda x: x[2])
    raw_text = best_result[1]
    confidence = best_result[2]
    
    # Parse the text
    try:
        value = int(raw_text)
        if 0 <= value <= 40:
            return True, value, confidence, raw_text
    except ValueError:
        pass
    
    return False, -1, confidence, raw_text


def analyze_color_distribution(region: np.ndarray) -> Dict[str, Any]:
    """
    Analyze the color distribution in the play clock region.
    
    Returns info about whether digits appear white, yellow, or red.
    """
    # Convert to HSV for color analysis
    hsv = cv2.cvtColor(region, cv2.COLOR_BGR2HSV)
    
    # Get the brightest pixels (likely the digits)
    gray = cv2.cvtColor(region, cv2.COLOR_BGR2GRAY)
    bright_mask = gray > 150  # Threshold for bright pixels
    
    if not np.any(bright_mask):
        return {"has_bright_pixels": False}
    
    # Analyze hue of bright pixels
    bright_hues = hsv[:, :, 0][bright_mask]
    bright_sats = hsv[:, :, 1][bright_mask]
    
    mean_hue = np.mean(bright_hues)
    mean_sat = np.mean(bright_sats)
    
    # Classify color
    # White: low saturation
    # Yellow: hue ~20-40, moderate saturation
    # Red: hue ~0-10 or ~170-180, high saturation
    is_red = detect_red_digits(region)
    
    if mean_sat < 30:
        color_class = "white"
    elif 15 <= mean_hue <= 45 and mean_sat > 30:
        color_class = "yellow"
    elif (mean_hue < 15 or mean_hue > 165) and mean_sat > 50:
        color_class = "red"
    else:
        color_class = "other"
    
    return {
        "has_bright_pixels": True,
        "mean_hue": float(mean_hue),
        "mean_saturation": float(mean_sat),
        "color_class": color_class,
        "is_red_detected": is_red,
    }


def test_template_vs_ocr(
    video_path: str,
    playclock_coords: Tuple[int, int, int, int],
    scorebug_coords: Tuple[int, int, int, int],
    template_dir: str,
    template_path: str,
    timestamps: List[float],
    output_dir: str,
) -> List[Dict[str, Any]]:
    """
    Compare template matching results against OCR ground truth.
    """
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    # Load template library
    library = DigitTemplateLibrary()
    if not library.load(template_dir):
        raise RuntimeError(f"Failed to load templates from {template_dir}")

    # Create play clock reader
    pc_w, pc_h = playclock_coords[2], playclock_coords[3]
    reader = ReadPlayClock(library, region_width=pc_w, region_height=pc_h)

    # Create scorebug detector
    scorebug_detector = DetectScoreBug(
        template_path=template_path,
        fixed_region=scorebug_coords,
        use_split_detection=True,
    )

    # Open video
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError(f"Failed to open video: {video_path}")

    fps = cap.get(cv2.CAP_PROP_FPS)
    results = []

    logger.info("Comparing template matching vs OCR on %d timestamps...", len(timestamps))
    logger.info("(OCR is slow, this may take a minute)")

    for i, ts in enumerate(timestamps):
        frame_num = int(ts * fps)
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_num)
        ret, frame = cap.read()

        if not ret or frame is None:
            continue

        # Check scorebug
        sb_result = scorebug_detector.detect(frame)
        if not sb_result.detected:
            continue

        # Extract play clock region
        pc_x, pc_y, pc_w, pc_h = playclock_coords
        playclock_region = frame[pc_y : pc_y + pc_h, pc_x : pc_x + pc_w].copy()

        # Analyze color
        color_info = analyze_color_distribution(playclock_region)

        # Template matching (with padding=10 as currently used)
        tmpl_result = reader.read_from_fixed_location(frame, playclock_coords, padding=10)

        # OCR ground truth
        ocr_detected, ocr_value, ocr_conf, ocr_raw = read_clock_with_ocr(playclock_region)

        # Compare results
        match = False
        if tmpl_result.detected and ocr_detected:
            match = (tmpl_result.value == ocr_value)

        result = {
            "timestamp": ts,
            "timestamp_str": seconds_to_timestamp(ts),
            "color_info": color_info,
            "template_detected": tmpl_result.detected,
            "template_value": tmpl_result.value,
            "template_confidence": tmpl_result.confidence,
            "ocr_detected": ocr_detected,
            "ocr_value": ocr_value,
            "ocr_confidence": ocr_conf,
            "ocr_raw": ocr_raw,
            "match": match,
            "mismatch": tmpl_result.detected and ocr_detected and not match,
        }
        results.append(result)

        # Save debug image for mismatches
        if result["mismatch"]:
            save_mismatch_image(
                frame,
                playclock_region,
                playclock_coords,
                scorebug_coords,
                result,
                output_path / f"mismatch_{i:02d}_t{int(ts)}_tmpl{tmpl_result.value}_ocr{ocr_value}.png",
            )

        status = "MATCH" if match else ("MISMATCH" if result["mismatch"] else "incomplete")
        logger.info(
            "[%d/%d] %s: tmpl=%s, ocr=%s, color=%s [%s]",
            i + 1,
            len(timestamps),
            seconds_to_timestamp(ts),
            tmpl_result.value if tmpl_result.detected else "X",
            ocr_value if ocr_detected else "X",
            color_info.get("color_class", "?"),
            status,
        )

    cap.release()
    return results


def save_mismatch_image(
    frame: np.ndarray,
    playclock_region: np.ndarray,
    playclock_coords: Tuple[int, int, int, int],
    scorebug_coords: Tuple[int, int, int, int],
    result: Dict[str, Any],
    output_path: Path,
) -> None:
    """Save a debug image for a mismatch case."""
    pc_x, pc_y, pc_w, pc_h = playclock_coords
    sb_x, sb_y, sb_w, sb_h = scorebug_coords

    # Extract scorebug area
    margin = 30
    crop_y1 = max(0, sb_y - margin)
    crop_y2 = min(frame.shape[0], sb_y + sb_h + margin)
    crop_x1 = max(0, sb_x - margin)
    crop_x2 = min(frame.shape[1], sb_x + sb_w + margin)
    scorebug_crop = frame[crop_y1:crop_y2, crop_x1:crop_x2].copy()

    # Draw play clock region
    pc_rel_x = pc_x - crop_x1
    pc_rel_y = pc_y - crop_y1
    cv2.rectangle(scorebug_crop, (pc_rel_x, pc_rel_y), (pc_rel_x + pc_w, pc_rel_y + pc_h), (0, 0, 255), 2)

    # Scale images
    pc_scaled = cv2.resize(playclock_region, None, fx=12, fy=12, interpolation=cv2.INTER_NEAREST)
    preprocessed = preprocess_playclock_region(playclock_region, scale_factor=4)
    preprocessed_bgr = cv2.cvtColor(preprocessed, cv2.COLOR_GRAY2BGR)
    preprocessed_scaled = cv2.resize(preprocessed_bgr, (pc_scaled.shape[1], pc_scaled.shape[0]), interpolation=cv2.INTER_NEAREST)
    scorebug_scaled = cv2.resize(scorebug_crop, None, fx=2, fy=2)

    # Create composite
    composite_height = max(scorebug_scaled.shape[0], pc_scaled.shape[0] + preprocessed_scaled.shape[0] + 20)
    composite_width = scorebug_scaled.shape[1] + max(pc_scaled.shape[1], preprocessed_scaled.shape[1]) + 40
    composite = np.zeros((composite_height + 80, composite_width, 3), dtype=np.uint8)

    composite[0 : scorebug_scaled.shape[0], 0 : scorebug_scaled.shape[1]] = scorebug_scaled
    x_offset = scorebug_scaled.shape[1] + 20
    composite[0 : pc_scaled.shape[0], x_offset : x_offset + pc_scaled.shape[1]] = pc_scaled
    y_offset = pc_scaled.shape[0] + 10
    composite[y_offset : y_offset + preprocessed_scaled.shape[0], x_offset : x_offset + preprocessed_scaled.shape[1]] = preprocessed_scaled

    # Labels
    font = cv2.FONT_HERSHEY_SIMPLEX
    cv2.putText(composite, f"MISMATCH at {result['timestamp_str']}", (10, composite_height + 25), font, 0.6, (0, 0, 255), 1)
    cv2.putText(composite, f"Template: {result['template_value']} (conf={result['template_confidence']:.2f})", (10, composite_height + 50), font, 0.5, (255, 255, 255), 1)
    cv2.putText(composite, f"OCR: {result['ocr_value']} (conf={result['ocr_confidence']:.2f}, raw='{result['ocr_raw']}')", (10, composite_height + 70), font, 0.5, (0, 255, 0), 1)

    cv2.imwrite(str(output_path), composite)


def main():
    """Main function to test yellow digit handling."""
    config = load_texas_config()

    video_path = config["video_path"]
    template_path = config["template_path"]

    playclock_coords = (
        config["scorebug_x"] + config["playclock_x_offset"],
        config["scorebug_y"] + config["playclock_y_offset"],
        config["playclock_width"],
        config["playclock_height"],
    )

    scorebug_coords = (
        config["scorebug_x"],
        config["scorebug_y"],
        config["scorebug_width"],
        config["scorebug_height"],
    )

    logger.info("=" * 70)
    logger.info("TEXAS YELLOW DIGIT COLOR TEST")
    logger.info("=" * 70)
    logger.info("Video: %s", video_path)
    logger.info("")

    # Sample timestamps throughout the video
    # Mix of early game, mid game, and late game (where short plays were detected)
    timestamps = [
        # Early game
        500.0, 600.0, 700.0, 800.0, 900.0,
        # Mid game
        2000.0, 2100.0, 2200.0, 2300.0, 2400.0,
        # Late game (around short plays)
        5930.0, 5935.0, 5940.0, 5945.0, 5950.0,
        5955.0, 5960.0, 5965.0, 5970.0, 5975.0,
    ]

    results = test_template_vs_ocr(
        video_path=video_path,
        playclock_coords=playclock_coords,
        scorebug_coords=scorebug_coords,
        template_dir="output/debug/digit_templates",
        template_path=template_path,
        timestamps=timestamps,
        output_dir="output/debug/texas_yellow_test",
    )

    # Summary
    logger.info("")
    logger.info("=" * 70)
    logger.info("SUMMARY")
    logger.info("=" * 70)

    total = len(results)
    both_detected = sum(1 for r in results if r["template_detected"] and r["ocr_detected"])
    matches = sum(1 for r in results if r["match"])
    mismatches = sum(1 for r in results if r["mismatch"])

    # Color distribution
    color_counts = {}
    for r in results:
        color = r["color_info"].get("color_class", "unknown")
        color_counts[color] = color_counts.get(color, 0) + 1

    logger.info("Total samples: %d", total)
    logger.info("Both detected: %d (%.1f%%)", both_detected, 100 * both_detected / total if total > 0 else 0)
    logger.info("Matches: %d (%.1f%% of both detected)", matches, 100 * matches / both_detected if both_detected > 0 else 0)
    logger.info("MISMATCHES: %d", mismatches)
    logger.info("")
    logger.info("Color distribution:")
    for color, count in sorted(color_counts.items()):
        logger.info("  %s: %d (%.1f%%)", color, count, 100 * count / total if total > 0 else 0)

    # Check if mismatches correlate with color
    if mismatches > 0:
        logger.info("")
        logger.info("MISMATCH DETAILS:")
        for r in results:
            if r["mismatch"]:
                logger.info(
                    "  %s: tmpl=%s, ocr=%s, color=%s, hue=%.1f, sat=%.1f",
                    r["timestamp_str"],
                    r["template_value"],
                    r["ocr_value"],
                    r["color_info"].get("color_class", "?"),
                    r["color_info"].get("mean_hue", 0),
                    r["color_info"].get("mean_saturation", 0),
                )

    # Save full results
    summary_path = Path("output/debug/texas_yellow_test_summary.json")
    with open(summary_path, "w") as f:
        json.dump(results, f, indent=2)
    logger.info("")
    logger.info("Full results saved to: %s", summary_path)
    logger.info("Mismatch images saved to: output/debug/texas_yellow_test/")


if __name__ == "__main__":
    main()