brahe 1.6.2

"""Shared helpers for the GMAT benchmark implementations.

Provides:
  - _ensure_gmat: precondition-validating gmatpy initializer (idempotent).
  - gmat_clear, time_iterations, mj2000_to_gcrf, km_to_m_state, m_to_km_state,
    mu_si_to_gmat, build_task_result, ensure_eop (added in later tasks).
"""

import os
import sys
import time
from pathlib import Path

import numpy as np

import brahe

from benchmarks.comparative.results import TaskResult


_GMAT_INITIALIZED = False


def _ensure_gmat() -> None:
    """Validate preconditions and initialize gmatpy (idempotent).

    Raises ImportError with a precise diagnostic on any failure. The runner's
    `_run_gmat` catches this and emits the yellow "GMAT not ready" skip message.

    Order of checks:
      1. GMAT_ROOT_PATH is set.
      2. The _pyXY/_gmat_py.so exists for the running Python.
      3. api_startup_file.txt exists (generated by `just bench-compare-setup`).
      4. sys.path includes <GMAT_ROOT_PATH>/bin.
      5. import gmatpy succeeds.
      6. gmat.Setup(<absolute startup path>) succeeds.
    """
    global _GMAT_INITIALIZED
    if _GMAT_INITIALIZED:
        return

    gmat_root = os.environ.get("GMAT_ROOT_PATH")
    if not gmat_root:
        raise ImportError("GMAT_ROOT_PATH not set")

    root = Path(gmat_root)
    py_tag = f"_py{sys.version_info.major}{sys.version_info.minor}"
    so_path = root / "bin" / "gmatpy" / py_tag / "_gmat_py.so"
    if not so_path.exists():
        raise ImportError(
            f"gmatpy missing binary for Python {py_tag} at {so_path}"
        )

    startup = root / "bin" / "api_startup_file.txt"
    if not startup.exists():
        raise ImportError(
            f"api_startup_file.txt missing at {startup} — "
            f"run: just bench-compare-setup"
        )

    bin_dir = str(root / "bin")
    if bin_dir not in sys.path:
        sys.path.insert(1, bin_dir)

    import gmatpy as gmat  # noqa: F401  (side-effect: triggers FileManager init)

    gmat.Setup(str(startup))
    _GMAT_INITIALIZED = True


def time_iterations(func, iterations: int):
    """Time a function over multiple iterations.

    Returns (times_seconds, results_from_first_iteration). Identical signature
    to the brahe-py and Basilisk implementations.
    """
    times: list[float] = []
    first_results = None
    for i in range(iterations):
        start = time.perf_counter()
        results = func()
        elapsed = time.perf_counter() - start
        times.append(elapsed)
        if i == 0:
            first_results = results
    return times, first_results


def gmat_clear() -> None:
    """Reset the GMAT Moderator singleton's named-object registry.

    Required at the start of every iteration of any task that constructs
    Spacecraft, Propagator, ForceModel, or other named objects. Without this,
    the second iteration's Construct() either errors on name collision or
    returns the previous iteration's object.
    """
    import gmatpy as gmat
    gmat.Clear()


def km_to_m_state(state_km) -> list[float]:
    """Scale a 6-vector [r (km), v (km/s)] -> [r (m), v (m/s)]."""
    return [float(x) * 1000.0 for x in state_km]


def m_to_km_state(state_m) -> list[float]:
    """Scale a 6-vector [r (m), v (m/s)] -> [r (km), v (km/s)]."""
    return [float(x) / 1000.0 for x in state_m]


def mu_si_to_gmat(mu_si: float) -> float:
    """Scale a gravitational parameter from m^3/s^2 to km^3/s^2 (GMAT default)."""
    return mu_si * 1.0e-9


def _find_orekit_eop_file() -> str | None:
    orekit_data = os.environ.get(
        "OREKIT_DATA", str(Path.home() / ".orekit" / "orekit-data")
    )
    eop_path = (
        Path(orekit_data)
        / "Earth-Orientation-Parameters"
        / "IAU-2000"
        / "finals2000A.all"
    )
    if eop_path.exists():
        return str(eop_path)
    return None


def ensure_eop() -> None:
    """Initialize brahe EOP using OreKit's IERS data if available, else fallback.

    Identical behavior to implementations/python/base.py:ensure_eop and
    implementations/basilisk/base.py:ensure_eop so the J2000->GCRF transform
    applied to GMAT output uses the same IERS data the other baselines use.
    """
    if not brahe.get_global_eop_initialization():
        eop_path = _find_orekit_eop_file()
        if eop_path:
            provider = brahe.FileEOPProvider.from_file(eop_path, True, "Hold")
            brahe.set_global_eop_provider(provider)
        else:
            brahe.initialize_eop()


def mj2000_to_gcrf(r_m, v_mps) -> list[float]:
    """Transform a 6-vector state from EarthMJ2000Eq (= EME2000 = J2000-equatorial)
    inertial to GCRF.

    GMAT's `EarthMJ2000Eq` inertial frame is equivalent to Orekit's EME2000 and
    Basilisk's J2000-equatorial. brahe's state_eme2000_to_gcrf applies the
    IAU 2006 frame bias (time-independent, sub-meter at LEO). EOP is not strictly
    required but we ensure it for consistency with the other baselines.

    Inputs are in METERS / METERS-PER-SECOND. Callers must scale GMAT output
    (km, km/s) using km_to_m_state() before calling this.
    """
    ensure_eop()
    state = np.array([*r_m, *v_mps], dtype=float)
    gcrf = brahe.state_eme2000_to_gcrf(state)
    return list(map(float, gcrf))


def build_task_result(
    task_name: str,
    iterations: int,
    times_seconds: list,
    results: list,
    extra_metadata: dict | None = None,
) -> TaskResult:
    """Build a TaskResult with standard GMAT metadata."""
    metadata = {
        "library": "gmat",
        "language": "gmat",
        "version": _detect_gmat_version(),
    }
    if extra_metadata:
        metadata.update(extra_metadata)
    return TaskResult(
        task_name=task_name,
        language="gmat",
        library="gmat",
        iterations=iterations,
        times_seconds=times_seconds,
        results=results,
        metadata=metadata,
    )


def _detect_gmat_version() -> str:
    """Best-effort GMAT version string from GMAT_ROOT_PATH directory name."""
    root = os.environ.get("GMAT_ROOT_PATH", "")
    name = Path(root).name  # e.g. "GMAT R2026a"
    return name or "unknown"


def quat_brahe_to_gmat(q_brahe) -> list[float]:
    """Reorder a quaternion from brahe scalar-first [w, x, y, z]
    to GMAT scalar-last [q1, q2, q3, q4]."""
    return [float(q_brahe[1]), float(q_brahe[2]), float(q_brahe[3]), float(q_brahe[0])]


def quat_gmat_to_brahe(q_gmat) -> list[float]:
    """Reorder a quaternion from GMAT scalar-last [q1, q2, q3, q4]
    to brahe scalar-first [w, x, y, z]."""
    return [float(q_gmat[3]), float(q_gmat[0]), float(q_gmat[1]), float(q_gmat[2])]