engine.py

"""Simulation engine: core runtime for Heliosail-RX.

KERNEL-CENTRIC ARCHITECTURE:
  The engine is the ONLY authority for:
    • Global state (position, velocity, attitude, power)
    • Time stepping and loop control
    • Module execution order
    • Error monitoring and recovery
    • Data aggregation and coupling

  NO direct communication between modules.
  NO global variables. 
  NO hidden dependencies.
  
  Modules are STATELESS coordinators that:
    • Receive PhysicsInput (read-only copy of state + params)
    • Perform computation independently
    • Return PhysicsOutput (acceleration, diagnostics, events)
    • Never call other modules
    • Never mutate state (frozen dataclass prevents it anyway)

The engine manages:
  1. Time stepping (fixed or adaptive)
  2. State propagation (calls physics modules)
  3. Event detection and handling
  4. Logging and checkpointing
  5. Error handling and recovery
"""
import time
import logging
import os
import sys
from typing import Dict, List, Callable, Optional, Any
from dataclasses import replace
import sys
import numpy as np

from models.core import (
    SimulationConfig,
    SimulationStep,
    SimulationResult,
    SpacecraftState,
    EnvironmentState,
    PhysicsInput,
    PhysicsOutput,
    StateDerivative,
    Event,
    EventType,
    TimeStep,
    AuditEntry,
    SolverType,
)
from kernel.validation import (
    ValidationFramework, ValidationLevel, ValidationReport
)
from kernel.safety import (
    FailureSafetyMonitor, SafetyThresholds, SafetyStatus
)


logger = logging.getLogger(__name__)

# Add environment-models into path to allow normal execution
_env_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "environment-models"))
if _env_path not in sys.path:
    sys.path.insert(0, _env_path)

try:
    from environment_models.solar_cycle import SolarCycleModel
    from environment_models.cme import CMEModel
    HAS_SPACE_WEATHER = True
except ImportError as e:
    logger.warning(f"Could not import space weather models: {e}")
    HAS_SPACE_WEATHER = False

# Enforcement: Only kernel can access this
_KERNEL_AUTHORITY = True  # Insurance check (will always be True)


class SimulationEngine:
    """Main simulation runtime.
    
    Usage:
        engine = SimulationEngine(config, physics_modules)
        engine.initialize(initial_state)
        result = engine.run(t_end=86400.0)
    """
    
    def __init__(
        self,
        config: SimulationConfig,
        physics_modules: Optional[Dict[str, Any]] = None,
        validation_level: ValidationLevel = ValidationLevel.OFF,
        enable_safety_monitoring: bool = True,
        safety_thresholds: Optional[SafetyThresholds] = None,
    ):
        """Initialize engine with configuration.
        
        Args:
            config: SimulationConfig defining Mission parameters
            physics_modules: Dict mapping module names to callable objects
                implementing PhysicsModule.compute(PhysicsInput) -> PhysicsOutput
            validation_level: ValidationLevel for convergence/energy checks
            enable_safety_monitoring: Enable real-time failure safety monitoring
            safety_thresholds: Custom safety thresholds (uses defaults if None)
        """
        self.config = config
        self.physics_modules = physics_modules or {}
        self.state = config.initial_state
        self.t_epoch = config.t_start
        self.step_num = 0
        self.trajectory: List[SimulationStep] = []
        self.events_history: List[Event] = []
        self.telemetry: Dict[str, List[tuple]] = {}
        self.audit_log: List[Dict[str, Any]] = []
        
        # Validation framework
        self.validation_level = validation_level
        self.validation_framework = ValidationFramework(validation_level)
        self.validation_report: Optional[ValidationReport] = None
        
        # Safety monitoring
        self.safety_monitoring_enabled = enable_safety_monitoring
        self.safety_monitor = FailureSafetyMonitor(
            thresholds=safety_thresholds or SafetyThresholds(),
            auto_halt=True
        )
        self.safety_monitor.on_critical.append(self._on_safety_critical)
        self.safety_monitor.on_halt.append(self._on_safety_halt)
        
        # Environmental Models
        self.solar_cycle_model = None
        self.cme_model = None
        if HAS_SPACE_WEATHER:
            self.solar_cycle_model = SolarCycleModel()
            self.cme_model = CMEModel()
            # If CME is enabled, maybe inject a random/realistic event for visual plotting testing
            if self.config.environment and getattr(self.config.environment, 'use_cme', False):
                # Add a strong CME event arriving around day 15 lasting 3 days
                self.cme_model.add_cme_event(
                    t_arrival_sec=15 * 86400,
                    duration_sec=3 * 86400,
                    speed_max=1200.0, # km/s
                    density_max=30.0  # cm^-3
                )
        
        # Callbacks registered by users
        self.callbacks: Dict[str, List[Callable]] = {
            "on_step_start": [],
            "on_physics_computed": [],
            "on_events_detected": [],
            "on_step_complete": [],
            "on_error": [],
        }
        
        # Logging
        self.last_telemetry_log_time = config.t_start
        self.last_checkpoint_time = config.t_start
    
    def register_callback(self, event_name: str, callback: Callable) -> None:
        """Register a callback to be invoked at a simulation event.
        
        Event names: on_step_start, on_physics_computed, on_events_detected,
                    on_step_complete, on_error
        """
        if event_name not in self.callbacks:
            raise ValueError(f"Unknown event: {event_name}")
        self.callbacks[event_name].append(callback)
    
    def initialize(self, initial_state: Optional[SpacecraftState] = None) -> None:
        """Initialize simulation to ready state.
        
        Args:
            initial_state: Initial spacecraft state. If None, uses config.initial_state.
        """
        if initial_state is None:
            if self.config.initial_state is None:
                raise ValueError("No initial state provided")
            initial_state = self.config.initial_state
        
        self.state = initial_state
        self.t_epoch = self.config.t_start
        self.step_num = 0
        self.trajectory = []
        self.events_history = []
        self.telemetry = {}
        self.audit_log = []
        
        logger.info(f"Simulation initialized: {self.config.name}")
        self._log_audit("init", f"Started simulation: {self.config.name}")
    
    def run(
        self,
        t_end: Optional[float] = None,
        callbacks: Optional[Dict[str, Callable]] = None,
    ) -> SimulationResult:
        """Execute simulation from current time to t_end.
        
        Args:
            t_end: End time [s epoch]. If None, uses config.t_end.
            callbacks: Dict of event_name -> callback to register before running
        
        Returns:
            SimulationResult with trajectory, telemetry, events, diagnostics
        """
        if t_end is None:
            t_end = self.config.t_end
        
        # Register callbacks
        if callbacks:
            for event_name, callback in callbacks.items():
                self.register_callback(event_name, callback)
        
        run_start = time.time()
        success = True
        error_msg = ""
        
        try:
            logger.info(f"Running simulation from t={self.t_epoch} to t={t_end}")
            
            while self.t_epoch < t_end and self.step_num < self.config.solver.max_steps:
                self._step()
            
            logger.info(f"Simulation complete: {len(self.trajectory)} steps")
            
        except Exception as e:
            success = False
            error_msg = str(e)
            logger.error(f"Simulation error: {e}")
            self._invoke_callbacks("on_error", {"exception": e})
        
        run_end = time.time()
        elapsed = run_end - run_start
        
        result = SimulationResult(
            config=self.config,
            trajectory=self.trajectory,
            telemetry=self.telemetry,
            events=self.events_history,
            audit_log=self.audit_log,
            run_start_time=run_start,
            run_end_time=run_end,
            elapsed_wall_time_sec=elapsed,
            n_steps=len(self.trajectory),
            success=success,
            error_message=error_msg,
        )
        
        # Run validation if enabled
        if self.validation_level != ValidationLevel.OFF and success:
            self._run_validation()
        
        return result
    
    def _step(self) -> None:
        """Execute one time step.
        
        SIMULATION EXECUTION PIPELINE (6 Stages):
        ═════════════════════════════════════════════════════════════════════
        
        Time advances from t → t+dt through these stages (in order):
        
        ┌─────────────────────────────────────────────────────────────────┐
        │ STAGE 1: ENVIRONMENT UPDATE                                     │
        │ Compute solar flux, magnetic field, plasma at current time      │
        └─────────────────────────────────────────────────────────────────┘
        
        ┌─────────────────────────────────────────────────────────────────┐
        │ STAGE 2: MODULE EXECUTION                                       │
        │ Call all physics modules with SAME state snapshot (no coupling) │
        │ Each module computes independently → StateDerivative            │
        └─────────────────────────────────────────────────────────────────┘
        
        ┌─────────────────────────────────────────────────────────────────┐
        │ STAGE 3: AGGREGATION                                            │
        │ Sum all accelerations from modules                              │
        │ Collect all events and diagnostics                              │
        └─────────────────────────────────────────────────────────────────┘
        
        ┌─────────────────────────────────────────────────────────────────┐
        │ STAGE 4: STATE INTEGRATION                                      │
        │ Use aggregated derivatives + ODE solver (RK4)                   │
        │ Produce new_state at t+dt                                       │
        └─────────────────────────────────────────────────────────────────┘
        
        ┌─────────────────────────────────────────────────────────────────┐
        │ STAGE 5: STATE UPDATE (KERNEL ONLY)                             │
        │ Replace internal state (only place state is updated)            │
        │ Increment time and step counter                                 │
        └─────────────────────────────────────────────────────────────────┘
        
        ┌─────────────────────────────────────────────────────────────────┐
        │ STAGE 6: RECORDING                                              │
        │ Archive step to trajectory, events to history                   │
        │ Log telemetry, checkpoint                                       │
        └─────────────────────────────────────────────────────────────────┘
        
        CRITICAL: Modules are called in Stage 2, NEVER elsewhere.
        CRITICAL: State is updated in Stage 5, ONLY place.
        CRITICAL: All modules see SAME state in Stage 2.
        """
        # ═════════════════════════════════════════════════════════════════
        # ANNOUNCE: Step starting
        # ═════════════════════════════════════════════════════════════════
        self._invoke_callbacks("on_step_start", {"t": self.t_epoch})
        
        # Determine time step
        dt = self.config.solver.dt_nominal
        if self.t_epoch + dt > self.config.t_end:
            dt = self.config.t_end - self.t_epoch
        
        if dt <= 0:
            logger.warning("Time step too small, stopping")
            return
        
        # ═════════════════════════════════════════════════════════════════
        # STAGE 1: ENVIRONMENT UPDATE
        # ═════════════════════════════════════════════════════════════════
        # Compute environmental conditions at current time
        environment = self._compute_environment(self.state, self.t_epoch)
        
        # ═════════════════════════════════════════════════════════════════
        # STAGE 2: MODULE EXECUTION
        # ═════════════════════════════════════════════════════════════════
        # Call all physics modules (kernel-orchestrated, independent)
        try:
            physics_outputs, diagnostics = self._compute_physics()
        except Exception as e:
            logger.error(f"Physics computation failed: {e}")
            raise
        
        self._invoke_callbacks("on_physics_computed", {"outputs": physics_outputs})
        
        # ═════════════════════════════════════════════════════════════════
        # STAGE 3: AGGREGATION
        # ═════════════════════════════════════════════════════════════════
        # Collect events from all modules
        events = []
        for module_name, output in physics_outputs.items():
            events.extend(output.events)
        
        self._invoke_callbacks("on_events_detected", {"events": events})
        
        # ═════════════════════════════════════════════════════════════════
        # STAGE 4: STATE INTEGRATION
        # ═════════════════════════════════════════════════════════════════
        # Propagate state using aggregated derivatives
        try:
            new_state = self._propagate_state(dt, physics_outputs)
        except Exception as e:
            logger.error(f"State propagation failed: {e}")
            raise
        
        # ═════════════════════════════════════════════════════════════════
        # STAGE 5: STATE UPDATE (KERNEL ONLY)
        # ═════════════════════════════════════════════════════════════════
        # This is the ONLY place where kernel state is updated
        self.state = new_state        # ← KERNEL AUTHORITY
        self.t_epoch += dt            # ← KERNEL AUTHORITY
        self.step_num += 1            # ← KERNEL AUTHORITY
        
        # ═════════════════════════════════════════════════════════════════
        # SAFETY MONITORING (Real-time failure detection)
        # ═════════════════════════════════════════════════════════════════
        if self.safety_monitoring_enabled:
            # Compute total acceleration magnitude for CFL check
            total_accel = np.zeros(3)
            for output in physics_outputs.values():
                if output.state_derivative and output.state_derivative.acceleration is not None:
                    total_accel += output.state_derivative.acceleration
            accel_mag = np.linalg.norm(total_accel)
            
            # Run safety monitoring
            safety_ok = self.safety_monitor.on_step_complete(
                state=self.state,
                dt=dt,
                acceleration_magnitude=accel_mag,
                energy=None,  # Optional: can compute from state if needed
                time_value=self.t_epoch
            )
            
            if not safety_ok:
                logger.critical(f"Safety halt at step {self.step_num}, t={self.t_epoch:.1f}s")
                raise RuntimeError(f"Safety violation: {self.safety_monitor.halt_reason}")
        
        # ═════════════════════════════════════════════════════════════════
        # STAGE 6: RECORDING
        # ═════════════════════════════════════════════════════════════════
        # Archive this step
        time_info = TimeStep(
            t_epoch=self.t_epoch,
            mission_elapsed=self.t_epoch - self.config.t_start,
            dt=dt,
            step_num=self.step_num,
        )
        
        step = SimulationStep(
            time_info=time_info,
            state=new_state,
            accelerations={name: output.state_derivative.acceleration for name, output in physics_outputs.items()},
            events=events,
            diagnostics=diagnostics,
        )
        
        # Kernel maintains trajectory (only place these are modified)
        # Downsample trajectory saving based on telemetry interval to ~1000 points max
        save_interval = self.config.telemetry_step_s
        if not self.trajectory or (self.t_epoch - self.trajectory[-1].time_info.t_epoch >= save_interval) or self.t_epoch >= self.config.t_end:
            self.trajectory.append(step)
            
        self.events_history.extend(events)
        
        # Logging (kernel manages telemetry)
        self._log_telemetry(step)
        self._checkpoint(step)
        
        self._invoke_callbacks("on_step_complete", {"step": step})
    
    def _compute_physics(self) -> tuple:
        """STAGE 2: MODULE EXECUTION
        
        Compute forces/torques from all enabled physics modules.
        
        KERNEL PRINCIPLE:
          • Each module receives the SAME state snapshot (t_epoch)
          • Each module computes independently
          • Modules execute in deterministic order (dict keys)
          • No module sees output of other modules at compute time
          • Kernel aggregates all results AFTER all modules return
          • No inter-module communication
          
        PIPELINE GUARANTEE:
          All modules see state at time t (frozen).
          No module sees evolved state until next timestep.
          No module couples to another module's output this step.
        
        Returns:
            (physics_outputs: Dict[str, PhysicsOutput],
             aggregated_diagnostics: Dict[str, float])
        """
        outputs = {}
        diagnostics = {}
        
        # Create snapshot: this is what ALL modules see (kernel controls this)
        snapshot_state = self.state
        snapshot_time = self.t_epoch
        
        # All modules receive same environment (kernel computes once)
        environment = self._compute_environment(snapshot_state, snapshot_time)
        
        # Execute modules in stable, deterministic order
        module_names = sorted(self.physics_modules.keys())
        
        for module_name in module_names:
            module = self.physics_modules[module_name]
            
            try:
                # Create isolated input for this module
                # Module can ONLY see this, nothing else
                inp = PhysicsInput(
                    state=snapshot_state,           # Read-only frozen dataclass
                    environment=environment,        # Current environmental conditions
                    t_epoch=snapshot_time,
                    params=self.config.physics.module_params.get(module_name, {}),
                )
                
                # Module computes independently with no side effects
                # Uses new interface: compute_derivatives(state, environment, t)
                # Falls back to legacy: compute(PhysicsInput) if needed
                if hasattr(module, 'compute_derivatives'):
                    out_deriv = module.compute_derivatives(snapshot_state, environment, snapshot_time)
                    if hasattr(module, 'compute_outputs'):
                        diags = module.compute_outputs(snapshot_state, environment, snapshot_time)
                    elif hasattr(module, 'diagnostics'):
                        diags = module.diagnostics()
                    else:
                        diags = {}
                        
                    out = PhysicsOutput(
                        state_derivative=out_deriv,
                        diagnostics=diags,
                        valid=True
                    )
                else:
                    out = module.compute(inp)
                
                # Kernel collects results
                outputs[module_name] = out
                diagnostics.update({
                    f"{module_name}__{k}": v
                    for k, v in out.diagnostics.items()
                })
                
            except Exception as e:
                logger.warning(f"Physics module '{module_name}' failed: {e}")
                # Continue with other modules (error isolation)
                outputs[module_name] = PhysicsOutput(
                    state_derivative=StateDerivative(
                        velocity_derivative=(0, 0, 0),
                        acceleration=(0, 0, 0),
                        angular_acceleration=(0, 0, 0)
                    ),
                    diagnostics={"error": str(e)},
                    valid=False
                )
        
        return outputs, diagnostics
    
    def _compute_environment(self, state: SpacecraftState, t_epoch: float) -> EnvironmentState:
        """Compute environmental state at given time.
        
        Args:
            state: Current spacecraft state
            t_epoch: Absolute time [s]
            
        Returns:
            EnvironmentState with current conditions
        """
        import math
        pos = getattr(state, 'position', (0.0, 0.0, 0.0))
        r = max(1.0, math.sqrt(sum(x*x for x in pos)))
        
        # Base flux based on 1 AU distance
        au = 1.495978707e11
        
        # Calculate solar flux using SolarCycleModel if enabled
        solar_flux_1au = 1361.0
        if self.solar_cycle_model and self.config.environment and getattr(self.config.environment, 'use_solar_cycle', False):
            solar_flux_1au = self.solar_cycle_model.irradiance(t_epoch, mode="stochastic")
        
        # Scale to current distance
        solar_flux = solar_flux_1au * (au / r)**2
        
        # Interplanetary Magnetic Field (Parker Spiral Approximation)
        # B0 at 1 AU is approx 5 nanoTesla (5e-9 T)
        b0 = 5e-9
        r_au = r / au
        # Radial decay (1/r^2), Azimuthal decay (1/r)
        Br = b0 / (r_au**2)
        # Simplified azimuthal component assuming nominal solar rotation and solar wind speed
        Bphi = -b0 / r_au 
        
        # Convert to Cartesian (assuming r vector is position pos)
        # For simplicity in this engine, we rough it as radial + azimuthal in XY plane
        xy_dist = math.sqrt(pos[0]**2 + pos[1]**2)
        if xy_dist > 0:
            cos_theta = pos[0] / xy_dist
            sin_theta = pos[1] / xy_dist
            Bx = Br * cos_theta - Bphi * sin_theta
            By = Br * sin_theta + Bphi * cos_theta
            Bz = 0.0
            magnetic_field = (Bx, By, Bz)
        else:
            magnetic_field = (Br, 0.0, 0.0)
        
        # Calculate solar wind density using CMEModel if enabled
        solar_wind_density = 5e6  # kg/m^3 baseline rough default
        if self.cme_model and self.config.environment and getattr(self.config.environment, 'use_cme', False):
            wind_state = self.cme_model.solar_wind_state(t_epoch, mode="stochastic")
            # wind_state["n_sw"] is in cm^-3 proton density
            # Density in kg/m^3: n * 1e6 * proton_mass
            proton_mass = 1.67262192e-27 # kg
            solar_wind_density = wind_state["n_sw"] * 1e6 * proton_mass
            
        radiation_index = 2.0  # Kp-like index
        
        return EnvironmentState(
            solar_flux=solar_flux,
            magnetic_field=magnetic_field,
            solar_wind_density=solar_wind_density,
            radiation_index=radiation_index,
            epoch_sec=t_epoch
        )
    
    def _propagate_state(
        self,
        dt: float,
        physics_outputs: Dict[str, PhysicsOutput],
    ) -> SpacecraftState:
        """STAGE 4: STATE INTEGRATION
        
        Propagate state using derivatives from physics modules.
        
        INTEGRATION ALGORITHM (Euler - fast but less accurate):
        
          position(t+dt) = position(t) + velocity(t) * dt
          velocity(t+dt) = velocity(t) + acceleration(t) * dt
        
        For higher accuracy, use RK4 integration (future enhancement).
        
        KERNEL AUTHORITY:
          • ONLY the kernel can modify state
          • ONLY the kernel decides integration method
          • State is immutable (frozen dataclass)
          • Modules never see new state until next timestep
          • Temporal ordering is guaranteed by kernel
        
        AGGREGATION (Stage 3/4 boundary):
          • Sum all accelerations from all modules
          • Modules never see each other's outputs
          • Aggregation is deterministic (sorted order)
        
        Returns:
            new_state: SpacecraftState at time t+dt (never None)
        """
        # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
        # AGGREGATION: Sum all accelerations and mass rates from all modules
        # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
        total_accel = [0.0, 0.0, 0.0]
        total_mass_rate = 0.0
        total_degradation_rate = 0.0
        
        for module_name in sorted(physics_outputs.keys()):  # Deterministic order
            output = physics_outputs[module_name]
            if output.valid:
                for i in range(3):
                    total_accel[i] += output.state_derivative.acceleration[i]
                
                # Check for legacy fallback vs new interface
                if hasattr(output.state_derivative, 'mass_rate'):
                    total_mass_rate += output.state_derivative.mass_rate
                elif hasattr(output.state_derivative, 'mass_derivative'):
                    total_mass_rate += output.state_derivative.mass_derivative
                    
                if hasattr(output.state_derivative, 'degradation_rate'):
                    deg_val = getattr(output.state_derivative, 'degradation_rate', 0.0)
                    total_degradation_rate += deg_val
                    if deg_val != 0.0:
                        print(f"DEBUG {module_name}: deg_val={deg_val:.2e}, dt={dt}")
        
        # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
        # INTEGRATION: Euler method
        # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
        r = self.state.position
        v = self.state.velocity
        
        # Position step: r(t+dt) = r(t) + v(t)*dt
        new_r = (
            r[0] + v[0] * dt,
            r[1] + v[1] * dt,
            r[2] + v[2] * dt,
        )
        
        # Velocity step: v(t+dt) = v(t) + a(t)*dt
        new_v = (
            v[0] + total_accel[0] * dt,
            v[1] + total_accel[1] * dt,
            v[2] + total_accel[2] * dt,
        )
        
        # Mass step: m(t+dt) = m(t) + dm(t)*dt (mass_rate is negative for consumption)
        new_mass = max(1.0, self.state.mass + (total_mass_rate * dt))
        
        # Degradation step: deg(t+dt) = deg(t) + d(deg)(t)*dt (rate is negative for deterioration)
        new_degradation = max(0.0, min(1.0, getattr(self.state, 'degradation_factor', 1.0) + (total_degradation_rate * dt)))
        
        # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
        # CREATE NEW STATE (immutable - kernel authority only)
        # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
        new_state = replace(
            self.state,
            position=new_r,
            velocity=new_v,
            mass=new_mass,
            degradation_factor=new_degradation,
            epoch_sec=self.t_epoch + dt,
            mission_elapsed_sec=self.t_epoch + dt - self.config.t_start,
        )
        
        # Sanity check: state is immutable (frozen dataclass)
        assert isinstance(new_state, SpacecraftState), "Kernel should always return SpacecraftState"
        
        return new_state
    
    def _log_telemetry(self, step: SimulationStep) -> None:
        """Log telemetry if interval has elapsed."""
        if step.time_info.t_epoch - self.last_telemetry_log_time >= self.config.telemetry_step_s:
            t = step.time_info.t_epoch
            
            # Log key scalars
            for key in ["r", "v"]:
                if key not in self.telemetry:
                    self.telemetry[key] = []
            
            # Magnitude of position/velocity
            r_mag = sum(x**2 for x in step.state.r) ** 0.5
            v_mag = sum(x**2 for x in step.state.v) ** 0.5
            
            self.telemetry["r"].append((t, r_mag))
            self.telemetry["v"].append((t, v_mag))
            
            # Add diagnostics
            for diag_name, diag_value in step.diagnostics.items():
                if diag_name not in self.telemetry:
                    self.telemetry[diag_name] = []
                self.telemetry[diag_name].append((t, diag_value))
            
            self.last_telemetry_log_time = t
            logger.debug(f"Logged telemetry at t={t:.1f}s: r={r_mag:.1e} m, v={v_mag:.1e} m/s")
    
    def _checkpoint(self, step: SimulationStep) -> None:
        """Save checkpoint if interval has elapsed."""
        if step.time_info.t_epoch - self.last_checkpoint_time >= self.config.checkpoint_step_s:
            logger.debug(f"Checkpoint at t={step.time_info.t_epoch:.1f}s, step={step.time_info.step_num}")
            # In production, write to disk (HDF5, JSON, etc.)
            self._log_audit(
                "checkpoint",
                f"Saved checkpoint at t={step.time_info.t_epoch}",
            )
            self.last_checkpoint_time = step.time_info.t_epoch
    
    
    def _run_validation(self) -> None:
        """Run validation checks on completed trajectory.
        
        Performs convergence tests, energy conservation checks, and
        reference mission comparisons based on validation_level.
        """
        if len(self.trajectory) == 0:
            logger.warning("No trajectory to validate")
            return
        
        try:
            # Extract trajectory states and times
            trajectory_states = [step.state for step in self.trajectory]
            times = [step.time_info.t_epoch for step in self.trajectory]
            
            # Run validation framework
            self.validation_report = self.validation_framework.run_validation(
                trajectory=trajectory_states,
                times=times,
                mission_name=self.config.name
            )
            
            # Log validation results
            logger.info(self.validation_report.summary())
            
            if self.validation_level in [ValidationLevel.DEBUG]:
                logger.info(self.validation_report.detailed_report())
            
            # Check for validation failures
            if not self.validation_report.overall_pass():
                logger.warning(f"Validation FAILED: {self.validation_report.failed_tests} test(s)")
            else:
                logger.info("Validation PASSED")
            
            self._log_audit(
                "validation",
                f"Validation complete: {self.validation_report.summary()}",
                {"report": self.validation_report}
            )
        
        except Exception as e:
            logger.error(f"Validation error: {e}")
            self._log_audit("validation_error", f"Validation failed: {e}")
    
    def _on_safety_critical(self, alert) -> None:
        """Callback when safety system detects critical issue."""
        logger.critical(
            f"Safety critical at step {alert.step_number}: {alert.check_name}\n"
            f"  {alert.message}"
        )
        self._log_audit(
            "safety_critical",
            f"{alert.check_name}: {alert.message}",
            {"alert": alert}
        )
    
    def _on_safety_halt(self, reason: str) -> None:
        """Callback when safety system triggers halt."""
        logger.critical(f"SAFETY HALT: {reason}")
        self._log_audit("safety_halt", reason)
    
    def _invoke_callbacks(self, event_name: str, data: Dict[str, Any]) -> None:
        """Invoke all registered callbacks for an event."""
        for callback in self.callbacks.get(event_name, []):
            try:
                callback(data)
            except Exception as e:
                logger.warning(f"Callback error in {event_name}: {e}")
    
    def _log_audit(self, action_type: str, description: str, data: Optional[Dict] = None) -> None:
        """Record action to audit log."""
        entry = {
            "timestamp": time.time(),
            "t_epoch": self.t_epoch,
            "action_type": action_type,
            "description": description,
            "data": data or {},
        }
        self.audit_log.append(entry)