"""Simplified surface distance calculation.

Use the average latitude for flattening and pythagoras for combining north and east
distances. Useful for short distances and very fast.

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

from __future__ import annotations

import math
from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


# 1 nautical mile = 1852m = 1 arc-minute at the equator
# So 1 radian = (1852 * 60) meters * (180/pi) = degrees(1852 * 60) meters
RAD_2_M = math.degrees(1852 * 60)


def surface_dist_simple(
    lla1: NDArray[np.floating] | Sequence,
    lla2: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating] | float:
    """Get an approximate surface distance based on average latitudes.

    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.

    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)

    lat1_rad = np.radians(arr1[..., 0])
    lat2_rad = np.radians(arr2[..., 0])
    lon1_rad = np.radians(arr1[..., 1])
    lon2_rad = np.radians(arr2[..., 1])

    lat_med_rad = (lat1_rad + lat2_rad) / 2
    delta_lat_rad = lat2_rad - lat1_rad
    delta_lon = lon2_rad - lon1_rad

    lat_m = delta_lat_rad * RAD_2_M
    lon_m = np.cos(lat_med_rad) * delta_lon * RAD_2_M

    result = np.sqrt(lat_m**2 + lon_m**2)

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