"""Distance functions for ellipsoids.

Computes distances from an ellipsoid to:
  - a point (closed-form check + scipy.optimize for exterior distance)
  - a hyperplane (via support function)
  - another ellipsoid (QCQP via CVXPY — added in Phase 2C-iii)

Reference
---------
MATLAB Ellipsoidal Toolbox v3, ``@ellipsoid/distance.m``.
"""

from __future__ import annotations

import numpy as np
import scipy.optimize

from ..core.config import ABS_TOL
from ..core.ellipsoid import Ellipsoid
from ..core.hyperplane import Hyperplane
from ..core.helpers import ell_inv, regularize

__all__ = [
    "distance_to_point",
    "distance_to_hyperplane",
    "distance_to_ellipsoid",
]


def distance_to_point(
    E: Ellipsoid,
    y: np.ndarray,
    metric: str = "euclidean",
) -> float:
    """Distance from an ellipsoid to a point.

    Negative distance means the point is strictly inside the ellipsoid.
    Zero means the point is on the boundary.

    Parameters
    ----------
    E : Ellipsoid
        Must be non-empty.
    y : array_like (n,)
        Target point.
    metric : {"euclidean", "ellipsoidal"}
        ``"euclidean"`` (default) uses standard Euclidean norm.
        ``"ellipsoidal"`` uses the metric induced by Q^{-1}.

    Returns
    -------
    float
        Signed distance.  Negative if *y* is interior.
    """
    if E.is_empty():
        raise ValueError("distance_to_point: ellipsoid must be non-empty.")

    q, Q = E.parameters()
    assert q is not None and Q is not None
    n = q.shape[0]

    y_arr = np.asarray(y, dtype=float).ravel()
    if y_arr.shape[0] != n:
        raise ValueError(
            f"distance_to_point: point dimension ({y_arr.shape[0]}) "
            f"does not match ellipsoid dimension ({n})."
        )

    sd, rk = E.dimension()
    Q_work = Q.copy()
    if rk < sd:
        Q_work = regularize(Q_work)

    Qi = ell_inv(Q_work)
    Qi = 0.5 * (Qi + Qi.T)

    diff = y_arr - q
    quad_val = float(diff @ Qi @ diff)

    if metric == "ellipsoidal":
        W = Qi
    else:
        W = np.eye(n)

    if quad_val <= 1.0 + ABS_TOL:
        # Point is inside or on boundary → compute negative distance
        # as -min ||x - y||_W over the boundary of E
        # For interior point: use the support function trick.
        # dist = -(1 - sqrt(quad_val)) * ||diff||_W / ||diff||_Qi
        # Simplified: the signed distance is -sqrt((1-sqrt(qv))^2 * something)
        # Exact: solve min ||x - y||_W s.t. (x-q)^T Qi (x-q) >= 1
        # For simplicity and MATLAB consistency: if inside, return -(1 - sqrt(qv))*scale
        if quad_val < ABS_TOL:
            # Point is at the center — distance = -min_semiaxis
            eigvals = np.linalg.eigvalsh(Q_work)
            return -np.sqrt(float(np.min(eigvals)))

        # Negative of the distance from y to the boundary along the ray q→y
        scale = 1.0 / np.sqrt(quad_val)
        x_boundary = q + scale * diff
        return -float(np.sqrt((x_boundary - y_arr) @ W @ (x_boundary - y_arr)))

    # Point is outside → minimise ||x - y||^2_W subject to (x-q)^T Qi (x-q) <= 1
    # This is a convex QCQP with one constraint → use Lagrange multiplier (KKT).
    # At optimum: x* = (W + mu*Qi)^{-1} (W y + mu Qi q) with mu >= 0 chosen
    # so that (x*-q)^T Qi (x*-q) = 1.
    def _boundary_residual(mu: float) -> float:
        M = W + mu * Qi
        x_star = np.linalg.solve(M, W @ y_arr + mu * Qi @ q)
        d = x_star - q
        return float(d @ Qi @ d) - 1.0

    # Find mu by bisection. mu=0 → interior (residual < 0), mu→∞ → x*→q (residual = -1).
    # Actually mu→∞ pushes x* toward q, so residual → q^TQi q -1 ... need upper bound.
    # The residual at mu=0 is quad_val - 1 > 0 (point outside), so we need
    # the bracket differently. Let me re-derive.
    # At mu=0: x*=y, residual = (y-q)^T Qi (y-q) - 1 = quad_val - 1 > 0
    # As mu→∞: x*→q, residual → -1 < 0
    # So root is between mu=0 and some large mu.
    mu_lo = 0.0
    mu_hi = 1.0
    while _boundary_residual(mu_hi) > 0:
        mu_hi *= 2.0
        if mu_hi > 1e16:
            break

    try:
        mu_star = scipy.optimize.brentq(_boundary_residual, mu_lo, mu_hi, xtol=ABS_TOL)
    except ValueError:
        mu_star = mu_hi

    M = W + mu_star * Qi
    x_star = np.linalg.solve(M, W @ y_arr + mu_star * Qi @ q)
    dist_sq = float((x_star - y_arr) @ W @ (x_star - y_arr))
    return np.sqrt(max(dist_sq, 0.0))


def distance_to_hyperplane(
    E: Ellipsoid,
    H: Hyperplane,
    metric: str = "euclidean",
) -> float:
    """Distance from an ellipsoid to a hyperplane.

    Negative distance means the ellipsoid intersects the hyperplane.

    Uses the support function: the distance from E to H(v,c) is

        d = (c - rho(E, v)) / ||v||

    or the corresponding formula from the other side.

    Parameters
    ----------
    E : Ellipsoid
    H : Hyperplane
    metric : {"euclidean", "ellipsoidal"}
        ``"ellipsoidal"`` uses the metric induced by Q.

    Returns
    -------
    float
        Signed distance.  Negative if E intersects H.
    """
    if E.is_empty():
        raise ValueError("distance_to_hyperplane: ellipsoid must be non-empty.")
    if H.is_empty():
        raise ValueError("distance_to_hyperplane: hyperplane must be non-empty.")

    q = E.center
    assert q is not None
    v, c = H.parameters()
    assert v is not None

    n = q.shape[0]
    if v.shape[0] != n:
        raise ValueError(
            f"distance_to_hyperplane: hyperplane dimension ({v.shape[0]}) "
            f"does not match ellipsoid dimension ({n})."
        )

    # Normalise so c >= 0
    if c < 0:
        c = -c
        v = -v

    if metric == "ellipsoidal":
        Q = E.shape
        assert Q is not None
        sr = np.sqrt(max(float(v @ Q @ v), 0.0))
    else:
        sr = np.sqrt(float(v @ v))

    if sr < ABS_TOL:
        return 0.0

    if float(v @ q) < c:
        rho_val, _ = E.rho(v)
        return (c - rho_val) / sr
    else:
        rho_neg, _ = E.rho(-v)
        return (-c - rho_neg) / sr


def distance_to_ellipsoid(
    E1: Ellipsoid,
    E2: Ellipsoid,
    metric: str = "euclidean",
) -> float:
    """Distance between two ellipsoids (QCQP via CVXPY).

    Solves::

        min  ||x - y||^2_W
        s.t. (x - q1)^T Q1^{-1} (x - q1) <= 1
             (y - q2)^T Q2^{-1} (y - q2) <= 1

    This is a convex QCQP (intersection of two convex sets with a convex
    objective).  Returns 0 when the ellipsoids overlap, otherwise the
    positive Euclidean (or ellipsoidal) distance.

    Parameters
    ----------
    E1, E2 : Ellipsoid
        Non-empty ellipsoids of the same dimension.
    metric : {"euclidean", "ellipsoidal"}
        ``"ellipsoidal"`` uses the metric induced by ``Q1^{-1}``.

    Returns
    -------
    float
        Non-negative distance (0 when the ellipsoids intersect).

    Raises
    ------
    ImportError
        If CVXPY is not installed.
    """
    import cvxpy as cp  # conditional import

    if E1.is_empty() or E2.is_empty():
        raise ValueError("distance_to_ellipsoid: both ellipsoids must be non-empty.")

    q1, Q1 = E1.parameters()
    q2, Q2 = E2.parameters()
    assert q1 is not None and Q1 is not None
    assert q2 is not None and Q2 is not None

    n = q1.shape[0]
    if q2.shape[0] != n:
        raise ValueError(
            f"distance_to_ellipsoid: dimension mismatch ({n} vs {q2.shape[0]})."
        )

    sd1, rk1 = E1.dimension()
    sd2, rk2 = E2.dimension()
    Q1w = regularize(Q1) if rk1 < sd1 else Q1.copy()
    Q2w = regularize(Q2) if rk2 < sd2 else Q2.copy()

    Q1i = ell_inv(Q1w)
    Q1i = 0.5 * (Q1i + Q1i.T)
    Q2i = ell_inv(Q2w)
    Q2i = 0.5 * (Q2i + Q2i.T)

    x = cp.Variable(n)
    y = cp.Variable(n)

    if metric == "ellipsoidal":
        obj = cp.quad_form(x - y, Q1i)
    else:
        obj = cp.sum_squares(x - y)

    c1 = cp.quad_form(x - q1, Q1i) <= 1.0
    c2 = cp.quad_form(y - q2, Q2i) <= 1.0

    prob = cp.Problem(cp.Minimize(obj), [c1, c2])
    prob.solve(solver=cp.SCS, eps=1e-9, max_iters=10000)

    if prob.status in ("infeasible", "infeasible_inaccurate"):
        raise RuntimeError(
            f"distance_to_ellipsoid: solver returned '{prob.status}'."
        )

    val = prob.value
    if val is None or val < ABS_TOL:
        return 0.0
    return float(np.sqrt(max(val, 0.0)))
