# ============================
# fitting.py
# ============================

"""
====================================

* **Filename**:          fitting.py
* **Author**:            Frank Myhre
* **Description**:       Geometric fitters (circle, ellipse, limacon) and a unified general_fit interface.

====================================

**Notes**
* Fitters assume a fixed center (xs, ys) and estimate only shape parameters.
* `general_fit` standardizes outputs for downstream use and tests:
  returns "shape", "params", "center", "success", "cost", and optional "desired".
"""

from __future__ import annotations

import numpy as np

TAU = 2.0 * np.pi

__all__ = [
    "fit_circle",
    "fit_ellipse",
    "fit_limacon",
    "general_fit",
]


def _validate_points(points):
    """
    Validate input points array.

    **Args**:
    * points (array-like): expected shape (N, 2).

    **Returns**:
    * pts (np.ndarray): float array of shape (N, 2).

    """
    pts = np.asarray(points, dtype=float)
    if pts.ndim != 2 or pts.shape[1] != 2:
        raise ValueError(f"`points` must be (N,2), got shape {pts.shape}")
    if pts.shape[0] < 3:
        raise ValueError("Need at least 3 points.")
    return pts


def fit_circle(points, xs, ys):
    """
    Fit a circle radius with fixed center.

    **Args**:
    * points (array-like): (N, 2) sample points.
    * xs (float): fixed x-center.
    * ys (float): fixed y-center.

    **Returns**:
    * r_est (float): estimated radius.

    """
    pts = _validate_points(points)
    xs = float(xs)
    ys = float(ys)
    dx = pts[:, 0] - xs
    dy = pts[:, 1] - ys
    rr = np.sqrt(dx * dx + dy * dy)
    return float(np.mean(rr))


def fit_ellipse(points, xs, ys):
    """
    Fit axis-aligned ellipse diameters (2a, 2b) with fixed center.

    Points assumed generated by:
      x = xs + a*cos(theta) + noise
      y = ys + b*sin(theta) + noise

    **Args**:
    * points (array-like): (N, 2) sample points.
    * xs (float): fixed x-center.
    * ys (float): fixed y-center.

    **Returns**:
    * d2a_est (float): estimated diameter along x (2a).
    * d2b_est (float): estimated diameter along y (2b).

    """
    pts = _validate_points(points)
    xs = float(xs)
    ys = float(ys)
    dx = pts[:, 0] - xs
    dy = pts[:, 1] - ys

    # For uniform theta: Var(a cos θ) = a^2/2  => a ≈ sqrt(2)*std(dx)
    a_est = np.sqrt(2.0) * float(np.std(dx))
    b_est = np.sqrt(2.0) * float(np.std(dy))

    return float(2.0 * a_est), float(2.0 * b_est)


def fit_limacon(points, xs, ys):
    """
    Fit limacon parameters (c, L2, phi) with fixed center.

    Model:
      r(θ) = c * (1 + L2 * cos(θ - phi))

    Linearized regression:
      r = A + B*cosθ + D*sinθ
    where:
      A = c
      B = c*L2*cos(phi)
      D = c*L2*sin(phi)

    **Args**:
    * points (array-like): (N, 2) sample points.
    * xs (float): fixed x-center.
    * ys (float): fixed y-center.

    **Returns**:
    * c_est (float): estimated c.
    * L2_est (float): estimated L2.
    * phi_est (float): estimated phi in radians.

    """
    pts = _validate_points(points)
    xs = float(xs)
    ys = float(ys)
    dx = pts[:, 0] - xs
    dy = pts[:, 1] - ys

    th = np.arctan2(dy, dx) % TAU
    r = np.sqrt(dx * dx + dy * dy)

    M = np.column_stack([np.ones_like(th), np.cos(th), np.sin(th)])
    coef, *_ = np.linalg.lstsq(M, r, rcond=None)
    A, B, D = coef

    c_est = float(A)
    amp = float(np.hypot(B, D))
    if not np.isfinite(c_est) or c_est <= 0:
        c_est = float(np.median(r))
    L2_est = float(amp / c_est) if c_est != 0 else 0.0
    phi_est = float(np.arctan2(D, B))

    return c_est, L2_est, phi_est


def _circle_cost(points, xs, ys, r):
    """
    Mean squared radial residual for circle.

    **Args**:
    * points (np.ndarray): (N,2) points.
    * xs (float): x-center.
    * ys (float): y-center.
    * r (float): radius.

    **Returns**:
    * cost (float): mean squared error.

    """
    dx = points[:, 0] - xs
    dy = points[:, 1] - ys
    rr = np.sqrt(dx * dx + dy * dy)
    return float(np.mean((rr - r) ** 2))


def _ellipse_cost(points, xs, ys, d2a, d2b):
    """
    Mean squared implicit residual for axis-aligned ellipse.

    **Args**:
    * points (np.ndarray): (N,2) points.
    * xs (float): x-center.
    * ys (float): y-center.
    * d2a (float): x-diameter.
    * d2b (float): y-diameter.

    **Returns**:
    * cost (float): mean squared error.

    """
    a = 0.5 * d2a
    b = 0.5 * d2b
    dx = points[:, 0] - xs
    dy = points[:, 1] - ys
    val = (dx / a) ** 2 + (dy / b) ** 2
    return float(np.mean((val - 1.0) ** 2))


def _limacon_cost(points, xs, ys, c, L2, phi):
    """
    Mean squared radial residual for limacon.

    **Args**:
    * points (np.ndarray): (N,2) points.
    * xs (float): x-center.
    * ys (float): y-center.
    * c (float): c parameter.
    * L2 (float): L2 parameter.
    * phi (float): phi phase.

    **Returns**:
    * cost (float): mean squared error.

    """
    dx = points[:, 0] - xs
    dy = points[:, 1] - ys
    th = np.arctan2(dy, dx) % TAU
    rr = np.sqrt(dx * dx + dy * dy)
    model = c * (1.0 + L2 * np.cos(th - phi))
    return float(np.mean((rr - model) ** 2))


_DEFAULT_DESIRED = {
    "circle": {
        "radius": lambda p, xs, ys: float(p[0]),
    },
    "ellipse": {
        "area": lambda p, xs, ys: float(np.pi * (0.5 * p[0]) * (0.5 * p[1])),
    },
    "limacon": {
        "c": lambda p, xs, ys: float(p[0]),
        "phi": lambda p, xs, ys: float(p[2]),
    },
}


def general_fit(points, xs, ys, shape="circle", method="Powell",
                bounds=None, return_desired=False, desired=None, **kwargs):
    """
    General dispatcher for fitting supported shapes.

    **Args**:
    * points (array-like): (N,2) sample points.
    * xs (float): fixed x-center.
    * ys (float): fixed y-center.
    * shape (str): "circle", "ellipse", or "limacon".
    * method (str): optimization method (kept for API compatibility).
    * bounds (list or None): parameter bounds (kept for API compatibility).
    * return_desired (bool): if True attach diagnostics under "desired".
    * desired (dict or None): optional custom desired dict of callables.
    * **kwargs: ignored extra options for forward compatibility.

    **Returns**:
    * res (dict): standardized fit result with keys:
      shape, params, center, success, cost[, desired]

    """
    pts = _validate_points(points)
    xs = float(xs)
    ys = float(ys)
    shape = str(shape).lower().strip()

    if shape == "circle":
        r_est = fit_circle(pts, xs, ys)
        params = (float(r_est),)
        cost = _circle_cost(pts, xs, ys, r_est)

    elif shape == "ellipse":
        d2a_est, d2b_est = fit_ellipse(pts, xs, ys)
        params = (float(d2a_est), float(d2b_est))
        cost = _ellipse_cost(pts, xs, ys, d2a_est, d2b_est)

    elif shape == "limacon":
        c_est, L2_est, phi_est = fit_limacon(pts, xs, ys)
        params = (float(c_est), float(L2_est), float(phi_est))
        cost = _limacon_cost(pts, xs, ys, c_est, L2_est, phi_est)

    else:
        raise ValueError(f"Unknown shape '{shape}'")

    res = {
        "shape": shape,
        "params": params,
        "center": (xs, ys),
        "success": True,
        "cost": float(max(0.0, cost)),
    }

    if return_desired:
        desired_funcs = desired if desired is not None else _DEFAULT_DESIRED.get(shape, {})
        desired_vals = {}
        p_arr = np.asarray(params, float)
        for k, fn in desired_funcs.items():
            try:
                desired_vals[k] = fn(p_arr, xs, ys)
            except Exception:
                desired_vals[k] = None
        res["desired"] = desired_vals

    return res
