#!/usr/bin/env python3
"""moon_detect.py — find the moon disk in a sky-cam frame.

Returns centroid, radius, and a quality score derived from:
  - roundness  (area / (pi * r^2))
  - isolation  (no comparable bright blob within IGNORE_RADIUS_PX)
  - sky        (low halo extent → clearer sky around the moon)

Used as a library:

  from moon_detect import detect_moon
  result = detect_moon('/path/to/frame.jpg')
  if result is not None:
      cx, cy = result['centroid']
      r      = result['radius']
      score  = result['quality']

CLI (handy for tuning):

  python3 moon_detect.py /path/to/frame.jpg
  python3 moon_detect.py /path/to/frame.jpg --debug debug.png

The detector ignores the bottom OVERLAY_PX rows because capture frames carry a
burnt-in timestamp that contains saturated pixels.
"""
from __future__ import annotations

import argparse
import json
import math
import sys
from dataclasses import dataclass

import numpy as np
from PIL import Image, ImageDraw
from scipy import ndimage

# --- Tuning constants --------------------------------------------------------
SATURATED_THRESHOLD = 240   # 0..255 — pixels at or above this count as "moon disk"
HALO_THRESHOLD = 100        # 0..255 — pixels above this count toward halo extent
OVERLAY_PX = 200            # bottom rows to ignore (timestamp overlay)
MIN_DIAMETER_PX = 10        # blob smaller than this is noise
MAX_DIAMETER_PX = 120       # blob bigger than this is probably the sun, not the moon
MIN_ROUNDNESS = 0.55        # area / (pi * r^2) — perfect circle = 1
ISOLATION_PX = 200          # other comparable blob this close → reject
# Halo extent — typical moon halo is ~3.5x disk radius; >5x means thick clouds
MAX_HALO_RATIO = 6.0


@dataclass
class MoonDetection:
    centroid_xy: tuple[float, float]
    radius_px: float
    diameter_px: float
    blob_pixels: int
    roundness: float
    halo_radius_px: float
    halo_ratio: float
    isolation_px: float
    quality: float

    def asdict(self) -> dict:
        return {
            'centroid': list(self.centroid_xy),
            'radius': self.radius_px,
            'diameter': self.diameter_px,
            'blob_pixels': self.blob_pixels,
            'roundness': self.roundness,
            'halo_radius': self.halo_radius_px,
            'halo_ratio': self.halo_ratio,
            'isolation': self.isolation_px,
            'quality': self.quality,
        }


def _grayscale_array(image_path: str) -> np.ndarray:
    im = Image.open(image_path).convert('L')
    return np.asarray(im)


def _largest_round_blob(
    mask: np.ndarray,
    min_roundness: float = MIN_ROUNDNESS,
) -> tuple[int, np.ndarray, np.ndarray] | None:
    labels, n = ndimage.label(mask)
    if n == 0:
        return None
    sizes = ndimage.sum(mask, labels, range(1, n + 1))
    order = np.argsort(sizes)[::-1]
    for idx in order[:8]:
        label_id = idx + 1
        ys, xs = np.where(labels == label_id)
        if len(xs) < 4:
            continue
        bbw = xs.max() - xs.min() + 1
        bbh = ys.max() - ys.min() + 1
        diam = max(bbw, bbh)
        if diam < MIN_DIAMETER_PX or diam > MAX_DIAMETER_PX:
            continue
        radius = diam / 2.0
        roundness = len(xs) / (math.pi * radius * radius)
        if roundness < min_roundness:
            continue
        return label_id, ys, xs
    return None


def detect_moon(
    image_path: str,
    saturated_threshold: int = SATURATED_THRESHOLD,
    min_roundness: float = MIN_ROUNDNESS,
    max_halo_ratio: float = MAX_HALO_RATIO,
) -> MoonDetection | None:
    """Return MoonDetection or None if no acceptable moon is found.

    saturated_threshold / min_roundness / max_halo_ratio can be relaxed for
    crescent phases, which produce a dim, non-circular arc rather than a bright
    near-perfect disk.
    """
    a = _grayscale_array(image_path)
    if OVERLAY_PX > 0:
        a = a.copy()
        a[-OVERLAY_PX:, :] = 0
    mask = a >= saturated_threshold
    if not mask.any():
        return None

    found = _largest_round_blob(mask, min_roundness)
    if found is None:
        return None
    _, ys, xs = found

    cx = float(xs.mean())
    cy = float(ys.mean())
    blob_pixels = int(len(xs))
    bbw = xs.max() - xs.min() + 1
    bbh = ys.max() - ys.min() + 1
    diameter = float(max(bbw, bbh))
    radius = diameter / 2.0
    roundness = blob_pixels / (math.pi * radius * radius)

    # Isolation: any other saturated blob nearby of comparable size?
    labels_full, n_full = ndimage.label(mask)
    own_label = labels_full[int(round(cy)), int(round(cx))]
    isolation = float('inf')
    for label_id in range(1, n_full + 1):
        if label_id == own_label:
            continue
        ys2, xs2 = np.where(labels_full == label_id)
        if len(xs2) < blob_pixels * 0.3:
            continue
        d = math.hypot(xs2.mean() - cx, ys2.mean() - cy)
        if d < isolation:
            isolation = d
    if isolation < ISOLATION_PX:
        return None

    # Halo: connected component above HALO_THRESHOLD that contains the centroid
    halo_mask = a >= HALO_THRESHOLD
    halo_labels, _ = ndimage.label(halo_mask)
    halo_id = halo_labels[int(round(cy)), int(round(cx))]
    if halo_id == 0:
        halo_radius = radius
    else:
        yh, xh = np.where(halo_labels == halo_id)
        halo_radius = max(xh.max() - xh.min(), yh.max() - yh.min()) / 2.0
    halo_ratio = halo_radius / radius if radius > 0 else 1.0
    if halo_ratio > max_halo_ratio:
        return None  # too much glow → probably thick cloud cover

    # Quality score: roundness (0..1), low halo (1 = clear, 0 = thick cloud),
    # isolation factor (1 if very isolated, less if close to other lights).
    halo_clean = max(0.0, min(1.0, (max_halo_ratio - halo_ratio) / (max_halo_ratio - 1.5)))
    iso_factor = 1.0 if isolation == float('inf') else min(1.0, isolation / 600.0)
    quality = 0.5 * roundness + 0.35 * halo_clean + 0.15 * iso_factor

    return MoonDetection(
        centroid_xy=(cx, cy),
        radius_px=radius,
        diameter_px=diameter,
        blob_pixels=blob_pixels,
        roundness=roundness,
        halo_radius_px=halo_radius,
        halo_ratio=halo_ratio,
        isolation_px=isolation if isolation != float('inf') else -1.0,
        quality=quality,
    )


def _draw_debug(image_path: str, detection: MoonDetection | None, out_path: str):
    im = Image.open(image_path).convert('RGB')
    draw = ImageDraw.Draw(im)
    if detection is None:
        draw.text((20, 20), 'NO MOON DETECTED', fill='red')
    else:
        cx, cy = detection.centroid_xy
        r = detection.radius_px
        hr = detection.halo_radius_px
        draw.ellipse((cx - r, cy - r, cx + r, cy + r), outline='yellow', width=4)
        draw.ellipse((cx - hr, cy - hr, cx + hr, cy + hr), outline='cyan', width=2)
        draw.text((20, 20), f'q={detection.quality:.2f} r={r:.1f} halo={hr:.1f}',
                  fill='yellow')
    im.save(out_path)


def _cli():
    p = argparse.ArgumentParser()
    p.add_argument('image')
    p.add_argument('--debug', help='write annotated PNG to this path')
    p.add_argument('--json', action='store_true', help='emit JSON for scripting')
    args = p.parse_args()

    det = detect_moon(args.image)
    if args.json:
        print(json.dumps(det.asdict() if det else None))
    elif det is None:
        print('no moon detected')
    else:
        d = det.asdict()
        for k, v in d.items():
            print(f'{k}={v}')
    if args.debug:
        _draw_debug(args.image, det, args.debug)
        print(f'debug image written: {args.debug}', file=sys.stderr)
    return 0 if det else 2


if __name__ == '__main__':
    sys.exit(_cli())