"""Minkowski difference approximations for pairs of ellipsoids.

External approximation (minkdiff_ea):
    q = q1 - q2
    Q = M^T M,   M = Q1^{1/2} - S Q2^{1/2}
    where S aligns Q2^{1/2} l to Q1^{1/2} l.

Internal approximation (minkdiff_ia):
    q = q1 - q2
    Q = (1 - 1/p) Q1 + (1 - p) Q2
    where p = sqrt(<l, Q1 l>) / sqrt(<l, Q2 l>).

The geometric difference E1 - E2 is non-empty only when E1 >= E2
(i.e. E1 would contain E2 if both were centered at the origin).

Reference
---------
MATLAB Ellipsoidal Toolbox v3,
``@ellipsoid/minkdiff_ea.m``, ``minkdiff_ia.m``, ``isbaddirection.m``,
``private/rm_bad_directions.m``.
"""

from __future__ import annotations

import numpy as np
import scipy.linalg

from ..core.config import ABS_TOL
from ..core.ellipsoid import Ellipsoid
from ..core.helpers import ell_simdiag, ell_valign, regularize

__all__ = ["minkdiff_ea", "minkdiff_ia", "isbaddirection"]


# ------------------------------------------------------------------
# Public helpers
# ------------------------------------------------------------------

def isbaddirection(
    E1: Ellipsoid,
    E2: Ellipsoid,
    L: np.ndarray,
) -> np.ndarray:
    """Check which directions are "bad" for the geometric difference E1 - E2.

    A direction *l* is bad when the ellipsoidal approximation of E1 - E2
    in that direction cannot be computed.  This happens when

        p = sqrt(<l, Q1 l>) / sqrt(<l, Q2 l>)  >=  a

    where *a* is the minimum generalised eigenvalue of Q1 w.r.t. Q2
    (computed via simultaneous diagonalisation).

    Parameters
    ----------
    E1, E2 : Ellipsoid
        The two operands of the geometric difference.
    L : array_like (n,) or (n, k)
        Direction vector(s).

    Returns
    -------
    ndarray of bool, shape (k,)
        ``True`` at positions corresponding to bad directions.
    """
    L_arr = np.asarray(L, dtype=float)
    if L_arr.ndim == 1:
        L_arr = L_arr.reshape(-1, 1)
    k = L_arr.shape[1]

    if E1.is_empty() or E2.is_empty():
        return np.ones(k, dtype=bool)
    if not E1.is_bigger(E2):
        return np.ones(k, dtype=bool)

    Q1, Q2 = E1.shape, E2.shape
    assert Q1 is not None and Q2 is not None
    n = Q1.shape[0]

    if L_arr.shape[0] != n:
        raise ValueError(
            f"isbaddirection: direction dimension ({L_arr.shape[0]}) "
            f"does not match ellipsoid dimension ({n})."
        )

    # Simultaneous diagonalisation: T Q2 T' = I,  T Q1 T' = diag
    T = ell_simdiag(Q2, Q1)
    a = float(np.min(np.abs(np.diag(T @ Q1 @ T.T))))

    bad = np.empty(k, dtype=bool)
    for i in range(k):
        l = L_arr[:, i]
        lQ1l = float(l @ Q1 @ l)
        lQ2l = float(l @ Q2 @ l)
        if lQ2l < ABS_TOL:
            bad[i] = True
            continue
        p = np.sqrt(max(lQ1l, 0.0)) / np.sqrt(max(lQ2l, 0.0))
        bad[i] = not (a > p)
    return bad


# ------------------------------------------------------------------
# Private helpers
# ------------------------------------------------------------------

def _rm_bad_directions(
    Q1: np.ndarray,
    Q2: np.ndarray,
    L: np.ndarray,
) -> np.ndarray:
    """Filter *L* to keep only directions where the approximation is valid.

    Mirrors MATLAB ``private/rm_bad_directions.m``.

    Parameters
    ----------
    Q1, Q2 : ndarray (n, n)
        Shape matrices (already regularised if needed).
    L : ndarray (n, k)
        Direction columns.

    Returns
    -------
    ndarray (n, m)
        Subset of columns of *L* that are "good".  May be empty (0 columns).
    """
    T = ell_simdiag(Q2, Q1)
    a = float(np.min(np.abs(np.diag(T @ Q1 @ T.T))))
    if a < 1.0:
        return np.empty((Q1.shape[0], 0))

    good_cols = []
    for i in range(L.shape[1]):
        l = L[:, i]
        lQ1l = float(l @ Q1 @ l)
        lQ2l = float(l @ Q2 @ l)
        if lQ2l < ABS_TOL:
            continue
        p = np.sqrt(max(lQ1l, 0.0)) / np.sqrt(max(lQ2l, 0.0))
        if p <= a:
            good_cols.append(l)

    if len(good_cols) == 0:
        return np.empty((Q1.shape[0], 0))
    return np.column_stack(good_cols)


# ------------------------------------------------------------------
# Main functions
# ------------------------------------------------------------------

def minkdiff_ea(
    E1: Ellipsoid,
    E2: Ellipsoid,
    L: np.ndarray,
) -> list[Ellipsoid]:
    """External approximating ellipsoids for the geometric difference E1 - E2.

    Parameters
    ----------
    E1, E2 : Ellipsoid
        Single ellipsoids of the same dimension.
    L : array_like (n,) or (n, k)
        Direction vector(s).

    Returns
    -------
    list of Ellipsoid
        One external approximation per *good* direction.
        Empty list if the difference is empty or no good directions exist.
    """
    if E1.is_empty() or E2.is_empty():
        raise ValueError("minkdiff_ea: both ellipsoids must be non-empty.")

    sd1, _ = E1.dimension()
    sd2, _ = E2.dimension()
    if sd1 != sd2:
        raise ValueError("minkdiff_ea: ellipsoids must have the same dimension.")

    if not E1.is_bigger(E2):
        return []

    L_arr = np.asarray(L, dtype=float)
    if L_arr.ndim == 1:
        L_arr = L_arr.reshape(-1, 1)
    n = sd1
    if L_arr.shape[0] != n:
        raise ValueError(
            f"minkdiff_ea: direction dimension ({L_arr.shape[0]}) "
            f"does not match ellipsoid dimension ({n})."
        )

    q = E1.center - E2.center
    Q1 = E1.shape.copy()
    Q2 = E2.shape.copy()

    L_good = _rm_bad_directions(Q1, Q2, L_arr)
    if L_good.shape[1] == 0:
        return []

    if np.linalg.matrix_rank(Q1, tol=ABS_TOL) < n:
        Q1 = regularize(Q1)
    if np.linalg.matrix_rank(Q2, tol=ABS_TOL) < n:
        Q2 = regularize(Q2)

    Q1_sqrt = np.real(scipy.linalg.sqrtm(Q1))
    Q2_sqrt = np.real(scipy.linalg.sqrtm(Q2))

    result = []
    for j in range(L_good.shape[1]):
        l = L_good[:, j]
        # MATLAB: T = ell_valign(Q1*l, Q2*l)  ->  T*(Q2*l) || Q1*l
        # Python: ell_valign(v1, v2) -> S*v1 || v2
        # So:     ell_valign(Q2_sqrt @ l, Q1_sqrt @ l) -> S*(Q2_sqrt@l) || Q1_sqrt@l
        S = ell_valign(Q2_sqrt @ l, Q1_sqrt @ l)
        M = Q1_sqrt - S @ Q2_sqrt
        Q_res = M.T @ M
        Q_res = 0.5 * (Q_res + Q_res.T)
        result.append(Ellipsoid(q, Q_res))

    return result


def minkdiff_ia(
    E1: Ellipsoid,
    E2: Ellipsoid,
    L: np.ndarray,
) -> list[Ellipsoid]:
    """Internal approximating ellipsoids for the geometric difference E1 - E2.

    Parameters
    ----------
    E1, E2 : Ellipsoid
        Single ellipsoids of the same dimension.
    L : array_like (n,) or (n, k)
        Direction vector(s).

    Returns
    -------
    list of Ellipsoid
        One internal approximation per *good* direction.
        Empty list if the difference is empty or no good directions exist.
    """
    if E1.is_empty() or E2.is_empty():
        raise ValueError("minkdiff_ia: both ellipsoids must be non-empty.")

    sd1, _ = E1.dimension()
    sd2, _ = E2.dimension()
    if sd1 != sd2:
        raise ValueError("minkdiff_ia: ellipsoids must have the same dimension.")

    if not E1.is_bigger(E2):
        return []

    L_arr = np.asarray(L, dtype=float)
    if L_arr.ndim == 1:
        L_arr = L_arr.reshape(-1, 1)
    n = sd1
    if L_arr.shape[0] != n:
        raise ValueError(
            f"minkdiff_ia: direction dimension ({L_arr.shape[0]}) "
            f"does not match ellipsoid dimension ({n})."
        )

    q = E1.center - E2.center
    Q1 = E1.shape.copy()
    Q2 = E2.shape.copy()

    if np.linalg.matrix_rank(Q1, tol=ABS_TOL) < n:
        Q1 = regularize(Q1)
    if np.linalg.matrix_rank(Q2, tol=ABS_TOL) < n:
        Q2 = regularize(Q2)

    L_good = _rm_bad_directions(Q1, Q2, L_arr)
    if L_good.shape[1] == 0:
        return []

    result = []
    for j in range(L_good.shape[1]):
        l = L_good[:, j]
        lQ1l = float(l @ Q1 @ l)
        lQ2l = float(l @ Q2 @ l)
        p = np.sqrt(max(lQ1l, 0.0)) / np.sqrt(max(lQ2l, ABS_TOL))
        Q_res = (1.0 - 1.0 / p) * Q1 + (1.0 - p) * Q2
        Q_res = 0.5 * (Q_res + Q_res.T)
        result.append(Ellipsoid(q, Q_res))

    return result
