murphy_unified/engine_bridge.py

"""
Unified 3-engine bridge for Murphy Unified dashboard.

Dispatches cycle simulations to one of three engines:
  - euler       : cryosim.engine.fast.ICV_open  (H2+N2, ~3-5s)
  - lightspeed  : cycle2mdot_10var_njit.ICV_open_10var  (H2 only, ~1s, Numba)
  - general_ode : cycle2mdot_ode_v2.ODE_driver_v2  (H2+N2, ~25s)

Includes background Numba warmup thread and automatic N2/readiness fallbacks.
"""

import os
import sys
import time
import threading
from typing import Any, Dict, Tuple

import numpy as np

# ── Ensure ode_reformulation is on sys.path ──────────────────────────────────
# __file__ = .../CSH2/murphy-unified/murphy_unified/engine_bridge.py
# Two parents up reaches CSH2/
_ODE_DIR = os.path.normpath(
    os.path.join(os.path.dirname(__file__), os.pardir, os.pardir,
                 "cycle2mdot_acceleration", "ode_reformulation")
)
if _ODE_DIR not in sys.path:
    sys.path.insert(0, _ODE_DIR)

# Also add the parent of ode_reformulation so break_coolprop is importable
_ACCEL_DIR = os.path.dirname(_ODE_DIR)
if _ACCEL_DIR not in sys.path:
    sys.path.insert(0, _ACCEL_DIR)

# ── Euler engine (always available via cryosim) ──────────────────────────────
from cryosim.engine.fast import ICV_open  # noqa: E402

# ── Lightspeed engine (needs Numba) ──────────────────────────────────────────
_HAS_LIGHTSPEED = False
try:
    from cycle2mdot_10var_njit import ICV_open_10var  # noqa: E402
    _HAS_LIGHTSPEED = True
except ImportError:
    ICV_open_10var = None

# ── General-ODE engine ───────────────────────────────────────────────────────
# Dual backend under one "general_ode" label:
#   - H2 → DP45 compiled (cycle2mdot_ode_compiled.solve_cycle_compiled), ~20s/cycle
#   - N2 → scipy RK45    (cycle2mdot_ode_v2.ODE_driver_v2),              ~30s/cycle
# Both have formal O(h^5) error bounds. DP45 validated to <0.07% mdot error
# vs ode_v2 across 50–1150 barg H2 sweep (2026-04-14).
_HAS_GENERAL_ODE = False
try:
    from cycle2mdot_ode_v2 import ODE_driver_v2  # noqa: E402
    _HAS_GENERAL_ODE = True
except ImportError:
    ODE_driver_v2 = None

_HAS_DP45 = False
try:
    from cycle2mdot_ode_compiled import solve_cycle_compiled  # noqa: E402
    _HAS_DP45 = True
except ImportError:
    solve_cycle_compiled = None

# ── Engine metadata ──────────────────────────────────────────────────────────
ENGINES: Dict[str, Dict[str, Any]] = {
    "euler": {
        "label": "Euler (cryosim)",
        "fluids": ["h2", "n2"],
    },
    "lightspeed": {
        "label": "Lightspeed (Numba 10-var)",
        "fluids": ["h2"],
    },
    "general_ode": {
        "label": "General ODE v2",
        "fluids": ["h2", "n2"],
    },
}

# ── Numba warmup ─────────────────────────────────────────────────────────────
_warmup_event = threading.Event()
_warmup_ok = False


def _resolve_h2_ps_lut() -> str | None:
    """Locate h2_ps_lut.npz across local checkout and HF container layouts.

    Local: cycle2mdot_acceleration/h2_ps_lut.npz (next to break_coolprop/).
    HF:    /app/engines/h2_ps_lut.npz (flat).
    Falls back to scanning sys.path entries.
    """
    candidates = [os.path.join(_ACCEL_DIR, "h2_ps_lut.npz")]
    for entry in sys.path:
        if entry:
            candidates.append(os.path.join(entry, "h2_ps_lut.npz"))
    for p in candidates:
        if os.path.isfile(p):
            return p
    return None


def _numba_warmup():
    """Background warmup for Numba JIT functions.

    Warms up both break_coolprop.h2_props (used by Lightspeed and DP45) and
    the DP45 compiled solver itself (solve_cycle_compiled). The DP45 warmup
    runs a tiny cycle so the first user-facing H2/general_ode click doesn't
    pay the ~30–60s JIT compilation cost.
    """
    global _warmup_ok
    try:
        print("[engine_bridge] Numba warmup starting...")
        from break_coolprop import h2_props  # noqa: E402

        lut_path = _resolve_h2_ps_lut()
        if lut_path is None:
            print("[engine_bridge] h2_ps_lut.npz not found on disk or sys.path; "
                  "skipping LUT warmup.")
            _warmup_ok = False
        else:
            h2_props.init(lut_path)
            _warmup_ok = True
            print(f"[engine_bridge] Numba warmup complete (LUT: {lut_path}).")
    except Exception as exc:
        print(f"[engine_bridge] Numba warmup failed: {exc}")
        _warmup_ok = False
    finally:
        _warmup_event.set()

    # Pre-compile Lightspeed _ICV_loop_10var kernel so first user click
    # doesn't pay JIT cost (~3-5s on cpu-basic).
    if _HAS_LIGHTSPEED:
        try:
            print("[engine_bridge] Lightspeed warmup starting...")
            _lightspeed_warmup_call()
            print("[engine_bridge] Lightspeed warmup complete.")
        except Exception as exc:
            print(f"[engine_bridge] Lightspeed warmup failed (non-fatal): {exc}")

    # Pre-compile DP45 kernel (does not gate _warmup_ok — Lightspeed must
    # still be declared ready even if DP45 warmup fails)
    if _HAS_DP45:
        try:
            print("[engine_bridge] DP45 warmup starting...")
            _dp45_warmup_call()
            print("[engine_bridge] DP45 warmup complete.")
        except Exception as exc:
            print(f"[engine_bridge] DP45 warmup failed (non-fatal): {exc}")


def _lightspeed_warmup_call():
    """Run one dummy Lightspeed cycle to JIT-compile _ICV_loop_10var."""
    ICVparam = [20.5, 48.0, 5.0, 779.3, 33.4, 3.75, 7.36e-5, 1.0, 100]
    DCVparam = [8.3, 25.0, 3.79, 78.54, 7.8, 1.053, 2.6e-6, 0.8, 200]
    pump_geom = [40.3, 60.0, 120.0, 300.0, 0.8, 0.8, 0.026, 15.0,
                 37.0, 760000.0, 500.0, 0.01]
    proc_param = [300.0, 10.0, 4.0, 30.0, 0.006, 0.0, 0.5, 0.2]
    ICV_open_10var(
        Pexit_barg=350.0, speed_f=0.65,
        Ptank_barg=7.0, Psat_barg=2.0,
        ICVparam=ICVparam, DCVparam=DCVparam,
        pump_geom=pump_geom, proc_param=proc_param,
        fluid="h2", prtMode=0, flash_eff=0.0,
    )


def _dp45_warmup_call():
    """Run a single tiny cycle on DP45 to trigger Numba compilation."""
    ICVparam = [20.5, 48.0, 5.0, 779.3, 33.4, 3.75, 7.36e-5, 1.0, 100]
    DCVparam = [8.3, 25.0, 3.79, 78.54, 7.8, 1.053, 2.6e-6, 0.8, 200]
    pump_geom = [40.3, 60.0, 120.0, 300.0, 0.8, 0.8, 0.026, 15.0,
                 37.0, 760000.0, 500.0, 0.01]
    proc_param = [300.0, 10.0, 4.0, 30.0, 0.006, 0.0, 0.5, 0.2]
    solve_cycle_compiled(
        Pexit_barg=350.0, speed_f=0.65,
        Ptank_barg=7.0, Psat_barg=2.0,
        ICVparam=ICVparam, DCVparam=DCVparam,
        pump_geom=pump_geom, proc_param=proc_param,
        fluid="h2", prtMode=0, flash_eff=0.0,
    )


# Launch warmup daemon
_warmup_thread = threading.Thread(target=_numba_warmup, daemon=True)
_warmup_thread.start()


def is_lightspeed_ready() -> bool:
    """Return True if the Numba warmup completed successfully."""
    return _warmup_event.is_set() and _warmup_ok and _HAS_LIGHTSPEED


# ── Main dispatch ────────────────────────────────────────────────────────────
def cycle_sim(
    engine: str = "euler",
    Pexit_barg: float = 350.0,
    speed_f: float = 0.65,
    Ptank_barg: float = 7.0,
    Psat_barg: float = 2.0,
    ICVparam: list | None = None,
    DCVparam: list | None = None,
    pump_geom: list | None = None,
    proc_param: list | None = None,
    fluid: str = "h2",
    flash_eff: float = 0.0,
    **kwargs,
) -> Tuple[np.ndarray, Dict[str, Any]]:
    """Run a single pump-cycle simulation on the selected engine.

    Parameters
    ----------
    engine : str
        One of "euler", "lightspeed", "general_ode".
    Pexit_barg, speed_f, Ptank_barg, Psat_barg : float
        Operating conditions.
    ICVparam, DCVparam, pump_geom, proc_param : list or None
        Valve / geometry / process arrays.  None uses MVP 1.0 defaults.
    fluid : str
        "h2" or "n2".
    flash_eff : float
        Thermal mass film-boiling efficiency (0 = disabled).
    **kwargs
        Forwarded to the underlying engine.

    Returns
    -------
    out : np.ndarray
        Simulation output array (rows x 2).
    hist : dict
        History dict with at least "engine", "cpu_s", and optionally "fallback".
    """
    # ── MVP 1.0 Simplex defaults ─────────────────────────────────────────
    if ICVparam is None:
        ICVparam = [20.5, 48, 5, 779.3, 33.4, 3.75, 0.0000736, 1.0, 100]
    if DCVparam is None:
        DCVparam = [8.3, 25, 3.79, 78.54, 7.8, 1.053, 0.0000026, 0.8, 200]
    if pump_geom is None:
        pump_geom = [40.3, 60, 120, 300, 0.8, 0.8, 0.026, 15, 37, 760000, 500, 0.01]
    if proc_param is None:
        proc_param = [300, 10, 4, 30, 0.006, 0, 0.5, 0.2]

    fallback_msg = None

    # ── Fallback logic ───────────────────────────────────────────────────
    if engine == "lightspeed":
        if fluid.lower() not in ("h2", "hydrogen"):
            fallback_msg = (
                f"Lightspeed does not support {fluid}; falling back to general_ode."
            )
            engine = "general_ode"
        elif not is_lightspeed_ready():
            fallback_msg = (
                "Lightspeed Numba warmup not ready; falling back to euler."
            )
            engine = "euler"

    # ── Common kwargs for all engines ────────────────────────────────────
    common = dict(
        Pexit_barg=Pexit_barg,
        speed_f=speed_f,
        Ptank_barg=Ptank_barg,
        Psat_barg=Psat_barg,
        ICVparam=ICVparam,
        DCVparam=DCVparam,
        pump_geom=pump_geom,
        proc_param=proc_param,
        fluid=fluid,
        prtMode=0,
        flash_eff=flash_eff,
    )

    merged = {**common, **kwargs}

    t0 = time.perf_counter()

    # ── Dispatch ─────────────────────────────────────────────────────────
    if engine == "lightspeed":
        out, hist = ICV_open_10var(**merged)

    elif engine == "general_ode":
        # Fluid-based sub-dispatch: DP45 compiled for H2 (faster), ode_v2 for N2
        is_h2 = fluid.lower() in ("h2", "hydrogen")
        if is_h2 and _HAS_DP45:
            out, hist = solve_cycle_compiled(**merged)
            hist.setdefault("backend", "dp45_compiled")
        elif _HAS_GENERAL_ODE:
            if is_h2 and not _HAS_DP45:
                fallback_msg = (fallback_msg or "") + (
                    " | DP45 compiled backend unavailable; using scipy ode_v2"
                )
            out, hist = ODE_driver_v2(**merged)
            hist.setdefault("backend", "ode_v2_scipy")
        else:
            raise ImportError(
                "Neither DP45 compiled (cycle2mdot_ode_compiled) nor "
                "scipy ode_v2 (cycle2mdot_ode_v2) is available."
            )

    else:  # euler (default)
        # Euler (cryosim) does not accept exit_param / n_cycles — strip them
        euler_kw = {k: v for k, v in merged.items()
                    if k not in ("exit_param", "n_cycles")}
        out, hist = ICV_open(**euler_kw)

    cpu_s = time.perf_counter() - t0

    # ── Normalise output ─────────────────────────────────────────────────
    hist["engine"] = engine
    hist["cpu_s"] = cpu_s
    if fallback_msg is not None:
        hist["fallback"] = fallback_msg

    return out, hist