"""PCI tracker.

Finds the Plate Coverage Indicator (the green circle in MLB The Show Zone+PCI

batting). Uses template matching if configs/templates/pci_circle.png exists;

otherwise falls back to HSV segmentation using the green range in runtime.yaml.

Emits pci_track events each frame (x, y, r, score, method).

"""

from __future__ import annotations

import argparse

import sys

import time

from pathlib import Path

import cv2

import numpy as np

from rich.console import Console

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from cv._common import (

    iter_latest_frames,

    load_config,

    make_frame_subscriber,

    make_pub,

    send_event,

)

console = Console()

def load_template(path: Path) -> np.ndarray | None:

    if not path.exists():

        return None

    tpl = cv2.imread(str(path), cv2.IMREAD_COLOR)

    if tpl is None:

        console.print(f"[yellow]PCI template at {path} could not be read.[/yellow]")

        return None

    return tpl

def detect_by_template(frame: np.ndarray, tpl: np.ndarray, threshold: float):

    res = cv2.matchTemplate(frame, tpl, cv2.TM_CCOEFF_NORMED)

    _, max_val, _, max_loc = cv2.minMaxLoc(res)

    if max_val < threshold:

        return None

    th, tw = tpl.shape[:2]

    x = max_loc[0] + tw / 2.0

    y = max_loc[1] + th / 2.0

    r = 0.5 * min(tw, th)

    return {"x": float(x), "y": float(y), "r": float(r), "score": float(max_val), "method": "template"}

def detect_by_hsv(frame: np.ndarray, cfg_pci: dict):

    """Centroid-based PCI detection.

    The PCI in MLB 26 renders as a multi-part shape (brackets + inner + center),

    so fitting a single contour circle doesn't work reliably. Instead we mask

    the configured green range, restrict to the central strike-zone region,

    and return the centroid of all matching pixels as the PCI position.

    Radius is estimated from the spread of matching pixels.

    """

    hsv_low = np.array(cfg_pci["hsv_low"], dtype=np.uint8)

    hsv_high = np.array(cfg_pci["hsv_high"], dtype=np.uint8)

    min_px = int(cfg_pci.get("min_green_pixels", 150))

    x_lo_f = float(cfg_pci.get("search_x_min_frac", 0.3))

    x_hi_f = float(cfg_pci.get("search_x_max_frac", 0.7))

    y_lo_f = float(cfg_pci.get("search_y_min_frac", 0.3))

    y_hi_f = float(cfg_pci.get("search_y_max_frac", 0.8))

    h, w = frame.shape[:2]

    x0, x1 = int(w * x_lo_f), int(w * x_hi_f)

    y0, y1 = int(h * y_lo_f), int(h * y_hi_f)

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

    mask = cv2.inRange(hsv, hsv_low, hsv_high)

    window = np.zeros_like(mask)

    window[y0:y1, x0:x1] = 255

    mask = cv2.bitwise_and(mask, window)

    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((3, 3), np.uint8), iterations=1)

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

    n = len(xs)

    if n < min_px:

        return None

    cx = float(xs.mean())

    cy = float(ys.mean())

    dx = xs - cx

    dy = ys - cy

    dists = np.sqrt(dx * dx + dy * dy)

    r = float(np.percentile(dists, 80))

    x_bb_min, x_bb_max = int(xs.min()), int(xs.max())

    y_bb_min, y_bb_max = int(ys.min()), int(ys.max())

    bbox_area = max(1, (x_bb_max - x_bb_min + 1) * (y_bb_max - y_bb_min + 1))

    score = float(n / bbox_area)

    return {"x": cx, "y": cy, "r": r, "score": score, "method": "hsv_centroid", "n_px": n}

def main() -> int:

    ap = argparse.ArgumentParser(description="PCI tracker.")

    ap.add_argument("--duration", type=float, default=0.0, help="Stop after N seconds (0 = run forever).")

    ap.add_argument("--quiet", action="store_true")

    args = ap.parse_args()

    cfg = load_config()

    capture_ep = cfg["capture"]["publish_endpoint"]

    pci_ep = cfg["cv"]["pci_events_endpoint"]

    pci_cfg = cfg["cv"]["pci"]

    tpl_path = Path(__file__).resolve().parents[1] / pci_cfg["template"]

    tpl = load_template(tpl_path)

    tpl_thresh = float(pci_cfg.get("conf_threshold", 0.7))

    sub = make_frame_subscriber(capture_ep)

    pub = make_pub(pci_ep)

    method_label = "template" if tpl is not None else "hsv"

    console.print(

        f"[green]pci_tracker[/green] sub={capture_ep} pub={pci_ep}  method={method_label}"

        + ("" if tpl is not None else f"  (no template at {tpl_path.name}; using HSV fallback)")

    )

    t_end = time.perf_counter() + args.duration if args.duration > 0 else None

    frames = 0

    hits = 0

    t_report = time.perf_counter()

    try:

        for meta, frame in iter_latest_frames(sub, timeout_ms=3000):

            if t_end is not None and time.perf_counter() >= t_end:

                break

            frames += 1

            det = None

            if tpl is not None:

                det = detect_by_template(frame, tpl, tpl_thresh)

            if det is None:

                det = detect_by_hsv(frame, pci_cfg)

            if det is not None:

                hits += 1

                event = {

                    "type": "pci_track",

                    "seq": int(meta["seq"]),

                    "ts_ns": int(meta["ts_ns"]),

                    **det,

                }

                send_event(pub, event)

            else:

                send_event(pub, {

                    "type": "pci_miss",

                    "seq": int(meta["seq"]),

                    "ts_ns": int(meta["ts_ns"]),

                })

            now = time.perf_counter()

            if not args.quiet and now - t_report >= 5.0:

                console.print(f"[dim]  pci:  {frames} frames, {hits} hits ({method_label})[/dim]")

                t_report = now

    except TimeoutError as e:

        console.print(f"[red]pci_tracker: {e}. Is capture/ingest.py running?[/red]")

        return 2

    except KeyboardInterrupt:

        console.print("[yellow]pci_tracker interrupted.[/yellow]")

    finally:

        sub.close(0)

        pub.close(0)

    console.print(f"[bold]pci_tracker summary:[/bold] frames={frames} hits={hits} method={method_label}")

    return 0

if __name__ == "__main__":

    sys.exit(main())