"""Scan a frame for bright saturated green clusters at several HSV tightnesses.

Reports clusters as (x, y, w, h, pixel_count) and saves a mask per tightness.

"""

from __future__ import annotations

import sys

from pathlib import Path

import cv2

import numpy as np

from rich.console import Console

console = Console()

def report_mask(frame: np.ndarray, hsv: np.ndarray, lo, hi, label: str, out_dir: Path, stem: str) -> None:

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

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

    mask = cv2.inRange(hsv, lo_a, hi_a)

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

    if total == 0:

        console.print(f"[dim]{label}: 0 pixels match {list(lo)}..{list(hi)}[/dim]")

        return

    mask_c = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8), iterations=1)

    num, labels_img, stats, _ = cv2.connectedComponentsWithStats(mask_c, connectivity=8)

    console.print(f"[bold]{label}: total={total}, clusters={num-1} at range {list(lo)}..{list(hi)}[/bold]")

    clusters = []

    for i in range(1, num):

        x, y, w, h, area = stats[i]

        if area < 8:

            continue

        clusters.append((area, x, y, w, h))

    clusters.sort(reverse=True)

    for area, x, y, w, h in clusters[:8]:

        console.print(f"  cluster: x={x}..{x+w} y={y}..{y+h} area={area}")

    out = out_dir / f"green_{label}_{stem}.png"

    cv2.imwrite(str(out), mask)

    console.print(f"  mask -> {out}")

    ann = frame.copy()

    for area, x, y, w, h in clusters[:8]:

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

    ann_out = out_dir / f"green_{label}_{stem}_ann.png"

    cv2.imwrite(str(ann_out), ann)

def main() -> int:

    if len(sys.argv) < 2:

        console.print("usage: find_green.py <frame.png>")

        return 2

    fp = Path(sys.argv[1])

    frame = cv2.imread(str(fp), cv2.IMREAD_COLOR)

    if frame is None:

        console.print(f"[red]can't read {fp}[/red]")

        return 2

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

    out_dir = Path(__file__).resolve().parents[1] / "logs"

    stem = fp.stem

    ranges = [

        ("strict",   [50, 240, 240], [70, 255, 255]),

        ("tight",    [48, 200, 200], [72, 255, 255]),

        ("medium",   [45, 150, 150], [75, 255, 255]),

        ("loose",    [40, 100, 100], [80, 255, 255]),

    ]

    for label, lo, hi in ranges:

        report_mask(frame, hsv, lo, hi, label, out_dir, stem)

        console.print("")

    return 0

if __name__ == "__main__":

    sys.exit(main())