"""Calculate the osculating sphere at a point on the WGS84 ellipsoid.

The osculating sphere is the sphere that best approximates the ellipsoid
at a given point, having the same curvature in all directions.

Supports vectorized operations: pass arrays of shape (N, 2+) for multiple points.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, NamedTuple, overload

import numpy as np

from gri_utils.constants import ECCENTRICITY_SQUARED, ER_M
from gri_utils.conversion.xyz_lla import lla_to_xyz

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


class OsculatingSphereScalar(NamedTuple):
    """Result of osculating sphere calculation for a single point.

    Attributes:
        center: 3D coordinates of sphere center in meters, shape (3,).
        radius: Radius of osculating sphere in meters.
    """

    center: NDArray[np.floating]
    radius: float


class OsculatingSphereVector(NamedTuple):
    """Result of osculating sphere calculation for multiple points.

    Attributes:
        center: 3D coordinates of sphere centers in meters, shape (N, 3).
        radius: Radii of osculating spheres in meters, shape (N,).
    """

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


@overload
def osculating_sphere(
    lla_ddm: Sequence[float],
) -> OsculatingSphereScalar: ...


@overload
def osculating_sphere(
    lla_ddm: NDArray[np.floating],
) -> OsculatingSphereScalar | OsculatingSphereVector: ...


def osculating_sphere(
    lla_ddm: NDArray[np.floating] | Sequence,
) -> OsculatingSphereScalar | OsculatingSphereVector:
    """Calculate the osculating sphere at a point on the WGS84 ellipsoid.

    The osculating sphere is the sphere that best approximates the ellipsoid at a given
    point, having the same curvature in all directions. Uses 0 for altitude, so lla_ddm
    can be passed in as two floats or include alt (lla_ddm[2] is ignored).

    Supports vectorized operations for multiple points.

    Args:
        lla_ddm: LLA coordinate(s). Shape (2+,) for single point or (N, 2+) for
            N points. Only lat/lon are used; altitude is ignored.

    Returns:
        OsculatingSphereScalar for single point (center shape (3,), radius float) or
        OsculatingSphereVector for multiple points (center (N, 3), radius (N,)).
    """
    arr = np.asarray(lla_ddm)

    lat_rad = np.radians(arr[..., 0])
    lon_rad = np.radians(arr[..., 1])
    sin_lat = np.sin(lat_rad)
    cos_lat = np.cos(lat_rad)
    sin_lon = np.sin(lon_rad)
    cos_lon = np.cos(lon_rad)

    # prime vertical radius of curvature
    n = ER_M / np.sqrt(1 - ECCENTRICITY_SQUARED * sin_lat**2)

    # meridional radius of curvature
    m = (
        ER_M
        * (1 - ECCENTRICITY_SQUARED)
        / (1 - ECCENTRICITY_SQUARED * sin_lat**2) ** (3 / 2)
    )

    # osculating sphere radius
    radius = np.sqrt(m * n)

    # Build LLA array with zero altitude for lla_to_xyz
    if arr.ndim == 1:
        lla_zero_alt = np.array([arr[0], arr[1], 0.0])
    else:
        lla_zero_alt = np.column_stack(
            [arr[..., 0], arr[..., 1], np.zeros(arr.shape[0])],
        )

    # surface point in ECEF coordinates
    surface_point = lla_to_xyz(lla_zero_alt)

    # normal vector - shape (3,) for scalar, (N, 3) for vectorized
    if arr.ndim == 1:
        normal = np.array([cos_lat * cos_lon, cos_lat * sin_lon, sin_lat])
    else:
        normal = np.column_stack([cos_lat * cos_lon, cos_lat * sin_lon, sin_lat])

    # center is surface point minus radius along normal
    if arr.ndim == 1:
        center = surface_point - radius * normal
        return OsculatingSphereScalar(center, float(radius))

    center = surface_point - radius[:, np.newaxis] * normal
    return OsculatingSphereVector(center, radius)
