"""Transform XYZ (ECEF) covariance matrices <=> ENU covariance matrices.

These functions handle 3x3 covariance matrices (or stacks of them) and transform
them between XYZ and ENU coordinate frames using similarity transforms.

For single matrices:
    xyz_to_enu_cov(xyz_pos, xyz_cov) -> enu_cov
    enu_to_xyz_cov(xyz_pos, enu_cov) -> xyz_cov

For vectorized stacks of matrices (N, 3, 3):
    xyz_to_enu_cov(xyz_pos, xyz_covs) -> enu_covs
    enu_to_xyz_cov(xyz_pos, enu_covs) -> xyz_covs
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from .xyz_enu_rotation import enu_xyz_rotation, xyz_enu_rotation

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def xyz_to_enu_cov(
    xyz_pos: NDArray[np.floating] | Sequence,
    xyz_matrix: NDArray[np.floating],
) -> NDArray[np.floating]:
    """Convert XYZ covariance matrix(s) to ENU at the given position.

    Uses the similarity transform: enu_cov = T @ xyz_cov @ T.T
    where T is the XYZ->ENU rotation matrix.

    Args:
        xyz_pos(np.ndarray|Sequence): XYZ reference location. Shape (3,).
        xyz_matrix(np.ndarray): XYZ covariance matrix(s). Shape (3, 3) for single
            matrix or (N, 3, 3) for N matrices.

    Returns:
        np.ndarray: ENU covariance matrix(s). Shape (3, 3) or (N, 3, 3) matching
            xyz_matrix.
    """
    t_enu = xyz_enu_rotation(xyz_pos)
    xyz_matrix = np.asarray(xyz_matrix)
    single = xyz_matrix.ndim == 2  # noqa: PLR2004
    if single:
        xyz_matrix = xyz_matrix[np.newaxis, ...]

    # Vectorized: (N, 3, 3) -> use einsum for batch matrix multiplication
    # Result[i] = T @ xyz_matrix[i] @ T.T
    temp = np.einsum("ij,njk->nik", t_enu, xyz_matrix)
    result = np.einsum("nij,kj->nik", temp, t_enu)

    return result[0] if single else result


def enu_to_xyz_cov(
    xyz_pos: NDArray[np.floating] | Sequence,
    enu_matrix: NDArray[np.floating],
) -> NDArray[np.floating]:
    """Convert ENU covariance matrix(s) to XYZ at the given position.

    Uses the similarity transform: xyz_cov = T @ enu_cov @ T.T
    where T is the ENU->XYZ rotation matrix.

    Args:
        xyz_pos(np.ndarray|Sequence): XYZ reference location. Shape (3,).
        enu_matrix(np.ndarray): ENU covariance matrix(s). Shape (3, 3) for single
            matrix or (N, 3, 3) for N matrices.

    Returns:
        np.ndarray: XYZ covariance matrix(s). Shape (3, 3) or (N, 3, 3) matching
            enu_matrix.
    """
    t_xyz = enu_xyz_rotation(xyz_pos)
    enu_matrix = np.asarray(enu_matrix)
    single = enu_matrix.ndim == 2  # noqa: PLR2004
    if single:
        enu_matrix = enu_matrix[np.newaxis, ...]

    # Vectorized: (N, 3, 3) -> use einsum for batch matrix multiplication
    # Result[i] = T @ enu_matrix[i] @ T.T
    temp = np.einsum("ij,njk->nik", t_xyz, enu_matrix)
    result = np.einsum("nij,kj->nik", temp, t_xyz)

    return result[0] if single else result
