"""NeQuickG ionospheric model calculations for min/max corrections."""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
from gri_iono import nequickg_ionospheric_correction
from gri_utils import conversion
from tqdm import tqdm

if TYPE_CHECKING:
    from numpy.typing import NDArray


# Time parameters for minimal/maximal ionospheric effects (UTC decimal hours)
HOUR_MIN = 1.0  # Early morning (minimal ionosphere)
HOUR_MAX = 14.0  # Afternoon (maximal ionosphere)

# Broadcast ionospheric coefficients (a0, a1, a2) for min/max conditions
# These are representative values for low and high ionospheric activity
# COEFFS_MIN = (100.0, 0.0, 0.0)  # Low ionospheric activity
# COEFFS_MAX = (200.0, 0.1, 0.01)  # High ionospheric activity
COEFFS_MIN = (200.0, 0.1, 0.01)
COEFFS_MAX = (200.0, 0.1, 0.01)

# Month parameters for seasonal variations
MONTH_MIN = 6  # June (northern summer, minimal ionosphere)
MONTH_MAX = 12  # December (northern winter, maximal ionosphere)

# Earth radius and LEO satellite altitude
_EARTH_RADIUS_M = 6378137.0
_LEO_ALTITUDE_M = 1_000_000.0  # 1000 km LEO orbit


def _slant_range_for_altitude(elevation_deg: float, altitude_m: float) -> float:
    """Calculate slant range to achieve a specific satellite altitude.

    Args:
        elevation_deg: Elevation angle in degrees.
        altitude_m: Desired satellite altitude above Earth surface in meters.

    Returns:
        Slant range in meters.
    """
    elev_rad = np.radians(elevation_deg)
    sat_radius = _EARTH_RADIUS_M + altitude_m

    # Quadratic formula: slant_range^2 + 2*R*sin(elev)*slant_range + (R^2 - r_sat^2) = 0
    b = 2 * _EARTH_RADIUS_M * np.sin(elev_rad)
    c = _EARTH_RADIUS_M**2 - sat_radius**2
    discriminant = b**2 - 4 * c

    return (-b + np.sqrt(discriminant)) / 2


def compute_nequickg_corrections(
    source_lla: NDArray,
    elevations: NDArray,
    frequency: float,
    *,
    use_max: bool,
) -> tuple[NDArray, NDArray]:
    """Compute NeQuickG ionospheric corrections for a range of elevation angles.

    Args:
        source_lla: Source position as [lat, lon, alt] in degrees and meters.
        elevations: Array of elevation angles in degrees (0 to 90).
        frequency: Signal frequency in Hz.
        use_max: If True, use maximal ionosphere parameters; if False, use minimal.

    Returns:
        Tuple of (angular_corrections, time_corrections) where:
            - angular_corrections: Array of angular corrections in radians.
            - time_corrections: Array of ionospheric time delays in seconds.
    """
    receiver_xyz = conversion.lla_to_xyz(source_lla)

    # Select parameters based on min/max ionospheric conditions
    utc_hours = HOUR_MAX if use_max else HOUR_MIN
    coefficients = COEFFS_MAX if use_max else COEFFS_MIN
    month = MONTH_MAX if use_max else MONTH_MIN

    angular_corrections = []
    time_corrections = []
    for elev in tqdm(elevations, desc="NeQuickG elevations"):
        # Compute satellite position at fixed LEO altitude
        slant_range = _slant_range_for_altitude(elev, _LEO_ALTITUDE_M)
        collector_xyz = conversion.translate_aer(
            receiver_xyz,
            (0, elev, slant_range),
        )

        # Compute ionospheric correction using NeQuickG model
        # Returns (angular_correction, time_delay)
        angular_corr, time_delay = nequickg_ionospheric_correction(
            source_lla,
            collector_xyz,
            frequency,
            coefficients,
            utc_hours,
            month,
        )
        angular_corrections.append(angular_corr)
        time_corrections.append(time_delay)

    return np.array(angular_corrections), np.array(time_corrections)
