"""Haversine surface distance.

https://en.wikipedia.org/wiki/Haversine_formula

Haversine is a relatively simple calculation but has known accuracy errors of up to
0.5% across distances.

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

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from gri_utils import constants

from .central_angle import central_angle

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


MEAN_EARTH_RADIUS = (2 * constants.ER_M + constants.PR_M) / 3
"""Arithmetic mean radius"""


def surface_dist_haversine(
    lla1: NDArray[np.floating] | Sequence,
    lla2: NDArray[np.floating] | Sequence,
    *,
    use_geocentric_radius: bool = False,
) -> NDArray[np.floating] | float:
    """Haversine distance on the Earth.

    Returns the great circle distance in meters assuming a spherical Earth.

    Supports vectorized operations for pairwise distances between corresponding points.

    Args:
        lla1: LLA coordinate(s). Shape (2+,) for single point or (N, 2+) for N points.
            Only lat/lon (first two elements) are used.
        lla2: LLA coordinate(s). Must match shape of lla1.
        use_geocentric_radius: If True, calculate geocentric radius at the midpoint
            latitude. If False (default), use mean Earth radius.

    Returns:
        Distance(s) in meters. Returns a float for single point pairs,
        or ndarray of shape (N,) for vectorized input.
    """
    arr1 = np.asarray(lla1)
    arr2 = np.asarray(lla2)

    c_rad = central_angle(arr1, arr2)

    if not use_geocentric_radius:
        r = MEAN_EARTH_RADIUS
    else:
        # https://en.wikipedia.org/wiki/Earth_radius#Geocentric_radius
        lat_mid = np.radians(arr1[..., 0] + arr2[..., 0]) / 2
        a_cos_phi = constants.ER_M * np.cos(lat_mid)
        b_sin_phi = constants.PR_M * np.sin(lat_mid)
        num = (a_cos_phi * constants.ER_M) ** 2 + (b_sin_phi * constants.PR_M) ** 2
        den = a_cos_phi**2 + b_sin_phi**2
        r = np.sqrt(num / den)

    result = r * c_rad

    if np.ndim(result) == 0:
        return float(result)
    return result
