Core

def distance(
    lon0: t.AngleRad,
    lat0: t.AngleRad,
    lon1: t.AngleRad,
    lat1: t.AngleRad,
    *,
    xp: ArrayApiNamespace | None,
) -> t.LengthM:
    """
    Returns the [Haversine great circle distance](https://en.wikipedia.org/wiki/Haversine_formula)
    between two coordinates.
    """
    d_lon = lon1 - lon0
    d_lat = lat1 - lat0

    xpr = math if xp is None else xp
    a = (xpr.sin(d_lat / 2)) ** 2 + (
        xpr.cos(lat0) * xpr.cos(lat1) * xpr.sin(d_lon / 2)
    ) ** 2
    c = 2 * xpr.asin(xpr.sqrt(a))

    return RADIUS_EARTH_MEAN * c

def bearing(
    lon0: t.AngleRad,
    lat0: t.AngleRad,
    lon1: t.AngleRad,
    lat1: t.AngleRad,
    *,
    xp: ArrayApiNamespace | None,
) -> t.AngleRad:
    """
    Returns the initial bearing (from origin to destination) along a
    [great-circle](https://en.wikipedia.org/wiki/Great_circle).

    :return: initial bearing, radians, [$-\\pi$, $\\pi$], clockwise from north
    """
    d_lon = lon1 - lon0

    xpr = math if xp is None else xp
    y = xpr.sin(d_lon) * xpr.cos(lat1)
    x = xpr.cos(lat0) * xpr.sin(lat1) - (
        xpr.sin(lat0) * xpr.cos(lat1) * xpr.cos(d_lon)
    )

    return xpr.atan2(y, x)

def lla_to_ecef(
    lon: t.AngleRad,
    lat: t.AngleRad,
    alt: t.GeometricAltitudeM,
    *,
    xp: ArrayApiNamespace | None,
) -> Point3D[t.LengthM]:
    """
    Converts geodetic coordinates to Earth-centered, Earth-fixed coordinates.
    Equivalent to `epsg:4979 +proj=cart +ellps=WGS84` in PROJ.

    :return: (x, y, z) coordinates
    """
    xpr = math if xp is None else xp
    v = RADIUS_EARTH_MEAN / xpr.sqrt(1 - E2 * xpr.sin(lat) * xpr.sin(lat))

    x = (v + alt) * xpr.cos(lat) * xpr.cos(lon)
    y = (v + alt) * xpr.cos(lat) * xpr.sin(lon)
    z = (v * (1 - E2) + alt) * xpr.sin(lat)

    return Point3D(x, y, z)

def ecef_to_enu(
    dx: t.DeltaLengthM,
    dy: t.DeltaLengthM,
    dz: t.DeltaLengthM,
    lon_ref: t.AngleRad,
    lat_ref: t.AngleRad,
    *,
    xp: ArrayApiNamespace | None,
) -> Point3D[t.LengthM]:
    """
    Converts Earth-centered, Earth-fixed coordinates
    (x, y, z with respect to a reference point)
    to East-North-Up coordinates.

    :return: (east, north, up) coordinates
    """
    xpr = math if xp is None else xp
    s_lat = xpr.sin(lat_ref)
    c_lat = xpr.cos(lat_ref)
    s_lon = xpr.sin(lon_ref)
    c_lon = xpr.cos(lon_ref)

    east = -s_lon * dx + c_lon * dy
    north = -s_lat * c_lon * dx - s_lat * s_lon * dy + c_lat * dz
    up = c_lat * c_lon * dx + c_lat * s_lon * dy + s_lat * dz

    return Point3D(east, north, up)

def impact_pressure_from_cas(
    cas: t.CasMPS,
    gamma: t.RatioOfSpecificHeats = GAMMA_DRY_AIR,
) -> t.ImpactPressurePA:
    """Impact pressure, compressible flow"""
    return impact_pressure(cas, RHO_0, P_0, gamma)

def impact_pressure_from_cas_behind_normal_shock(
    cas: t.CasMPS,
    gamma: t.RatioOfSpecificHeats = GAMMA_DRY_AIR,
) -> t.ImpactPressurePA:
    """Impact pressure, behind normal shock wave, supersonic flow"""
    return impact_pressure_behind_normal_shock(cas, A_0, P_0, gamma)

def density_factor(
    rho: t.DensityKGM3,
) -> t.DensityFactor:
    return (rho / RHO_0) ** 0.5

def eas_from_tas(
    tas: t.TasMPS,
    rho: t.DensityKGM3,
) -> t.EasMPS:
    """Converts TAS to EAS"""
    return tas * density_factor(rho)

def tas_from_eas(
    eas: t.EasMPS,
    rho: t.DensityKGM3,
) -> t.TasMPS:
    """Converts EAS to TAS"""
    return eas / density_factor(rho)

def compressibility_factor(
    qc: t.ImpactPressurePA,
    p: t.StaticPressurePA,
    gamma: t.RatioOfSpecificHeats = GAMMA_DRY_AIR,
) -> t.CompressibilityFactor:
    """Assumption: subsonic speeds"""
    exponent = (gamma - 1) / gamma
    inner = (qc / p + 1) ** exponent - 1
    return (exponent * p / qc * inner) ** 0.5

def eas_from_cas(
    cas: t.CasMPS,
    p: t.StaticPressurePA,
    gamma: t.RatioOfSpecificHeats = GAMMA_DRY_AIR,
) -> t.EasMPS:
    """Assumption: subsonic speeds"""
    qc = impact_pressure_from_cas(cas)
    f = compressibility_factor(qc, p, gamma)
    f0 = compressibility_factor(qc, P_0, gamma)
    return cas * f / f0

def tas_from_cas(
    cas: t.CasMPS,
    rho: t.DensityKGM3,
    p: t.StaticPressurePA,
) -> t.TasMPS:
    """Assumption: subsonic speeds"""
    eas = eas_from_cas(cas, p)
    return tas_from_eas(eas, rho)

def cas_from_tas(
    tas: t.TasMPS,
    rho: t.DensityKGM3,
    p: t.StaticPressurePA,
    gamma: t.RatioOfSpecificHeats = GAMMA_DRY_AIR,
) -> t.CasMPS:
    """Assumption: subsonic speeds"""
    eas = eas_from_tas(tas, rho)
    qc = impact_pressure(tas, rho, p, gamma)
    f = compressibility_factor(qc, p, gamma)
    f0 = compressibility_factor(qc, P_0, gamma)
    return eas * f0 / f

def temperature_below_tropopause(
    altitude: t.GeopotentialAltitudeM,
    delta_temperature: t.DeltaTemperatureK = 0.0,
) -> t.StaticTemperatureKBelowTropo:
    """
    Temperature below the tropopause, lapsing linearly with altitude (3.1-13)
    """
    return T_0 + delta_temperature + altitude * BETA_BELOW_TROP

def temperature_above_tropopause(
    delta_temperature: t.DeltaTemperatureK,
) -> t.StaticPressurePABelowTropo:
    """Temperature above the tropopause, isothermal (3.1-15, 3.1-16)."""
    return T_11 + delta_temperature

def pressure_below_tropopause(
    temperature_below_trop: t.StaticTemperatureK,
    delta_temperature: t.DeltaTemperatureK = 0.0,
) -> t.StaticPressurePA:
    """Pressure below tropopause (3.1-18)"""
    return P_0 * (
        ((temperature_below_trop - delta_temperature) / T_0)
        ** (-G_0 / (R_SPECIFIC_DRY_AIR * BETA_BELOW_TROP))
    )

def pressure_above_tropopause(
    altitude: t.GeopotentialAltitudeM, *, xp: ArrayApiNamespace | None
) -> t.StaticPressurePA:
    """Pressure above tropopause (3.1-20)"""
    xpr = math if xp is None else xp
    return P_11 * (
        xpr.exp(
            (altitude - H_BELOW_TROP) * (-G_0 / (R_SPECIFIC_DRY_AIR * T_11))
        )
    )

def atmosphere(
    altitude: t.GeopotentialAltitudeM,
    delta_temperature: t.DeltaTemperatureK,
    *,
    xp: ArrayApiNamespace | None,
) -> GasState:
    """
    BADA3 atmospheric model.

    Non-standard atmosphere conditions are modelled with the
    ISA temperature deviation.

    Note that this is only valid for 0-20km.
    """
    tropo_mask = altitude <= H_BELOW_TROP
    temperature = where(
        tropo_mask,
        temperature_below_tropopause(altitude, delta_temperature),
        lambda: temperature_above_tropopause(delta_temperature),
        xp=xp,
    )
    return GasState(
        temperature=temperature,
        pressure=(
            where(
                tropo_mask,
                pressure_below_tropopause(temperature, delta_temperature),
                lambda: pressure_above_tropopause(altitude, xp=xp),
                xp=xp,
            )
        ),
    )

def mach_number(
    tas: t.TasMPS,
    temperature: t.StaticTemperatureK,
) -> t.MachNumber:
    """Mach/TAS conversion (3.1-26)"""
    a = speed_of_sound(temperature, KAPPA, R_SPECIFIC_DRY_AIR)
    return tas / a

def specific_gas_constant(
    molar_mass: t.MolarMassKGMol,
) -> t.SpecificGasConstantJKGK:
    return R / molar_mass

def density(
    temperature: t.StaticTemperatureK,
    pressure: t.StaticPressurePA,
    specific_gas_constant: t.SpecificGasConstantJKGK,
) -> t.DensityKGM3:
    """Density, perfect gas"""
    return pressure / (specific_gas_constant * temperature)

def speed_of_sound(
    temperature: t.StaticTemperatureK,
    adiabatic_index: t.RatioOfSpecificHeats,
    specific_gas_constant: t.SpecificGasConstantJKGK,
) -> t.SpeedOfSoundMPS:
    """Speed of sound, perfect gas"""
    return (adiabatic_index * specific_gas_constant * temperature) ** 0.5

def dynamic_pressure(
    rho: t.DensityKGM3,
    tas: t.TasMPS,
) -> t.DynamicPressurePA:
    """Dynamic pressure, incompressible flow"""
    return 0.5 * rho * tas**2

def total_pressure(
    tas: t.TasMPS,
    rho: t.DensityKGM3,
    p: t.StaticPressurePA,
    gamma: t.RatioOfSpecificHeats = GAMMA_DRY_AIR,
) -> t.TotalPressurePA:
    """Total pressure, compressible flow"""
    # NOTE: from bernoulli's formula
    inner = 1 + (gamma - 1) / (2 * gamma) * rho / p * tas**2
    return p * inner ** (gamma / (gamma - 1))

def total_pressure_behind_normal_shock(
    tas: t.TasMPS,
    rho: t.DensityKGM3,
    p: t.StaticPressurePA,
    gamma: t.RatioOfSpecificHeats = GAMMA_DRY_AIR,
) -> t.TotalPressurePA:
    """Total pressure, behind normal shock wave, supersonic flow"""
    common = rho / p * tas**2
    inner = ((gamma + 1) ** 2 / gamma * common) / (4 * common - 2 * (gamma - 1))
    return (1 + gamma) / (2 * gamma) * rho * tas**2 * inner ** (1 / (gamma - 1))

def impact_pressure(
    tas: t.TasMPS,
    rho: t.DensityKGM3,
    p: t.StaticPressurePA,
    gamma: t.RatioOfSpecificHeats = GAMMA_DRY_AIR,
) -> t.ImpactPressurePA:
    """Impact pressure, compressible flow"""
    return total_pressure(tas, rho, p, gamma) - p

def impact_pressure_behind_normal_shock(
    tas: t.TasMPS,
    a: t.SpeedOfSoundMPS,
    p: t.StaticPressurePA,
    gamma: t.RatioOfSpecificHeats = GAMMA_DRY_AIR,
) -> t.TotalPressurePA:
    """Impact pressure, behind normal shock wave, supersonic flow"""
    inner = (gamma + 1) ** 2 / (4 * gamma - 2 * (gamma - 1) * (a / tas) ** 2)
    return (
        (1 + gamma) / 2 * (tas / a) ** 2 * p * (inner ** (1 / (gamma - 1)) - 1)
    )