import math

import numpy as np
import polars as pl

_EARTH_RADIUS_KM = 6371.0


def haversine_km(
    lat1: np.ndarray, lon1: np.ndarray, lat2: float, lon2: float
) -> np.ndarray:
    """Compute haversine distance in km between arrays (lat1, lon1) and a single point (lat2, lon2)."""
    lat1_rad = np.radians(lat1)
    lon1_rad = np.radians(lon1)
    lat2_rad = np.radians(lat2)
    lon2_rad = np.radians(lon2)
    dlat = lat2_rad - lat1_rad
    dlon = lon2_rad - lon1_rad
    a = (
        np.sin(dlat / 2) ** 2
        + np.cos(lat1_rad) * np.cos(lat2_rad) * np.sin(dlon / 2) ** 2
    )
    c = 2 * np.arcsin(np.sqrt(a))
    return _EARTH_RADIUS_KM * c


def haversine_km_expr(
    lat_col: str, lon_col: str, dest_lat: float, dest_lon: float
) -> pl.Expr:
    """Polars expression computing haversine distance in km to a fixed point."""
    dest_lat_rad = math.radians(dest_lat)
    dest_lon_rad = math.radians(dest_lon)

    lat_rad = pl.col(lat_col).radians()
    lon_rad = pl.col(lon_col).radians()

    dlat = pl.lit(dest_lat_rad) - lat_rad
    dlon = pl.lit(dest_lon_rad) - lon_rad

    a = (dlat / 2).sin() ** 2 + pl.lit(dest_lat_rad).cos() * lat_rad.cos() * (
        dlon / 2
    ).sin() ** 2
    return 2 * _EARTH_RADIUS_KM * a.sqrt().arcsin()