#!/usr/bin/env python3
"""Circle candidate-generation strategies (the pluggable part).

Every finder differs only in *how it proposes candidate circles*; the
fit / refine / score steps are shared. So each proposal method lives behind one
interface, ``BaseDetector.detect(ctx) -> list[Circle]``, and is looked up by
name in ``DETECTOR_REGISTRY``. Add a new method for another project with::

    from ecos_core.circle_finder.detectors import BaseDetector, register_detector

    @register_detector("my_method")
    class MyDetector(BaseDetector):
        DEFAULTS = {...}
        def detect(self, ctx):
            ...
            return [Circle(...)]

...without touching the Model or the other detectors (Open/Closed).

Four strategies ship, covering what NITTO uses:
  - ``hough``   : cv2.HoughCircles on a (down-scaled) channel. Centre-agnostic.
  - ``contour`` : threshold → contours → radius/circularity-gated LSQ fit.
  - ``radial``  : radial gradient edge search from a seed → trimmed LSQ. Precise.
  - ``ransac``  : (optionally multi-channel) Canny edges → band-gated RANSAC.
"""
from __future__ import annotations

import math
from dataclasses import dataclass, field
from abc import ABC, abstractmethod
from typing import Optional

import cv2
import numpy as np

from ecos_core.circle_finder.geometry import (
    Circle, circularity, fit_circle_lsq, fit_circle_trimmed,
    fit_circle_ransac, refine_center_subpixel_radial,
)
from ecos_core.circle_finder import preprocess as pp


# ---------------------------------------------------------------------------
# Shared context passed to every detector
# ---------------------------------------------------------------------------
@dataclass
class DetectContext:
    """What a detector receives. ``gray`` is the working single-channel image
    the Model already prepared (channel select + BC + blur); ``bgr`` is the
    colour image (or None) for methods that need HSV. ``r_min``/``r_max`` is the
    resolved pixel radius band; ``seed`` is an optional coarse circle (centre /
    radius) for methods that refine from a starting guess (``radial``).
    ``params`` carries the method's own tunables, so a custom detector can read
    anything without changing this dataclass.
    """

    gray: np.ndarray
    bgr: Optional[np.ndarray] = None
    r_min: float = 0.0
    r_max: float = float("inf")
    seed: Optional[Circle] = None
    params: dict = field(default_factory=dict)

    @property
    def r_target(self) -> float:
        """Mid-band radius; a reasonable default target when no seed is given."""
        if math.isinf(self.r_max):
            return self.r_min if self.r_min > 0 else 0.0
        return 0.5 * (self.r_min + self.r_max)


# ---------------------------------------------------------------------------
# Base class + registry
# ---------------------------------------------------------------------------
class BaseDetector(ABC):
    """Common interface. Subclasses set ``NAME`` and ``DEFAULTS`` and implement
    ``detect``."""

    NAME: str = "base"
    DEFAULTS: dict = {}

    def __init__(self, **params):
        self.p = {**self.DEFAULTS, **(params or {})}

    @abstractmethod
    def detect(self, ctx: DetectContext) -> list:
        """Return a list of Circle candidates (may be empty)."""
        raise NotImplementedError

    # small shared helper: keep only circles inside the band
    def _in_band(self, r: float, ctx: DetectContext) -> bool:
        return ctx.r_min <= r <= ctx.r_max


DETECTOR_REGISTRY: dict = {}


def register_detector(name: str):
    """Class decorator that registers a detector under ``name``."""
    def _wrap(cls):
        cls.NAME = name
        DETECTOR_REGISTRY[name] = cls
        return cls
    return _wrap


def build_detector(name: str, params: Optional[dict] = None) -> BaseDetector:
    """Instantiate a registered detector by name."""
    if name not in DETECTOR_REGISTRY:
        raise KeyError(
            f"Unknown circle detector '{name}'. Registered: {sorted(DETECTOR_REGISTRY)}"
        )
    return DETECTOR_REGISTRY[name](**(params or {}))


def available_detectors() -> list:
    return sorted(DETECTOR_REGISTRY)


# ---------------------------------------------------------------------------
# Strategy 1: Hough
# ---------------------------------------------------------------------------
@register_detector("hough")
class HoughDetector(BaseDetector):
    """cv2.HoughCircles on the working channel, optionally down-scaled for speed
    and vote robustness. Centre-agnostic (no seed needed), but prone to phantom
    circles on smooth gradients - pair with a strong ``min_score`` gate or the
    Model's edge-support re-scoring.
    """

    DEFAULTS = {
        "downscale": 1,      # >1 = detect on a shrunk image, then rescale back
        "dp": 1.5,
        "min_dist_frac": 1.0,   # minDist as a fraction of r_target (or image/8)
        "param1": 120,
        "param2": 30,
        "blur_ksize": 0,      # extra blur applied to the (down-scaled) channel
    }

    def detect(self, ctx: DetectContext) -> list:
        p = self.p
        ds = max(1, int(p["downscale"]))
        img = ctx.gray
        if ds > 1:
            img = cv2.resize(img, (img.shape[1] // ds, img.shape[0] // ds),
                             interpolation=cv2.INTER_AREA)
        img = pp.blur(img, int(p["blur_ksize"]))

        r_min = 0 if math.isinf(ctx.r_min) else int(ctx.r_min / ds)
        r_max = 0 if math.isinf(ctx.r_max) else int(ctx.r_max / ds)
        r_t = ctx.r_target / ds if ctx.r_target else max(img.shape) / 8.0
        min_dist = max(1.0, float(p["min_dist_frac"]) * (r_t or max(img.shape) / 8.0))

        circles = cv2.HoughCircles(
            img, cv2.HOUGH_GRADIENT, dp=float(p["dp"]), minDist=min_dist,
            param1=float(p["param1"]), param2=float(p["param2"]),
            minRadius=int(r_min), maxRadius=int(r_max),
        )
        out = []
        if circles is None:
            return out
        circles = circles[0]  # (N, 3): x, y, r  - HoughCircles returns strongest first
        n = len(circles)
        for i, (x, y, r) in enumerate(circles):
            cx, cy, rr = float(x) * ds, float(y) * ds, float(r) * ds
            if not self._in_band(rr, ctx):
                continue
            # provisional score by vote rank; the Model re-scores with edge support
            out.append(Circle(cx, cy, rr, score=1.0 - i / max(1, n),
                              method=self.NAME, meta={"rank": i}))
        return out


# ---------------------------------------------------------------------------
# Strategy 2: Contour
# ---------------------------------------------------------------------------
@register_detector("contour")
class ContourDetector(BaseDetector):
    """Threshold → external contours → radius/circularity-gated selection, with
    a robust LSQ fit for the final circle. The workhorse: fast and reliable when
    the target segments cleanly against its background.

    Threshold modes: ``otsu`` (Otsu, optionally inverted), ``fixed`` (a fixed
    ``thresh`` level), ``adaptive`` (adaptive Gaussian), ``hsv`` (an HSV
    in-range colour mask - needs a BGR image).
    """

    DEFAULTS = {
        "thresh_mode": "otsu",     # otsu | fixed | adaptive | hsv
        "thresh": 127,             # for fixed
        "invert": True,            # target darker than background → invert so it is white
        "adaptive_block": 51,
        "adaptive_C": 5,
        "hsv": None,               # {"h":[..],"s":[..],"v":[..],"inv":false,"scale":"cv"}
        "close_ksize": 5,
        "open_ksize": 3,
        "min_area_frac": 0.0005,   # reject specks (fraction of image area)
        "min_circularity": 0.5,    # per-candidate gate before fitting
        "max_candidates": 5,
        "fit": "trimmed",          # trimmed | lsq | none (none = boundingRect circle)
    }

    def _binarize(self, ctx: DetectContext) -> np.ndarray:
        p = self.p
        mode = str(p["thresh_mode"]).lower()
        if mode == "hsv":
            if ctx.bgr is None or not p.get("hsv"):
                # cannot build a colour mask → fall back to Otsu on gray
                mode = "otsu"
            else:
                h = p["hsv"]
                return pp.hsv_mask(ctx.bgr, h.get("h", [0, 179]), h.get("s", [0, 255]),
                                   h.get("v", [0, 255]), inv=h.get("inv", False),
                                   hue_scale=h.get("scale", "cv"))
        gray = ctx.gray
        if mode == "fixed":
            flag = cv2.THRESH_BINARY_INV if p["invert"] else cv2.THRESH_BINARY
            _, b = cv2.threshold(gray, int(p["thresh"]), 255, flag)
            return b
        if mode == "adaptive":
            block = int(p["adaptive_block"])
            if block % 2 == 0:
                block += 1
            flag = cv2.THRESH_BINARY_INV if p["invert"] else cv2.THRESH_BINARY
            return cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                         flag, block, float(p["adaptive_C"]))
        # otsu (default)
        flag = (cv2.THRESH_BINARY_INV if p["invert"] else cv2.THRESH_BINARY) + cv2.THRESH_OTSU
        _, b = cv2.threshold(gray, 0, 255, flag)
        return b

    def detect(self, ctx: DetectContext) -> list:
        p = self.p
        mask = self._binarize(ctx)
        ck, ok = int(p["close_ksize"]), int(p["open_ksize"])
        if ck > 1:
            mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE,
                                    cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (ck, ck)))
        if ok > 1:
            mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN,
                                    cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (ok, ok)))
        cnts, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        area_img = ctx.gray.shape[0] * ctx.gray.shape[1]
        min_area = float(p["min_area_frac"]) * area_img
        min_circ = float(p["min_circularity"])

        scored = []
        for c in cnts:
            area = cv2.contourArea(c)
            if area < min_area:
                continue
            # Radius from boundingRect, NOT cv2.minEnclosingCircle: the latter has
            # a pathological ~26 s case on large jagged contours in OpenCV 4.x,
            # while boundingRect stays sub-ms and is accurate enough for gating.
            bx, by, bw, bh = cv2.boundingRect(c)
            r_est = 0.25 * (bw + bh)
            if not self._in_band(r_est, ctx):
                continue
            circ = circularity(area, cv2.arcLength(c, True))
            if circ < min_circ:
                continue
            scored.append((circ, c, (bx, by, bw, bh)))

        scored.sort(key=lambda t: t[0], reverse=True)
        out = []
        for circ, c, (bx, by, bw, bh) in scored[: int(p["max_candidates"])]:
            cx, cy, r = self._fit(c, bx, by, bw, bh)
            if r is None or not self._in_band(r, ctx):
                continue
            out.append(Circle(cx, cy, r, score=float(min(1.0, circ)),
                              method=self.NAME, meta={"circularity": round(circ, 3)}))
        return out

    def _fit(self, contour, bx, by, bw, bh):
        mode = str(self.p["fit"]).lower()
        pts = contour.reshape(-1, 2).astype(np.float64)
        if mode == "none":
            return bx + bw / 2.0, by + bh / 2.0, 0.25 * (bw + bh)
        if mode == "lsq":
            cx, cy, r = fit_circle_lsq(pts)
            return cx, cy, r
        f = fit_circle_trimmed(pts, rounds=2, k=2.5)
        if f is None:
            return bx + bw / 2.0, by + bh / 2.0, 0.25 * (bw + bh)
        return f[0], f[1], f[2]


# ---------------------------------------------------------------------------
# Strategy 3: Radial ray edge search
# ---------------------------------------------------------------------------
@register_detector("radial")
class RadialRayDetector(BaseDetector):
    """From a seed centre, cast rays and locate the edge crossing near a target
    radius via the |gradient| peak, then fit a trimmed LSQ circle. The most
    precise method for a well-localised boundary; needs a seed (centre + rough
    radius). The Model supplies one automatically by running a cheap ``contour``
    (or ``hough``) pass first when ``seed_method`` is set - a natural
    coarse-to-fine composition.
    """

    DEFAULTS = {
        "num_rays": 72,
        "r_lo_frac": 0.6,     # scan window around the seed radius
        "r_hi_frac": 1.4,
        "grad_frac": 0.4,     # peak must exceed grad_frac * max ray gradient
        "grad_min": 6.0,
        "seed_method": "contour",   # how the Model seeds this (None = image centre)
        "seed_params": None,
    }

    def detect(self, ctx: DetectContext) -> list:
        seed = ctx.seed
        if seed is None:
            # no seed provided: fall back to the image centre + mid-band radius
            h, w = ctx.gray.shape[:2]
            r_t = ctx.r_target or 0.25 * min(h, w)
            seed = Circle(w / 2.0, h / 2.0, r_t)
        p = self.p
        r_lo = max(1.0, seed.r * float(p["r_lo_frac"]))
        r_hi = seed.r * float(p["r_hi_frac"])
        fit = self._edge_fit(ctx.gray, seed.cx, seed.cy, r_lo, r_hi, seed.r,
                             int(p["num_rays"]), float(p["grad_frac"]), float(p["grad_min"]))
        if fit is None:
            return []
        cx, cy, r, inl = fit
        if not self._in_band(r, ctx):
            return []
        return [Circle(cx, cy, r, score=float(inl), method=self.NAME,
                       meta={"inlier_frac": round(inl, 3), "seed": seed.as_int()})]

    @staticmethod
    def _edge_fit(gray, cx, cy, r_lo, r_hi, r_target, num_rays, grad_frac, grad_min):
        rs = np.arange(float(r_lo), float(r_hi), 1.0)
        if len(rs) < 5:
            return None
        h, w = gray.shape[:2]
        edge_pts = []
        for i in range(num_rays):
            a = 2.0 * math.pi * i / num_rays
            ct, st = math.cos(a), math.sin(a)
            xi = (cx + rs * ct).astype(int)
            yi = (cy + rs * st).astype(int)
            ok = (xi >= 0) & (xi < w) & (yi >= 0) & (yi < h)
            if ok.sum() < 5:
                continue
            prof = gray[np.clip(yi, 0, h - 1), np.clip(xi, 0, w - 1)].astype(np.float64)
            rr = rs[ok]
            g = np.abs(np.gradient(prof[ok]))
            thr = max(grad_min, grad_frac * float(g.max()))
            peaks = np.where(g >= thr)[0]
            if len(peaks) == 0:
                continue
            j = peaks[int(np.argmin(np.abs(rr[peaks] - r_target)))]
            edge_pts.append([cx + rr[j] * ct, cy + rr[j] * st])
        if len(edge_pts) < 12:
            return None
        return fit_circle_trimmed(np.array(edge_pts))


# ---------------------------------------------------------------------------
# Strategy 4: Canny edges + RANSAC
# ---------------------------------------------------------------------------
@register_detector("ransac")
class RansacEdgeDetector(BaseDetector):
    """Canny edges (optionally a union of several channels) → band-gated RANSAC
    circle. Robust to heavy edge texture (brushed rings, striations) that pulls
    a contour/least-squares fit off-centre, because RANSAC ignores off-band
    inconsistent edges. Needs a valid radius band to gate hypotheses.
    """

    DEFAULTS = {
        "channels": ["gray"],   # union Canny across these channels
        "canny_lo": 80,
        "canny_hi": 160,
        "blur_ksize": 7,
        "tol": 2.5,             # inlier distance (px)
        "iters": 3000,
        "min_inliers": 40,
        "mask_center": None,    # optional {"cx","cy","r"} to keep only interior edges
    }

    def detect(self, ctx: DetectContext) -> list:
        p = self.p
        if math.isinf(ctx.r_max) or ctx.r_max <= 0:
            # RANSAC needs a bounded band to gate hypotheses
            return []
        edges = self._edges(ctx)
        mc = p.get("mask_center")
        if mc:
            m = np.zeros(edges.shape, np.uint8)
            cv2.circle(m, (int(mc["cx"]), int(mc["cy"])), int(mc["r"]), 255, -1)
            edges = cv2.bitwise_and(edges, m)
        ys, xs = np.nonzero(edges)
        if len(xs) < 10:
            return []
        pool = np.c_[xs, ys].astype(np.float64)
        res = fit_circle_ransac(pool, ctx.r_min, ctx.r_max,
                                tol=float(p["tol"]), iters=int(p["iters"]),
                                rng=np.random.default_rng(0))
        if res is None:
            return []
        cx, cy, r, n_inl = res
        if n_inl < int(p["min_inliers"]) or not self._in_band(r, ctx):
            return []
        # score by inlier density along the circumference (~2*pi*r px expected)
        density = min(1.0, n_inl / max(1.0, 2 * math.pi * r))
        return [Circle(cx, cy, r, score=float(density), method=self.NAME,
                       meta={"inliers": int(n_inl)})]

    def _edges(self, ctx: DetectContext) -> np.ndarray:
        p = self.p
        lo, hi, bk = int(p["canny_lo"]), int(p["canny_hi"]), int(p["blur_ksize"])
        channels = p.get("channels") or ["gray"]
        acc = None
        for ch in channels:
            if ch == "gray" or ctx.bgr is None:
                src = ctx.gray
            else:
                src = pp.select_channel(ctx.bgr, ch)
            e = cv2.Canny(pp.blur(src, bk), lo, hi)
            acc = e if acc is None else cv2.bitwise_or(acc, e)
        return acc
