"""2D ellipse to 3D covariance and back."""

from __future__ import annotations

from typing import TYPE_CHECKING, NamedTuple

import numpy as np

from gri_utils.matrices import decompose
from gri_utils.matrices.rotate import rotate_cw

from .cov_sigma_95 import scale_chi2

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


class EllipseParams(NamedTuple):
    """Result of 2D covariance to ellipse parameter conversion.

    Attributes:
        sma: Semi-major axis length.
        smi: Semi-minor axis length.
        ori_deg: Orientation in degrees clockwise from North.
    """

    sma: float
    smi: float
    ori_deg: float


class EllipseParams95(NamedTuple):
    """Result of covariance to 95% ellipse parameter conversion.

    Attributes:
        sma_95: Semi-major axis length at 95% confidence.
        smi_95: Semi-minor axis length at 95% confidence.
        ori_deg: Orientation in degrees clockwise from North.
        alt_95: Altitude at 95% confidence (0 if input was 2x2).
    """

    sma_95: float
    smi_95: float
    ori_deg: float
    alt_95: float


class CovSlice(NamedTuple):
    """Result of slicing a 3D covariance ellipsoid.

    Attributes:
        cov_2d: 2x2 covariance matrix of the ellipse slice.
        center: 3D center point of the ellipse slice.
    """

    cov_2d: NDArray[np.floating]
    center: NDArray[np.floating]


def params_95_to_cov(
    sma_95: float,
    smi_95: float,
    ori_deg: float,
    alt_95: float | None = None,
) -> NDArray[np.floating]:
    """Convert ellipse parameters to a covariance matrix in ENU.

    Standard for converting 95% parameters to a 2x2 or 3x3 covariance matrix.
    SMA and SMI are converted using 2D chi**2 scaling and alt_95 using 1D.

    Args:
        sma_95(float): Semi major 95th percent
        smi_95(float): Semi minor 95th percent
        ori_deg(float): SMA degrees clockwise from North
        alt_95(float,optional): Altitude 95th percent

    Returns:
        np.ndarray of covariance in length**2, ENU. 2x2 if no alt, else 3x3
    """
    sma = sma_95 / scale_chi2(1, dof=2)
    smi = smi_95 / scale_chi2(1, dof=2)
    alt = alt_95 / scale_chi2(1) if alt_95 is not None else None
    return params_to_cov(sma, smi, ori_deg, alt)


def params_to_cov(
    sma: float,
    smi: float,
    ori_deg: float,
    alt: float | None = None,
) -> NDArray[np.floating]:
    """Convert ellipse parameters to a covariance matrix in ENU.

    The covariance matrix will be in whichever distance units (squared) and scaling the
    sma and smi were passed in as. If no alt is passed in, a 2x2 will be returned,
    otherwise a 3x3 will be returned.

    Args:
        sma(float): Semi major axis length
        smi(float): Semi minor axis length
        ori_deg(float): SMA degrees clockwise from North
        alt(float,optional): Altitude semi-axis length

    Returns:
        np.ndarray of covariance in length**2, ENU. 2x2 if no alt, else 3x3
    """
    if alt is not None:
        cov = np.array(((smi**2, 0, 0), (0, sma**2, 0), (0, 0, alt**2)))
    else:
        cov = np.array(((smi**2, 0), (0, sma**2)))
    return rotate_cw(np.radians(ori_deg), cov)


def cov_to_params_95(
    cov: NDArray[np.floating] | Sequence[Sequence[float | int]],
) -> EllipseParams95:
    """Covariance (2x2 or 3x3) into SMA, SMI, ORI (and ALT).

    Get the parameters of Semi Major Axis, Semi Minor Axis, and Orientation (degrees
    clockwise from North) from a 2x2 covariance array. Include Alt 95 if a 3x3.

    Args:
        cov(np.ndarray|Sequence[Sequence[float|int]]): Covariance 2x2 or 3x3 array in
            1-sigma.

    Returns:
        EllipseParams95 named tuple with sma_95, smi_95, ori_deg, alt_95.
    """
    cov = np.asarray(cov)
    if cov.shape == (3, 3):
        cov_2d = cov_projection(cov)
        alt_95 = scale_chi2(float(cov[2, 2]) ** 0.5)
    elif cov.shape == (2, 2):
        cov_2d = cov
        alt_95 = 0.0
    else:
        raise ValueError("Covariance shape is not (2,2) or (3,3)")
    sma, smi, ori_deg = cov_to_params(cov_2d)
    sma_95 = scale_chi2(sma, dof=2)
    smi_95 = scale_chi2(smi, dof=2)
    return EllipseParams95(sma_95, smi_95, ori_deg, alt_95)


def cov_to_params(
    cov_2d: NDArray[np.floating] | Sequence[Sequence[float | int]],
) -> EllipseParams:
    """2D Covariance (2x2) into SMA, SMI, ORI.

    Get the parameters of Semi Major Axis, Semi Minor Axis, and Orientation (degrees
    clockwise from North) from a 2x2 covariance array.

    Note: Units are same as cov_2d: Scale to 95% afterwards if desired.

    Args:
        cov_2d(np.ndarray|Sequence[Sequence[float|int]]): Covariance 2x2 array.

    Returns:
        EllipseParams named tuple with sma, smi, ori_deg. Units and scale will be the
            same as cov_2d.
    """
    cov_2d = np.asarray(cov_2d)
    if cov_2d.shape != (2, 2):
        raise ValueError("Covariance shape is not (2,2)")
    vals, uvecs = decompose(cov_2d)
    sma = float(abs(vals[0]) ** 0.5)
    smi = float(abs(vals[1]) ** 0.5)
    # clockwise direction
    cw_from_y = np.arctan2(uvecs[0][0], uvecs[0][1])
    rad = np.round(cw_from_y, 15) % np.pi
    ori_deg = float(np.degrees(rad))
    return EllipseParams(sma, smi, ori_deg)


def cov_to_alt_param(
    cov_3d: NDArray[np.floating] | Sequence[Sequence[float | int]],
) -> float:
    """Get an altitude parameter from a 3D covariance matrix.

    Use cov_projection or cov_slice to get the 2x2 matrix for SMA/SMI/ORI, and this for
    the unscaled alt.

    Args:
        cov_3d(np.ndarray|Sequence[Sequence[float|int]]): Covariance 3x3 array.

    Returns:
        altitude parameter to pair with SMA, SMI, ORI
    """
    cov_3d = np.asarray(cov_3d)
    if cov_3d.shape != (3, 3):
        raise ValueError("Covariance shape is not (3,3)")
    # Get 3rd axis height
    return abs(cov_3d[2, 2]) ** 0.5


def cov_projection(
    cov: NDArray[np.floating] | Sequence[Sequence[float | int]],
) -> NDArray[np.floating]:
    """Return the 2D projection of a 3x3 matrix in the first two axis.

    Args:
        cov(np.ndarray|Sequence[Sequence[float|int]]): A 3x3 covariance matrix

    Returns:
        2x2 projection of the 3x3 in the first two axis.
    """
    return np.asarray(cov)[:2, :2]


def cov_slice(
    cov_3d: NDArray[np.floating] | Sequence[Sequence[float | int]],
    z_diff: float | None = None,
    center: NDArray[np.floating] | Sequence[float | int] | None = None,
) -> CovSlice:
    """Return the 2x2 ellipse slice covariance from a 3x3 and its center point.

    Taking a 3x3 covariance matrix and an optional 3rd axis height, return the 2x2
    covariance representing the ellipse at that slice of the ellipsoid and the new
    center point of that ellipse, sized (3,).

    Args:
        cov_3d(np.ndarray): 3x3 covariance
        z_diff(float, optional): The 3rd axis height of the slice relative to center,
            defaults to 0
        center(np.ndarray, optional): 3d vector of the current center of the ellipsoid.
            Defaults to (0,0,0)

    Returns:
        CovSlice named tuple with cov_2d (2x2) and center (3,).
    """
    cov_3d = np.asarray(cov_3d)
    if cov_3d.shape != (3, 3):
        raise ValueError("Covariance shape is not (3,3)")
    center = np.array([0.0, 0.0, 0.0]) if center is None else np.asarray(center)
    if z_diff is None:
        z_diff = 0

    # Check if z_height is within the ellipsoid
    if abs(z_diff) >= np.sqrt(abs(cov_3d[2, 2])):
        raise ValueError(f"{z_diff} is outside the ellipsoid.")

    info = np.linalg.inv(cov_3d)
    # Extract blocks from the information covariance
    xy_mat = info[:2, :2]  # 2x2 block is the form of the ellipse
    xy_cov = np.linalg.inv(xy_mat)  # covariance of slice at center
    z_vec = info[:2, 2]  # 2x1 vector of z variance in x,y

    # The 2D ellipse center shifts based on height
    mean_2d = center[:2] - (xy_cov @ z_vec) * z_diff
    mean_3d = np.append(mean_2d, z_diff + center[2])

    # remaining_radius_sq = 1 - z_val * z_diff**2
    remaining_radius_sq = 1 - z_diff**2 / cov_3d[2, 2]
    if remaining_radius_sq <= 0:
        # Should be equal to above, but floating points ...
        raise ValueError(f"{z_diff} is outside the ellipsoid (near boundary).")

    # The 2D covariance of the slice
    cov_2d = xy_cov * remaining_radius_sq

    return CovSlice(cov_2d, mean_3d)
