"""Minkowski sum approximations for collections of ellipsoids.

External approximation (minksum_ea):
    q = q1 + ... + qm
    Q = (p1 + ... + pm) * ((1/p1)*Q1 + ... + (1/pm)*Qm)
    where p_i = sqrt(<l, Q_i l>)

Internal approximation (minksum_ia):
    q = q1 + ... + qm
    Q = M^T M,  M = Q1^{1/2} + S2 Q2^{1/2} + ... + Sm Qm^{1/2}
    where S_i are orthogonal matrices aligning Q_i^{1/2} l to Q_1^{1/2} l.

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

from __future__ import annotations

from typing import Sequence

import numpy as np
import scipy.linalg

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

__all__ = ["minksum_ea", "minksum_ia"]


def _validate_ellipsoids_and_directions(
    ellipsoids: Sequence[Ellipsoid],
    L: np.ndarray,
    name: str,
) -> tuple[int, np.ndarray]:
    """Shared validation for Minkowski sum functions.

    Returns
    -------
    n : int
        Common dimension of the ellipsoids.
    L_arr : ndarray (n, k)
        Direction matrix (always 2-D).
    """
    if len(ellipsoids) == 0:
        raise ValueError(f"{name}: need at least one ellipsoid.")
    if any(e.is_empty() for e in ellipsoids):
        raise ValueError(f"{name}: all ellipsoids must be non-empty.")

    dims = [e.dimension()[0] for e in ellipsoids]
    if len(set(dims)) != 1:
        raise ValueError(f"{name}: all ellipsoids must have the same dimension.")
    n = dims[0]

    L_arr = np.asarray(L, dtype=float)
    if L_arr.ndim == 1:
        L_arr = L_arr.reshape(-1, 1)
    if L_arr.shape[0] != n:
        raise ValueError(
            f"{name}: direction vectors must be in R^{n}, "
            f"got R^{L_arr.shape[0]}."
        )
    return n, L_arr


def minksum_ea(
    ellipsoids: Sequence[Ellipsoid],
    L: np.ndarray,
) -> list[Ellipsoid]:
    """External approximating ellipsoids for the Minkowski sum.

    Parameters
    ----------
    ellipsoids : sequence of Ellipsoid
        Ellipsoids to sum. All must share the same dimension.
    L : array_like (n,) or (n, k)
        Direction vector(s).

    Returns
    -------
    list of Ellipsoid
        One tight external approximation per direction column.
    """
    n, L_arr = _validate_ellipsoids_and_directions(
        ellipsoids, L, "minksum_ea"
    )

    if len(ellipsoids) == 1:
        return [ellipsoids[0]] * L_arr.shape[1]

    centers = []
    shapes = []
    for e in ellipsoids:
        q, Q = e.parameters()
        centers.append(q)
        sd, rk = e.dimension()
        if rk < sd:
            Q = regularize(Q)
        shapes.append(Q)

    q_sum = sum(centers)

    result = []
    for j in range(L_arr.shape[1]):
        l = L_arr[:, j]
        S = np.zeros((n, n))
        P_total = 0.0
        for Q in shapes:
            p = np.sqrt(max(float(l @ Q @ l), 0.0))
            if p < ABS_TOL:
                p = ABS_TOL
            S += (1.0 / p) * Q
            P_total += p
        S = 0.5 * P_total * (S + S.T)
        result.append(Ellipsoid(q_sum, S))

    return result


def minksum_ia(
    ellipsoids: Sequence[Ellipsoid],
    L: np.ndarray,
) -> list[Ellipsoid]:
    """Internal approximating ellipsoids for the Minkowski sum.

    Parameters
    ----------
    ellipsoids : sequence of Ellipsoid
        Ellipsoids to sum. All must share the same dimension.
    L : array_like (n,) or (n, k)
        Direction vector(s).

    Returns
    -------
    list of Ellipsoid
        One tight internal approximation per direction column.
    """
    n, L_arr = _validate_ellipsoids_and_directions(
        ellipsoids, L, "minksum_ia"
    )

    if len(ellipsoids) == 1:
        return [ellipsoids[0]] * L_arr.shape[1]

    centers = []
    sqrt_shapes = []
    for e in ellipsoids:
        q, Q = e.parameters()
        centers.append(q)
        sd, rk = e.dimension()
        if rk < sd:
            Q = regularize(Q)
        sqQ = np.real(scipy.linalg.sqrtm(Q))
        sqrt_shapes.append(sqQ)

    q_sum = sum(centers)

    result = []
    for j in range(L_arr.shape[1]):
        l = L_arr[:, j]
        v_ref = sqrt_shapes[0] @ l
        M = sqrt_shapes[0].copy()
        for i in range(1, len(sqrt_shapes)):
            v_i = sqrt_shapes[i] @ l
            # ell_valign(v_i, v_ref) -> T such that T*v_i || v_ref
            T = ell_valign(v_i, v_ref)
            M = M + T @ sqrt_shapes[i]
        Q_result = M.T @ M
        Q_result = 0.5 * (Q_result + Q_result.T)
        result.append(Ellipsoid(q_sum, Q_result))

    return result
