brahe 1.6.0

"""
Python (Brahe) propagation benchmarks.
"""

import numpy as np

import brahe

from benchmarks.comparative.implementations.python.base import (
    ensure_eop,
    ensure_sw,
    time_iterations,
)
from benchmarks.comparative.results import TaskResult


def keplerian_single(params: dict, iterations: int) -> TaskResult:
    """Benchmark Keplerian propagation of multiple orbits to single future epochs."""
    ensure_eop()
    cases = params["cases"]

    def run():
        results = []
        for case in cases:
            epc = brahe.Epoch.from_jd(case["jd"], brahe.TimeSystem.UTC)
            oe = np.array(case["elements"])
            dt = case["dt"]
            target = epc + dt

            prop = brahe.KeplerianPropagator.from_keplerian(
                epc, oe, brahe.AngleFormat.DEGREES, 60.0
            )
            state = prop.state_eci(target)
            results.append(state.tolist())
        return results

    times, results = time_iterations(run, iterations)

    return TaskResult(
        task_name="propagation.keplerian_single",
        language="python",
        library="brahe",
        iterations=iterations,
        times_seconds=times,
        results=results,
        metadata={
            "library": "brahe",
            "language": "python",
            "version": getattr(brahe, "__version__", "unknown"),
        },
    )


def keplerian_trajectory(params: dict, iterations: int) -> TaskResult:
    """Benchmark Keplerian propagation over trajectory steps.

    Two input shapes:

    - Single-IC (perf): ``{jd, elements, step_size, n_steps}`` — propagate
      one orbit and return every step's state, like before.
    - Multi-IC (accuracy): ``{cases: [{jd, elements}], step_size, n_steps}``
      — propagate each case over the shared horizon and return the
      final state per case, so accuracy stats reflect IC variance.
    """
    ensure_eop()
    step_size = params["step_size"]
    n_steps = params["n_steps"]
    cases = params.get("cases")

    if cases is None:
        # Perf path: single IC, full trajectory output.
        jd = params["jd"]
        oe = np.array(params["elements"])
        epc = brahe.Epoch.from_jd(jd, brahe.TimeSystem.UTC)

        def run():
            prop = brahe.KeplerianPropagator.from_keplerian(
                epc, oe, brahe.AngleFormat.DEGREES, step_size
            )
            results = []
            for step_idx in range(n_steps):
                target = epc + (step_idx + 1) * step_size
                state = prop.state_eci(target)
                results.append(state.tolist())
            return results
    else:
        # Accuracy path: IC sweep, return one final state per case.
        def run():
            results = []
            for case in cases:
                epc = brahe.Epoch.from_jd(case["jd"], brahe.TimeSystem.UTC)
                oe = np.array(case["elements"])
                prop = brahe.KeplerianPropagator.from_keplerian(
                    epc, oe, brahe.AngleFormat.DEGREES, step_size
                )
                target = epc + n_steps * step_size
                state = prop.state_eci(target)
                results.append(state.tolist())
            return results

    times, results = time_iterations(run, iterations)

    return TaskResult(
        task_name="propagation.keplerian_trajectory",
        language="python",
        library="brahe",
        iterations=iterations,
        times_seconds=times,
        results=results,
        metadata={
            "library": "brahe",
            "language": "python",
            "version": getattr(brahe, "__version__", "unknown"),
        },
    )


def sgp4_single(params: dict, iterations: int) -> TaskResult:
    """Benchmark SGP4 propagation to 50 future epochs."""
    ensure_eop()
    line1 = params["line1"]
    line2 = params["line2"]
    offsets = params["time_offsets_seconds"]

    prop = brahe.SGPPropagator.from_tle(line1, line2, 60.0)
    base_epoch = prop.epoch

    def run():
        results = []
        for dt in offsets:
            target = base_epoch + dt
            # Use .state() to get TEME output directly (matches Java's TEME frame)
            state = prop.state(target)
            results.append(state.tolist())
        return results

    times, results = time_iterations(run, iterations)

    return TaskResult(
        task_name="propagation.sgp4_single",
        language="python",
        library="brahe",
        iterations=iterations,
        times_seconds=times,
        results=results,
        metadata={
            "library": "brahe",
            "language": "python",
            "version": getattr(brahe, "__version__", "unknown"),
        },
    )


def sgp4_trajectory(params: dict, iterations: int) -> TaskResult:
    """Benchmark SGP4 propagation over 1 day at 60s steps."""
    ensure_eop()
    line1 = params["line1"]
    line2 = params["line2"]
    step_size = params["step_size"]
    n_steps = params["n_steps"]

    prop = brahe.SGPPropagator.from_tle(line1, line2, step_size)
    base_epoch = prop.epoch

    def run():
        results = []
        for step_idx in range(n_steps):
            target = base_epoch + (step_idx + 1) * step_size
            # Use .state() to get TEME output directly (matches Java's TEME frame)
            state = prop.state(target)
            results.append(state.tolist())
        return results

    times, results = time_iterations(run, iterations)

    return TaskResult(
        task_name="propagation.sgp4_trajectory",
        language="python",
        library="brahe",
        iterations=iterations,
        times_seconds=times,
        results=results,
        metadata={
            "library": "brahe",
            "language": "python",
            "version": getattr(brahe, "__version__", "unknown"),
        },
    )


def numerical_twobody(params: dict, iterations: int) -> TaskResult:
    """Benchmark numerical two-body propagation.

    See keplerian_trajectory for the single-IC vs multi-IC ``cases``
    parameter shape contract. Single-IC returns every trajectory step;
    multi-IC returns one final state per IC.
    """
    ensure_eop()
    step_size = params["step_size"]
    n_steps = params["n_steps"]
    cases = params.get("cases")

    # Fixed-step RK4 to match the Orekit adapter; both sides pinned to the
    # same integrator + step so the residual is a pure-numerics comparison.
    prop_config = brahe.NumericalPropagationConfig(
        brahe.IntegrationMethod.RK4,
        brahe.IntegratorConfig.fixed_step(step_size),
        brahe.VariationalConfig(),
    )
    force_config = brahe.ForceModelConfig.two_body()

    if cases is None:
        # Perf path: single IC, full trajectory.
        jd = params["jd"]
        oe = np.array(params["elements"])
        epc = brahe.Epoch.from_jd(jd, brahe.TimeSystem.UTC)
        cart = brahe.state_koe_to_eci(oe, brahe.AngleFormat.DEGREES)

        def run():
            prop = brahe.NumericalOrbitPropagator(
                epc, cart, prop_config, force_config
            )
            prop.set_trajectory_mode(brahe.TrajectoryMode.DISABLED)
            results = []
            for _ in range(n_steps):
                prop.step_by(step_size)
                results.append(prop.current_state().tolist())
            return results
    else:
        # Accuracy path: IC sweep, final state per case.
        def run():
            results = []
            for case in cases:
                epc = brahe.Epoch.from_jd(case["jd"], brahe.TimeSystem.UTC)
                oe = np.array(case["elements"])
                cart = brahe.state_koe_to_eci(oe, brahe.AngleFormat.DEGREES)
                prop = brahe.NumericalOrbitPropagator(
                    epc, cart, prop_config, force_config
                )
                prop.set_trajectory_mode(brahe.TrajectoryMode.DISABLED)
                for _ in range(n_steps):
                    prop.step_by(step_size)
                results.append(prop.current_state().tolist())
            return results

    times, results = time_iterations(run, iterations)

    return TaskResult(
        task_name="propagation.numerical_twobody",
        language="python",
        library="brahe",
        iterations=iterations,
        times_seconds=times,
        results=results,
        metadata={
            "library": "brahe",
            "language": "python",
            "version": getattr(brahe, "__version__", "unknown"),
        },
    )


def _force_model_from_params(params: dict) -> "brahe.ForceModelConfig":
    """Build a brahe ForceModelConfig matching the task parameter dict.

    The parameter vector layout is brahe's standard
    `[mass, drag_area, Cd, srp_area, Cr]`, so drag/SRP use parameter indices
    rather than fixed values — that way the same vector drives every step.
    """
    gravity = brahe.GravityConfiguration.spherical_harmonic(
        degree=params["gravity_degree"],
        order=params["gravity_order"],
    )

    third_body = None
    bodies = []
    if params.get("third_body_sun"):
        bodies.append(brahe.ThirdBody.SUN)
    if params.get("third_body_moon"):
        bodies.append(brahe.ThirdBody.MOON)
    if bodies:
        third_body = brahe.ThirdBodyConfiguration(
            ephemeris_source=brahe.EphemerisSource.DE440s,
            bodies=bodies,
        )

    drag = None
    if params.get("drag"):
        drag = brahe.DragConfiguration(
            model=brahe.AtmosphericModel.NRLMSISE00,
            area=brahe.ParameterSource.parameter_index(1),
            cd=brahe.ParameterSource.parameter_index(2),
        )

    srp = None
    if params.get("srp"):
        srp = brahe.SolarRadiationPressureConfiguration(
            area=brahe.ParameterSource.parameter_index(3),
            cr=brahe.ParameterSource.parameter_index(4),
            eclipse_model=brahe.EclipseModel.CONICAL,
        )

    needs_mass = drag is not None or srp is not None
    mass = brahe.ParameterSource.parameter_index(0) if needs_mass else None

    return brahe.ForceModelConfig(
        gravity=gravity,
        drag=drag,
        srp=srp,
        third_body=third_body,
        relativity=False,
        mass=mass,
    )


def _numerical_rk4_run(task_name: str, params: dict, iterations: int) -> TaskResult:
    """Shared driver for the RK4 force-model propagation tasks.

    Accepts the same single-IC/multi-IC parameter contract as the
    other propagation adapters (see ``keplerian_trajectory``).
    """
    ensure_eop()
    if params.get("drag"):
        ensure_sw()

    step_size = params["step_size"]
    n_steps = params["n_steps"]
    param_vec = np.array(params["params"])
    cases = params.get("cases")

    force_config = _force_model_from_params(params)
    prop_config = brahe.NumericalPropagationConfig(
        brahe.IntegrationMethod.RK4,
        brahe.IntegratorConfig.fixed_step(step_size),
        brahe.VariationalConfig(),
    )

    if cases is None:
        # Perf path: single IC, full trajectory output.
        jd = params["jd"]
        oe = np.array(params["elements_deg"])
        epc = brahe.Epoch.from_jd(jd, brahe.TimeSystem.UTC)
        cart = brahe.state_koe_to_eci(oe, brahe.AngleFormat.DEGREES)

        def run():
            prop = brahe.NumericalOrbitPropagator(
                epc, cart, prop_config, force_config, param_vec
            )
            prop.set_trajectory_mode(brahe.TrajectoryMode.DISABLED)
            results = []
            for _ in range(n_steps):
                prop.step_by(step_size)
                results.append(prop.current_state().tolist())
            return results
    else:
        # Accuracy path: IC sweep, final state per case.
        def run():
            results = []
            for case in cases:
                epc = brahe.Epoch.from_jd(case["jd"], brahe.TimeSystem.UTC)
                oe = np.array(case["elements"])
                cart = brahe.state_koe_to_eci(oe, brahe.AngleFormat.DEGREES)
                prop = brahe.NumericalOrbitPropagator(
                    epc, cart, prop_config, force_config, param_vec
                )
                prop.set_trajectory_mode(brahe.TrajectoryMode.DISABLED)
                for _ in range(n_steps):
                    prop.step_by(step_size)
                results.append(prop.current_state().tolist())
            return results

    times, results = time_iterations(run, iterations)

    return TaskResult(
        task_name=task_name,
        language="python",
        library="brahe",
        iterations=iterations,
        times_seconds=times,
        results=results,
        metadata={
            "library": "brahe",
            "language": "python",
            "version": getattr(brahe, "__version__", "unknown"),
            "frame": "GCRF",
            "integrator": "RK4",
            "step_size": step_size,
            "gravity": f"{params['gravity_degree']}x{params['gravity_order']}",
        },
    )


def numerical_rk4_grav5x5(params: dict, iterations: int) -> TaskResult:
    return _numerical_rk4_run("propagation.numerical_rk4_grav5x5", params, iterations)


def numerical_rk4_grav20x20_sun_moon(params: dict, iterations: int) -> TaskResult:
    return _numerical_rk4_run(
        "propagation.numerical_rk4_grav20x20_sun_moon", params, iterations
    )


def numerical_rk4_grav80x80_full(params: dict, iterations: int) -> TaskResult:
    return _numerical_rk4_run(
        "propagation.numerical_rk4_grav80x80_full", params, iterations
    )