"""Generate ionospheric correction min/max plots for ionospheric models.

This script computes and visualizes the range of ionospheric correction effects
by showing minimal and maximal angular and time corrections as a function of
elevation angle from horizon (0°) to zenith (90°).

Supports multiple ionospheric models:
- Bent: Traditional Bent ionospheric model
- BentAI: AI-created Bent model (experimental)
- NeQuick2: NeQuick2 ionospheric model using F10.7 solar flux
- NeQuickG: NeQuickG ionospheric model using Galileo broadcast coefficients
- IRI: International Reference Ionosphere model

Run with:
python -m gri_sandbox.iono_minmax.main [--MODEL] [freq]

where MODEL is one of: bent, bentai, nequick2, nequickg, iri
"""

from __future__ import annotations

# ruff: noqa: T201
import argparse
from time import time_ns

import numpy as np

from gri_sandbox.iono_minmax.bent import compute_bent_corrections
from gri_sandbox.iono_minmax.bent_ai import compute_bent_ai_corrections
from gri_sandbox.iono_minmax.iri import compute_iri_corrections
from gri_sandbox.iono_minmax.nequick2 import compute_nequick2_corrections
from gri_sandbox.iono_minmax.nequickg import compute_nequickg_corrections
from gri_sandbox.iono_minmax.plots import plot_iono_corrections


def _get_args() -> argparse.Namespace:
    """Parse command line arguments.

    Returns:
        Parsed arguments with frequency and model selection.
    """
    parser = argparse.ArgumentParser(
        description="Generate ionospheric correction min/max plots",
    )
    parser.add_argument(
        "frequency",
        nargs="?",
        type=float,
        default=500,
        help="Signal frequency in MHz. Default: 500 MHz",
    )

    # Model selection (mutually exclusive)
    model_group = parser.add_mutually_exclusive_group()
    model_group.add_argument(
        "--bent",
        action="store_true",
        help="Use the Bent ionospheric model (default)",
    )
    model_group.add_argument(
        "--bentai",
        action="store_true",
        help="Use the AI-generated / simplified Bent ionospheric model",
    )
    model_group.add_argument(
        "--nequick2",
        action="store_true",
        help="Use the NeQuick2 ionospheric model",
    )
    model_group.add_argument(
        "--nequickg",
        action="store_true",
        help="Use the NeQuickG ionospheric model (Galileo)",
    )
    model_group.add_argument(
        "--iri",
        action="store_true",
        help="Use the IRI (International Reference Ionosphere) model",
    )

    return parser.parse_args()


def main() -> None:
    """Generate ionospheric correction plots for the selected model."""
    args = _get_args()

    # Convert frequency from MHz to Hz
    frequency_mhz = args.frequency
    frequency = frequency_mhz * 1e6

    # Receiver at equator, prime meridian, sea level
    source_lla = np.array([0.0, 0.0, 0.0])

    # Elevation angles from horizon to zenith (exclusive to avoid singularities)
    step = 0.1
    elevations = np.arange(step, 90, step)  # (0, 90) exclusive

    # Select model and compute corrections
    if args.bentai:
        model_name = "BentAI"
        compute_func = compute_bent_ai_corrections
    elif args.nequick2:
        model_name = "NeQuick2"
        compute_func = compute_nequick2_corrections
    elif args.nequickg:
        model_name = "NeQuickG"
        compute_func = compute_nequickg_corrections
    elif args.iri:
        model_name = "IRI"
        compute_func = compute_iri_corrections
    else:
        # Default to traditional Bent model
        model_name = "Bent"
        compute_func = compute_bent_corrections

    print(f"Running {model_name} with {frequency_mhz:,} MHz")
    print(f"Computing {model_name} model corrections...")
    start = time_ns()

    angle_min, time_min = compute_func(
        source_lla,
        elevations,
        frequency,
        use_max=False,
    )
    angle_max, time_max = compute_func(
        source_lla,
        elevations,
        frequency,
        use_max=True,
    )

    duration_s = (time_ns() - start) / 1e9
    print(f"Duration: {duration_s:,.3f}s")

    # Plot combined corrections
    plot_iono_corrections(
        elevations,
        angle_min,
        angle_max,
        time_min,
        time_max,
        model_name,
        frequency,
        duration_s,
    )

    print("All plots generated successfully!")


if __name__ == "__main__":
    main()
