#!/usr/bin/env python3
"""Pure circle geometry primitives (numpy + OpenCV only, no torch).

These are the building blocks every circle finder re-implemented in its
own way: algebraic / robust circle fitting, sub-pixel centre refinement, a
circularity metric, and a method-agnostic edge-support quality score. Kept
dependency-light so they can be imported standalone (unit tests, notebooks,
other projects) without pulling in the ECOS BaseModel / torch stack.
"""
from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Optional

import cv2
import numpy as np


# ---------------------------------------------------------------------------
# Result type
# ---------------------------------------------------------------------------
@dataclass
class Circle:
    """A detected circle in image-pixel coordinates.

    Attributes:
        cx, cy, r: centre and radius in pixels.
        score: quality in [0, 1]. Its native meaning is method-specific
            (circularity for contour, inlier fraction for RANSAC/radial, vote
            rank for Hough); the Model re-scores every candidate with the
            common ``edge_support`` metric so a union of methods can be ranked
            fairly.
        method: name of the detector that produced it.
        meta: free-form extra info (inlier fraction, raw score, ...).
    """

    cx: float
    cy: float
    r: float
    score: float = 0.0
    method: str = ""
    meta: dict = field(default_factory=dict)

    @property
    def center(self) -> tuple:
        return (self.cx, self.cy)

    def as_int(self) -> tuple:
        return (int(round(self.cx)), int(round(self.cy)), int(round(self.r)))

    def shifted(self, dx: float, dy: float) -> "Circle":
        """Return a copy translated by (dx, dy) - used to lift a crop-local
        detection back into full-image coordinates."""
        return Circle(self.cx + dx, self.cy + dy, self.r, self.score, self.method, dict(self.meta))


# ---------------------------------------------------------------------------
# Basic metrics
# ---------------------------------------------------------------------------
def circularity(area: float, perimeter: float) -> float:
    """4*pi*A / P^2 - 1.0 for a perfect circle, →0 for a ragged shape."""
    if perimeter <= 1e-9:
        return 0.0
    return float(4.0 * math.pi * area / (perimeter * perimeter))


def bilinear_sample(gray: np.ndarray, xs: np.ndarray, ys: np.ndarray) -> np.ndarray:
    """Bilinearly sample ``gray`` at float coords; out-of-range → NaN.

    Accepts scalar or array xs/ys of any matching shape and returns a float
    array of that shape. Used by every radial/edge routine so sub-pixel
    sampling is consistent.
    """
    xs = np.asarray(xs, dtype=np.float64)
    ys = np.asarray(ys, dtype=np.float64)
    H, W = gray.shape[:2]
    x0 = np.floor(xs).astype(int)
    y0 = np.floor(ys).astype(int)
    valid = (x0 >= 0) & (y0 >= 0) & (x0 + 1 < W) & (y0 + 1 < H)
    xc = np.clip(x0, 0, W - 2)
    yc = np.clip(y0, 0, H - 2)
    fx = xs - xc
    fy = ys - yc
    g = gray.astype(np.float32)
    val = (g[yc, xc] * (1 - fx) * (1 - fy) + g[yc, xc + 1] * fx * (1 - fy) +
           g[yc + 1, xc] * (1 - fx) * fy + g[yc + 1, xc + 1] * fx * fy)
    val = np.asarray(val, dtype=np.float64)
    val[~valid] = np.nan
    return val


# ---------------------------------------------------------------------------
# Circle fitting
# ---------------------------------------------------------------------------
def fit_circle_lsq(points: np.ndarray) -> tuple:
    """Algebraic least-squares circle fit. Returns (cx, cy, r)."""
    pts = np.asarray(points, dtype=np.float64).reshape(-1, 2)
    x, y = pts[:, 0], pts[:, 1]
    A = np.column_stack([2 * x, 2 * y, np.ones(len(x))])
    b = x ** 2 + y ** 2
    result, _, _, _ = np.linalg.lstsq(A, b, rcond=None)
    cx, cy = result[0], result[1]
    r = float(np.sqrt(max(result[2] + cx ** 2 + cy ** 2, 0.0)))
    return float(cx), float(cy), r


def fit_circle_trimmed(points: np.ndarray, rounds: int = 2, k: float = 2.5) -> Optional[tuple]:
    """LSQ circle fit with iterative outlier trimming.

    Rejects points more than ``k`` standard deviations from the current circle
    each round - so a local edge defect (a notch/chip on the rim, a stray edge
    pixel) cannot drag the centre. Returns (cx, cy, r, inlier_frac) or None if
    there are too few points.
    """
    pts = np.asarray(points, dtype=np.float64).reshape(-1, 2)
    if len(pts) < 6:
        return None
    keep = np.ones(len(pts), dtype=bool)
    cx, cy, r = fit_circle_lsq(pts)
    for _ in range(rounds):
        d = np.abs(np.hypot(pts[:, 0] - cx, pts[:, 1] - cy) - r)
        s = d[keep].std() if keep.sum() > 5 else d.std()
        if s < 1e-6:
            break
        keep = d <= k * s
        if keep.sum() < 6:
            break
        cx, cy, r = fit_circle_lsq(pts[keep])
    return float(cx), float(cy), float(r), float(keep.mean())


def fit_circle_ransac(points: np.ndarray, r_min: float, r_max: float,
                      tol: float = 2.5, iters: int = 3000,
                      rng: Optional[np.random.Generator] = None) -> Optional[tuple]:
    """Vectorised RANSAC circle fit on edge points, gated to [r_min, r_max].

    Samples all 3-point hypotheses at once, keeps those whose circumscribed
    circle lands in the radius band, then counts inliers in memory-bounded
    chunks and refines the best by LSQ. Robust to heavy texture noise (brushed
    rings, striations) that defeats a plain contour fit. Returns
    (cx, cy, r, n_inliers) or None.
    """
    P = np.asarray(points, dtype=np.float64).reshape(-1, 2)
    n = len(P)
    if n < 10:
        return None
    if rng is None:
        rng = np.random.default_rng(0)
    iters = int(iters)
    idx = rng.integers(0, n, size=(iters, 3))
    ax, ay = P[idx[:, 0], 0], P[idx[:, 0], 1]
    bx, by = P[idx[:, 1], 0], P[idx[:, 1], 1]
    cx, cy = P[idx[:, 2], 0], P[idx[:, 2], 1]
    d = 2 * (ax * (by - cy) + bx * (cy - ay) + cx * (ay - by))
    good = np.abs(d) > 1e-6
    d = np.where(good, d, 1.0)
    a2, b2, c2 = ax * ax + ay * ay, bx * bx + by * by, cx * cx + cy * cy
    ux = (a2 * (by - cy) + b2 * (cy - ay) + c2 * (ay - by)) / d
    uy = (a2 * (cx - bx) + b2 * (ax - cx) + c2 * (bx - ax)) / d
    r = np.hypot(ax - ux, ay - uy)
    keep = good & (r >= r_min) & (r <= r_max)
    if not keep.any():
        return None
    ux, uy, r = ux[keep], uy[keep], r[keep]
    Px, Py = P[:, 0][None, :], P[:, 1][None, :]
    best_n, best_j = 0, -1
    for s in range(0, len(ux), 512):
        cxs, cys, rs = ux[s:s + 512, None], uy[s:s + 512, None], r[s:s + 512, None]
        cnt = (np.abs(np.hypot(Px - cxs, Py - cys) - rs) < tol).sum(1)
        j = int(cnt.argmax())
        if cnt[j] > best_n:
            best_n, best_j = int(cnt[j]), s + j
    if best_j < 0:
        return None
    bx0, by0, br0 = ux[best_j], uy[best_j], r[best_j]
    inl = np.abs(np.hypot(P[:, 0] - bx0, P[:, 1] - by0) - br0) < tol
    if inl.sum() < 6:
        return None
    cx0, cy0, r0 = fit_circle_lsq(P[inl])
    return cx0, cy0, r0, int(inl.sum())


# ---------------------------------------------------------------------------
# Sub-pixel refinement
# ---------------------------------------------------------------------------
def refine_center_subpixel_radial(gray: np.ndarray, cx: float, cy: float, radius: float,
                                  num_rays: int = 36, search_range: float = 6.0,
                                  trim: bool = True) -> tuple:
    """Refine a circle centre/radius via radial gradient sampling.

    Casts ``num_rays`` rays from (cx, cy); on each, samples intensity around
    ``radius`` ± ``search_range``, locates the |gradient| peak to sub-pixel via
    3-point parabolic interpolation, then fits the collected edge points. With
    ``trim=True`` the final fit rejects rays that landed on a local defect.
    Returns (cx_sub, cy_sub, r_sub); falls back to the input on too few edges.
    """
    h, w = gray.shape[:2]
    edge_points = []
    drs = np.arange(-search_range, search_range + 0.5, 0.5)
    for i in range(num_rays):
        angle = 2.0 * math.pi * i / num_rays
        ct, st = math.cos(angle), math.sin(angle)
        r = radius + drs
        px = cx + r * ct
        py = cy + r * st
        vals = bilinear_sample(gray, px, py)
        ok = ~np.isnan(vals)
        if ok.sum() < 5:
            continue
        v = vals[ok]
        d = drs[ok]
        grads = np.abs(np.gradient(v, d))
        peak_idx = int(np.argmax(grads))
        if 1 <= peak_idx < len(grads) - 1:
            y1, y2, y3 = grads[peak_idx - 1], grads[peak_idx], grads[peak_idx + 1]
            denom = 2 * y2 - y1 - y3
            delta = 0.5 * (y1 - y3) / denom if abs(denom) > 1e-9 else 0.0
            step = d[1] - d[0] if len(d) > 1 else 0.5
            dr_sub = d[peak_idx] + delta * step
        else:
            dr_sub = d[peak_idx]
        r_edge = radius + dr_sub
        edge_points.append([cx + r_edge * ct, cy + r_edge * st])

    if len(edge_points) < 6:
        return float(cx), float(cy), float(radius)
    pts = np.array(edge_points)
    if trim:
        fit = fit_circle_trimmed(pts, rounds=2, k=2.5)
        if fit is not None:
            return float(fit[0]), float(fit[1]), float(fit[2])
    cx0, cy0, r0 = fit_circle_lsq(pts)
    return cx0, cy0, r0


# ---------------------------------------------------------------------------
# Method-agnostic quality score
# ---------------------------------------------------------------------------
def edge_support_score(gray: np.ndarray, cx: float, cy: float, r: float,
                       num_rays: int = 72, band: float = 0.15,
                       min_grad: float = 8.0, tol_frac: float = 0.06) -> float:
    """Fraction of radial rays whose strongest intensity step sits *at* radius r.

    A common, detector-independent measure of "how much real edge sits on this
    circle, concentric with (cx, cy)". Because Hough votes, contour circularity
    and RANSAC inliers are not comparable across methods, the Model re-scores
    every candidate with this so a union of methods can be ranked on one scale.

    Each ray is scanned over ``[r(1-band), r(1+band)]``; it counts as support
    only when its gradient peak both exceeds ``min_grad`` *and* lands within
    ``tol_frac * r`` of ``r``. Requiring the peak *near r* (not merely somewhere
    in the band) is what makes the score reward radius/centre accuracy: an
    off-centre circle's edge crosses each ray at a different radius, so fewer
    peaks fall near the assumed r and its score drops below the concentric fit's.
    Returns a value in [0, 1].
    """
    if r <= 2:
        return 0.0
    lo = r * (1.0 - band)
    hi = r * (1.0 + band)
    rs = np.arange(max(1.0, lo), hi + 0.5, 0.5)
    if len(rs) < 3:
        return 0.0
    tol = max(1.0, tol_frac * r)
    supported = 0
    used = 0
    for i in range(num_rays):
        angle = 2.0 * math.pi * i / num_rays
        px = cx + rs * math.cos(angle)
        py = cy + rs * math.sin(angle)
        vals = bilinear_sample(gray, px, py)
        ok = ~np.isnan(vals)
        if ok.sum() < 3:
            continue
        used += 1
        v = vals[ok]
        rr = rs[ok]
        g = np.abs(np.gradient(v))
        peak = int(np.argmax(g))
        if g[peak] >= min_grad and abs(rr[peak] - r) <= tol:
            supported += 1
    if used == 0:
        return 0.0
    return float(supported) / float(used)
