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

from __future__ import annotations

from datetime import UTC, datetime
from typing import TYPE_CHECKING

import numpy as np
from gri_iono import iri_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)
# Using June 15 for minimal, December 15 for maximal (seasonal variations)
TIME_MIN = datetime(2024, 6, 15, 1, 0, 0, tzinfo=UTC)  # June, early morning
TIME_MAX = datetime(2024, 12, 15, 14, 0, 0, tzinfo=UTC)  # December, afternoon

# Solar flux (F10.7) parameters for minimal/maximal ionospheric effects
# SOLAR_FLUX_MIN = 70.0  # Low solar activity
# SOLAR_FLUX_MAX = 200.0  # High solar activity
SOLAR_FLUX_MIN = 200.0
SOLAR_FLUX_MAX = 200.0

# 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_iri_corrections(
    source_lla: NDArray,
    elevations: NDArray,
    frequency: float,
    *,
    use_max: bool,
) -> tuple[NDArray, NDArray]:
    """Compute IRI 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
    time = TIME_MAX if use_max else TIME_MIN
    solar_flux = SOLAR_FLUX_MAX if use_max else SOLAR_FLUX_MIN

    angular_corrections = []
    time_corrections = []
    for elev in tqdm(elevations, desc="IRI 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 IRI model
        # Returns (angular_correction, time_delay)
        angular_corr, time_delay = iri_ionospheric_correction(
            source_lla,
            collector_xyz,
            frequency,
            solar_flux,
            time,
        )
        angular_corrections.append(angular_corr)
        time_corrections.append(time_delay)

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