"""HSV color probe.

Given a frame PNG and a bounding box (or auto-detect the brightest green cluster),

print the HSV percentile stats so we can pick a tight HSV range that catches the

PCI overlay but not the grass field.

Usage:

  hsv_probe.py <frame.png>

  hsv_probe.py <frame.png> --box X Y W H

  hsv_probe.py <frame.png> --test-range "50,200,200-70,255,255"

"""

from __future__ import annotations

import argparse

import sys

from pathlib import Path

import cv2

import numpy as np

from rich.console import Console

console = Console()

ROOT = Path(__file__).resolve().parents[1]

LOG_DIR = ROOT / "logs"

def auto_pci_box(hsv: np.ndarray) -> tuple[int, int, int, int]:

    """Find the brightest, most-saturated green cluster in the central half."""

    h, w = hsv.shape[:2]

    mask = cv2.inRange(hsv, np.array([40, 200, 200], np.uint8), np.array([80, 255, 255], np.uint8))

    cx0, cy0 = w // 4, h // 4

    cx1, cy1 = 3 * w // 4, 3 * h // 4

    center = np.zeros_like(mask)

    center[cy0:cy1, cx0:cx1] = 255

    mask = cv2.bitwise_and(mask, center)

    ys, xs = np.where(mask > 0)

    if len(xs) == 0:

        return w // 3, h // 3, w // 3, h // 3

    return int(xs.min()), int(ys.min()), int(xs.max() - xs.min()), int(ys.max() - ys.min())

def main() -> int:

    ap = argparse.ArgumentParser(description="HSV probe.")

    ap.add_argument("frame_path")

    ap.add_argument("--box", nargs=4, type=int, metavar=("X", "Y", "W", "H"))

    ap.add_argument("--test-range", type=str, default=None,

                    help="'Hlo,Slo,Vlo-Hhi,Shi,Vhi' - save a mask with this range.")

    args = ap.parse_args()

    frame = cv2.imread(args.frame_path, cv2.IMREAD_COLOR)

    if frame is None:

        console.print(f"[red]Cannot read {args.frame_path}[/red]")

        return 2

    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

    if args.box is None:

        x, y, w, h = auto_pci_box(hsv)

        console.print(f"[cyan]auto-box: x={x} y={y} w={w} h={h}[/cyan]")

    else:

        x, y, w, h = args.box

    sub = hsv[y:y + h, x:x + w]

    if sub.size == 0:

        console.print("[red]Empty box.[/red]")

        return 2

    greens = cv2.inRange(sub, np.array([30, 80, 80], np.uint8), np.array([95, 255, 255], np.uint8))

    pci_pixels = sub[greens > 0]

    if len(pci_pixels) == 0:

        console.print("[yellow]No green-ish pixels in the box.[/yellow]")

        return 2

    hs = pci_pixels[:, 0]

    ss = pci_pixels[:, 1]

    vs = pci_pixels[:, 2]

    def stats(name, arr):

        console.print(

            f"  {name}: min={arr.min():3d}  p5={np.percentile(arr,5):5.1f}  "

            f"p50={np.percentile(arr,50):5.1f}  p95={np.percentile(arr,95):5.1f}  max={arr.max():3d}"

        )

    console.print(f"[bold]HSV stats of green-ish pixels in box (n={len(pci_pixels)}):[/bold]")

    stats("H", hs); stats("S", ss); stats("V", vs)

    suggested_low = (int(np.percentile(hs, 5)) - 2,

                     max(40, int(np.percentile(ss, 5)) - 20),

                     max(40, int(np.percentile(vs, 5)) - 20))

    suggested_high = (int(np.percentile(hs, 95)) + 2,

                      min(255, int(np.percentile(ss, 95)) + 20),

                      min(255, int(np.percentile(vs, 95)) + 20))

    console.print(

        f"[bold green]Suggested tight HSV:[/bold green] low={list(suggested_low)}  high={list(suggested_high)}"

    )

    if args.test_range:

        lo_str, hi_str = args.test_range.split("-")

        lo = np.array([int(x) for x in lo_str.split(",")], np.uint8)

        hi = np.array([int(x) for x in hi_str.split(",")], np.uint8)

    else:

        lo = np.array(suggested_low, np.uint8)

        hi = np.array(suggested_high, np.uint8)

    full_mask = cv2.inRange(hsv, lo, hi)

    LOG_DIR.mkdir(exist_ok=True)

    mask_out = LOG_DIR / f"hsv_probe_mask_{Path(args.frame_path).stem}.png"

    cv2.imwrite(str(mask_out), full_mask)

    total = int((full_mask > 0).sum())

    console.print(f"Mask with range {list(lo)}..{list(hi)}: {total} pixels  -> {mask_out}")

    ann = frame.copy()

    cv2.rectangle(ann, (x, y), (x + w, y + h), (255, 0, 255), 2)

    ann_out = LOG_DIR / f"hsv_probe_box_{Path(args.frame_path).stem}.png"

    cv2.imwrite(str(ann_out), ann)

    console.print(f"Annotated box -> {ann_out}")

    return 0

if __name__ == "__main__":

    sys.exit(main())