yksa-orbital/yksa_orbital/backends/orekit.py

"""Orekit propagation backend -- HTTP client to the sidecar service.

Implements the same :class:`~odm.PropagatorBackend` seam as
:class:`~yksa_orbital.backends.sgp4.Sgp4Backend`, but delegates the math to the
Orekit sidecar (``services/orekit``) over HTTP. The stored OMM goes over the wire
as a CCSDS message (see :mod:`yksa_orbital.wire`) and the JSON response maps back
into a frame-agnostic :class:`~odm.StateVector`. Set
``settings.ORBITAL_PROPAGATOR_BACKEND = "orekit"`` (and ``OREKIT_SERVICE_URL``)
to make it the default -- nothing downstream changes.

Requests carry only what this deployment means to override. The model itself
lives in the sidecar's ``config.py`` and comes back through :meth:`model_config`;
restating it here would be a second definition of the same thing.
"""

from __future__ import annotations

from datetime import datetime, timezone

import httpx
from django.conf import settings

from odm import (
    BackendBusy,
    PropagationError,
    PropagatorBackend,
    StateVector,
    parse_omm_epoch as _parse_iso,
)
from ..wire import omm_message, omm_messages

#: Frames every state_at/ephemeris call needs: TEME is the seam's source of
#: truth, ITRF supplies the Earth-fixed (ECEF) view.
_CORE_FRAMES = ("TEME", "ITRF")

#: Frames offered by the "show in frame" dropdown when this backend is active.
DROPDOWN_FRAMES = ("TEME", "GCRF", "EME2000", "TOD", "ITRF")

#: Route -> setting holding its timeout. A decay run is a minutes-long numerical
#: propagation and an ensemble repeats it per realization, so neither can share
#: the seconds-long budget the state/ephemeris queries use.
_TIMEOUTS = {
    "/decay": ("OREKIT_DECAY_TIMEOUT_S", 900),
    "/decay_ensemble": ("OREKIT_ENSEMBLE_TIMEOUT_S", 7200),
    # A download plus a file reload -- longer than the seconds-long state budget,
    # shorter than a propagation.
    "/refresh_space_weather": ("OREKIT_REFRESH_TIMEOUT_S", 120),
    # Read on a page render for a caption. Fail fast and let the caller fall
    # back rather than holding the response open.
    "/config": ("OREKIT_CONFIG_TIMEOUT_S", 2),
}


def _service_url() -> str:
    url = getattr(settings, "OREKIT_SERVICE_URL", "")
    if not url:
        raise PropagationError("OREKIT_SERVICE_URL is not configured")
    return url.rstrip("/")


def _timeout(path: str) -> float:
    name, default = _TIMEOUTS.get(path, ("OREKIT_TIMEOUT_S", 30))
    return float(getattr(settings, name, default))


def _iso(dt: datetime | None) -> str | None:
    if dt is None:
        return None
    return dt.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")


def _tuple3(seq) -> tuple[float, float, float]:
    return (float(seq[0]), float(seq[1]), float(seq[2]))


def _cartesian(states) -> list[dict]:
    return [
        {"epoch": _iso(sv.epoch), "r_km": list(sv.r_km), "v_kms": list(sv.v_kms)}
        for sv in states
    ]


class OrekitBackend(PropagatorBackend):
    name = "orekit"
    display_frames = DROPDOWN_FRAMES

    def model_config(self) -> dict:
        """The simulation model the sidecar is flying (its ``GET /config``).

        Read it rather than restating it: the model is defined in one place so
        that two services' forecasts are comparable, and a client that keeps its
        own copy has just made a second definition. See
        :mod:`yksa_orbital.model_config` for the cached accessor callers want.
        """
        return self._get("/config")

    def _get(self, path: str) -> dict:
        try:
            resp = httpx.get(_service_url() + path, timeout=_timeout(path))
        except httpx.HTTPError as exc:
            raise PropagationError(f"Orekit sidecar unreachable: {exc}") from exc
        if resp.status_code >= 400:
            raise PropagationError(f"Orekit sidecar error {resp.status_code}")
        return resp.json()

    def _post(self, path: str, payload: dict) -> dict:
        try:
            resp = httpx.post(
                _service_url() + path, json=payload, timeout=_timeout(path),
            )
        except httpx.HTTPError as exc:
            raise PropagationError(f"Orekit sidecar unreachable: {exc}") from exc
        if resp.status_code >= 400:
            try:
                detail = resp.json().get("error", resp.text)
            except Exception:  # noqa: BLE001
                detail = resp.text
            if resp.status_code == 503:
                # The sidecar keeps a worker free for interactive queries and
                # turns heavy work away rather than queueing it. Distinct from a
                # 400 so the caller retries instead of storing a failed run.
                raise BackendBusy(detail)
            raise PropagationError(f"Orekit sidecar error {resp.status_code}: {detail}")
        return resp.json()

    def _run(self, path: str, omm: dict, spacecraft: dict, options: dict) -> dict:
        """Shared payload assembly for the decay and ensemble routes."""
        payload = {"message": omm_message(omm)}
        for source in (spacecraft, options):
            payload.update({k: v for k, v in source.items() if v is not None})
        at = payload.get("at")
        if isinstance(at, datetime):
            payload["at"] = _iso(at)
        return self._post(path, payload)

    def state_at(self, omm: dict, at: datetime | None = None) -> StateVector:
        return _state_from_response(self._post("/state", {
            "message": omm_message(omm),
            "at": _iso(at),
            "frames": list(_CORE_FRAMES),
        }))

    def ephemeris(
        self, omm: dict, start: datetime, stop: datetime, step_s: float,
    ) -> list[StateVector]:
        data = self._post("/ephemeris", {
            "message": omm_message(omm),
            "start": _iso(start), "stop": _iso(stop),
            "step_s": float(step_s), "frame": "TEME",
        })
        element_epoch = _parse_iso(data.get("element_epoch"))
        return [_ephemeris_point(p, element_epoch) for p in data.get("states", [])]

    def transform(self, states, frame_in, frame_out):
        frame_in = frame_in.upper()
        frame_out = frame_out.upper()
        if frame_in == frame_out:
            return list(states)
        data = self._post("/transform", {
            "frame_in": frame_in,
            "frame_out": frame_out,
            "states": _cartesian(states),
        })
        return [
            StateVector(
                epoch=_parse_iso(res.get("epoch")) or src.epoch,
                frame=frame_out,
                r_km=_tuple3(res["r_km"]),
                v_kms=_tuple3(res["v_kms"]),
                element_epoch=src.element_epoch,
            )
            for src, res in zip(states, data.get("states", []))
        ]

    def fit_tle(self, states, frame, template_omm):
        return self._post("/fit_tle", {
            "states": _cartesian(states),
            "frame": frame.upper(),
            "template_message": omm_message(template_omm),
        })

    def decay(self, omm: dict, spacecraft: dict, **options) -> dict:
        """Propagate to re-entry (see ``services/orekit/decay.py``).

        Minutes-long, so it is only ever called from Celery.
        """
        return self._run("/decay", omm, spacecraft, options)

    def decay_ensemble(self, omm: dict, spacecraft: dict, **options) -> dict:
        """P10/P50/P90 lifetimes over resampled solar cycles.

        This is :meth:`decay` repeated once per realization, so it costs tens of
        minutes. Never call it from a request path.
        """
        return self._run("/decay_ensemble", omm, spacecraft, options)

    def refresh_space_weather(self, *, force: bool = True) -> dict:
        """Pull the current CSSI (and, if configured, MSAFE) file over the deployed
        one on the sidecar. Returns the sidecar's status dict; a failed download
        is reported there, not raised, so a scheduled refresh never errors on a
        transient network problem.
        """
        return self._post("/refresh_space_weather", {"force": bool(force)})

    def fit_drag(self, elements: list[dict], **options) -> dict:
        """Fit a ballistic coefficient from an element history.

        Cheap next to :meth:`decay` -- a linear fit plus one orbit of density
        evaluations -- so it uses the ordinary timeout.
        """
        payload = {"messages": omm_messages(elements)}
        payload.update({k: v for k, v in options.items() if v is not None})
        return self._post("/fit_drag", payload)

    def space_weather(
        self, start: datetime, stop: datetime, **options,
    ) -> dict:
        """Observed daily F10.7 / Ap between two dates.

        As cheap as :meth:`fit_drag` -- a provider lookup per day, no
        propagation -- so it uses the ordinary timeout.
        """
        payload = {"start": _iso(start), "stop": _iso(stop)}
        payload.update({k: v for k, v in options.items() if v is not None})
        return self._post("/space_weather", payload)

    def state_in_frames(
        self, omm: dict, at: datetime | None = None, frames=None,
    ) -> dict:
        """Multi-frame state for the coordinate-system dropdown.

        The sidecar's ``/state`` payload already matches the
        :meth:`PropagatorBackend.state_in_frames` contract, so nothing
        downstream special-cases Orekit.
        """
        return self._post("/state", {
            "message": omm_message(omm),
            "at": _iso(at),
            "frames": list(frames or DROPDOWN_FRAMES),
        })


def _state_from_response(data: dict) -> StateVector:
    states = data.get("states", {})
    teme = states.get("TEME")
    if not teme:
        raise PropagationError("Orekit response missing the TEME state")
    itrf = states.get("ITRF") or {}
    geodetic = data.get("geodetic")
    return StateVector(
        epoch=_parse_iso(data.get("epoch")),
        frame="TEME",
        r_km=_tuple3(teme["r_km"]),
        v_kms=_tuple3(teme["v_kms"]),
        ecef_km=_tuple3(itrf["r_km"]) if itrf.get("r_km") else None,
        ecef_v_kms=_tuple3(itrf["v_kms"]) if itrf.get("v_kms") else None,
        geodetic=_tuple3(geodetic) if geodetic else None,
        element_epoch=_parse_iso(data.get("element_epoch")),
        warnings=list(data.get("warnings") or []),
    )


def _ephemeris_point(point: dict, element_epoch: datetime | None) -> StateVector:
    ecef = point.get("ecef_km")
    ecef_v = point.get("ecef_v_kms")
    geodetic = point.get("geodetic")
    return StateVector(
        epoch=_parse_iso(point.get("epoch")),
        frame="TEME",
        r_km=_tuple3(point["r_km"]),
        v_kms=_tuple3(point["v_kms"]),
        ecef_km=_tuple3(ecef) if ecef else None,
        ecef_v_kms=_tuple3(ecef_v) if ecef_v else None,
        geodetic=_tuple3(geodetic) if geodetic else None,
        element_epoch=element_epoch,
    )