claudes_game/scripts/sim_engine.py

#!/usr/bin/env python3
"""
Generic orbital mechanics simulation engine.
Replicates the exact physics from src/orbital_mechanics.cpp and src/simulation.cpp.

Usage:
    from sim_engine import Simulator
    sim = Simulator("path/to/config.toml", dt=60.0)
    sim.run(steps=1000)
    for event in sim.events:
        print(event)
"""

import math
import tomllib
from dataclasses import dataclass, field, replace
from typing import Dict, Tuple, Any


# =============================================================================
# Constants
# =============================================================================

G = 6.67430e-11
PARABOLIC_TOLERANCE = 1e-3
KEPLER_TOLERANCE = 1e-10
KEPLER_MAX_ITER = 50
VEL_DRIFT_THRESHOLD = 1e-6  # m/s


# =============================================================================
# Vector operations
# =============================================================================

def vadd(a, b):
    return (a[0]+b[0], a[1]+b[1], a[2]+b[2])

def vsub(a, b):
    return (a[0]-b[0], a[1]-b[1], a[2]-b[2])

def vscale(v, s):
    return (v[0]*s, v[1]*s, v[2]*s)

def vmag(v):
    return math.sqrt(v[0]**2 + v[1]**2 + v[2]**2)

def vdot(a, b):
    return a[0]*b[0] + a[1]*b[1] + a[2]*b[2]

def vcross(a, b):
    return (
        a[1]*b[2] - a[2]*b[1],
        a[2]*b[0] - a[0]*b[2],
        a[0]*b[1] - a[1]*b[0]
    )

def vnorm(v):
    m = vmag(v)
    if m < 1e-15:
        return (0.0, 0.0, 0.0)
    return (v[0]/m, v[1]/m, v[2]/m)

def normalize_angle(angle):
    while angle < 0.0:
        angle += 2.0 * math.pi
    while angle >= 2.0 * math.pi:
        angle -= 2.0 * math.pi
    return angle


# =============================================================================
# Data structures
# =============================================================================

@dataclass
class OrbitalElements:
    a: float = 0.0            # semi-major axis (elliptical) / semi-latus rectum (parabolic)
    e: float = 0.0            # eccentricity
    nu: float = 0.0           # true anomaly
    inc: float = 0.0          # inclination
    Omega: float = 0.0        # longitude of ascending node
    omega: float = 0.0        # argument of periapsis
    p: float = 0.0            # semi-latus rectum (parabolic only)


@dataclass
class Body:
    name: str = ""
    mass: float = 0.0
    radius: float = 0.0
    parent_index: int = -1
    orbit: OrbitalElements = field(default_factory=OrbitalElements)
    local_pos: Tuple[float, float, float] = (0.0, 0.0, 0.0)
    local_vel: Tuple[float, float, float] = (0.0, 0.0, 0.0)
    global_pos: Tuple[float, float, float] = (0.0, 0.0, 0.0)
    global_vel: Tuple[float, float, float] = (0.0, 0.0, 0.0)


@dataclass
class Event:
    """Recorded simulation event."""
    kind: str = "state"
    time: float = 0.0
    data: Dict[str, Any] = field(default_factory=dict)


# =============================================================================
# Kepler equation solvers (exact C++ logic)
# =============================================================================

def get_initial_trial_value(mean_anomaly, eccentricity):
    """Initial guess for Kepler solver (C++ get_initial_trial_value)."""
    return (mean_anomaly + eccentricity * math.sin(mean_anomaly)
            + ((eccentricity ** 2 / 2.0) * math.sin(2.0 * mean_anomaly)))


def solve_kepler_elliptical(mean_anomaly, eccentricity):
    E = get_initial_trial_value(mean_anomaly, eccentricity)
    E_prev = E + 2.0 * KEPLER_TOLERANCE
    for _ in range(KEPLER_MAX_ITER):
        if abs(E - E_prev) < KEPLER_TOLERANCE:
            break
        E_prev = E
        sin_E = math.sin(E)
        E = E - (E - eccentricity * sin_E - mean_anomaly) / (1.0 - eccentricity * math.cos(E))
    return E


# =============================================================================
# Coordinate transforms
# =============================================================================

def orbital_to_cartesian(elements, parent_mass):
    """Convert orbital elements to local position/velocity vectors."""
    mu = G * parent_mass
    a = elements.a
    e = elements.e
    nu = elements.nu

    if abs(e - 1.0) < PARABOLIC_TOLERANCE:
        p = elements.p
    else:
        p = a * (1.0 - e * e)

    r = p / (1.0 + e * math.cos(nu))

    x_orb = r * math.cos(nu)
    y_orb = r * math.sin(nu)

    vx_orb = -math.sqrt(mu / p) * math.sin(nu)
    vy_orb = math.sqrt(mu / p) * (e + math.cos(nu))

    # z-x-z rotation: Rz(Omega) * Rx(inc) * Rz(omega)
    cos_w, sin_w = math.cos(elements.omega), math.sin(elements.omega)
    x1 = x_orb * cos_w - y_orb * sin_w
    y1 = x_orb * sin_w + y_orb * cos_w

    cos_i, sin_i = math.cos(elements.inc), math.sin(elements.inc)
    x2 = x1
    y2 = y1 * cos_i
    z2 = y1 * sin_i

    cos_O, sin_O = math.cos(elements.Omega), math.sin(elements.Omega)
    pos = (x2 * cos_O - y2 * sin_O,
           x2 * sin_O + y2 * cos_O,
           z2)

    vx1 = vx_orb * cos_w - vy_orb * sin_w
    vy1 = vx_orb * sin_w + vy_orb * cos_w
    vx2 = vx1
    vy2 = vy1 * cos_i
    vz2 = vy1 * sin_i

    vel = (vx2 * cos_O - vy2 * sin_O,
           vx2 * sin_O + vy2 * cos_O,
           vz2)

    return pos, vel


def cartesian_to_orbital_elements(pos, vel, parent_mass):
    """Convert local position/velocity to orbital elements."""
    mu = G * parent_mass
    r = vmag(pos)
    v = vmag(vel)
    v_sq = v * v

    specific_energy = -mu / r + v_sq / 2.0
    h_vec = vcross(pos, vel)
    h = vmag(h_vec)

    # Eccentricity vector: e_vec = (v² - μ/r)r - (r·v)v  all divided by μ
    r_dot_v = vdot(pos, vel)
    e_vec = ((v_sq - mu / r) * pos[0] - r_dot_v * vel[0]) / mu, \
            ((v_sq - mu / r) * pos[1] - r_dot_v * vel[1]) / mu, \
            ((v_sq - mu / r) * pos[2] - r_dot_v * vel[2]) / mu
    e = vmag(e_vec)

    # Semi-major axis
    if abs(specific_energy) < 1e-10:
        a = 1e10
    else:
        a = -mu / (2.0 * specific_energy)

    # True anomaly
    if e < 1e-10:
        # Nearly circular: use argument of latitude
        n_vec = vcross((0.0, 0.0, 1.0), h_vec)
        n_mag = vmag(n_vec)
        sin_i = (n_mag / h) if h > 1e-10 else 1.0
        if sin_i > 1e-6 and n_mag > 1e-10:
            # Well-defined ascending node: compute argument of latitude
            x_AN = n_vec[0] / n_mag
            y_AN = n_vec[1] / n_mag
            hcn = vcross(h_vec, n_vec)
            hcn_mag = vmag(hcn)
            if hcn_mag > 1e-10:
                hcn = vscale(hcn, 1.0 / hcn_mag)
            r_xAN = pos[0] * x_AN + pos[1] * y_AN
            r_yAN = pos[0] * hcn[0] + pos[1] * hcn[1] + pos[2] * hcn[2]
            nu = math.atan2(r_yAN, r_xAN)
        else:
            # Nearly coplanar: use atan2(y, x) as argument of latitude
            nu = math.atan2(pos[1], pos[0])
        nu = normalize_angle(nu)
    else:
        cos_nu = vdot(pos, e_vec) / (r * e)
        cos_nu = max(-1.0, min(1.0, cos_nu))
        sin_nu = None

        if abs(cos_nu) > 1.0 - 1e-10:
            h_cross_e = vcross(h_vec, e_vec)
            denom = r * e * h
            sin_nu = vdot(pos, h_cross_e) / denom if denom > 1e-10 else 0.0
        else:
            r_cross_h = vcross(pos, h_vec)
            denom = r * e * h
            sin_nu = vdot(r_cross_h, e_vec) / denom if denom > 1e-10 else 0.0

        nu = math.atan2(sin_nu, cos_nu)
        if nu == -math.pi:
            nu = math.pi
        nu = normalize_angle(nu)

    # Inclination
    i = math.acos(h_vec[2] / h) if h > 1e-10 else 0.0

    # RAAN
    n_vec = vcross((0.0, 0.0, 1.0), h_vec)
    n_mag = vmag(n_vec)
    if n_mag > 1e-10:
        Omega = math.acos(n_vec[0] / n_mag)
        if n_vec[1] < 0.0:
            Omega = 2.0 * math.pi - Omega
    else:
        Omega = 0.0

    # Argument of periapsis
    inclination_threshold = 0.01
    if e > 1e-10 and n_mag > 1e-10 and i > inclination_threshold:
        cos_omega = vdot(e_vec, n_vec) / (e * n_mag)
        n_cross_e = vcross(n_vec, e_vec)
        sin_omega = vdot(n_cross_e, h_vec) / (e * n_mag * h)
        omega = math.atan2(sin_omega, cos_omega)
        if omega < 0.0:
            omega += 2.0 * math.pi
    elif e > 1e-10:
        omega = math.atan2(e_vec[1], e_vec[0])
        if omega < 0.0:
            omega += 2.0 * math.pi
    else:
        omega = 0.0

    elements = OrbitalElements()
    if abs(e - 1.0) < 1e-3:
        elements.p = (h * h) / mu
    else:
        elements.a = a
    elements.e = e
    elements.nu = nu
    elements.inc = i
    elements.Omega = Omega
    elements.omega = omega

    return elements


# =============================================================================
# Propagation
# =============================================================================

def propagate(elements, dt, parent_mass):
    """Propagate orbital elements forward by dt. Returns new elements."""
    mu = G * parent_mass
    a = elements.a
    e = elements.e
    nu = elements.nu

    if abs(e - 1.0) < PARABOLIC_TOLERANCE:
        # Parabolic (Barker's equation)
        p = elements.p
        D = math.tan(nu / 2.0)
        M = D + (D * D * D) / 3.0
        n = math.sqrt(mu / (p ** 3.0))
        M = M + n * dt
        # Solve Barker's: D + D^3/3 = M
        c = 1.5 * M
        disc = c * c + 1.0
        sqrt_disc = math.sqrt(disc)
        D_new = math.cbrt(c + sqrt_disc) + math.cbrt(c - sqrt_disc)
        return replace(elements, nu=2.0 * math.atan(D_new))

    elif e < 1.0:
        # Elliptical
        n = math.sqrt(mu / (a ** 3.0))
        E = 2.0 * math.atan(math.sqrt((1.0 - e) / (1.0 + e)) * math.tan(nu / 2.0))
        M = E - e * math.sin(E)
        M = M + n * dt
        E_new = get_initial_trial_value(M, e)
        E_prev = E_new + 2.0 * KEPLER_TOLERANCE
        for _ in range(KEPLER_MAX_ITER):
            if abs(E_new - E_prev) < KEPLER_TOLERANCE:
                break
            E_prev = E_new
            sin_E = math.sin(E_new)
            E_new = E_new - (E_new - e * sin_E - M) / (1.0 - e * math.cos(E_new))
        nu_new = 2.0 * math.atan(math.sqrt((1.0 + e) / (1.0 - e)) * math.tan(E_new / 2.0))
        return replace(elements, nu=nu_new)

    else:
        # Hyperbolic
        raise NotImplementedError("hyperbolic propagation not yet implemented")


# =============================================================================
# Global coordinate computation
# =============================================================================

def compute_global_coordinates(bodies):
    """
    Compute global position/velocity for all bodies.
    Matches C++ compute_global_coordinates() exactly.
    """
    for body in bodies:
        if body.parent_index == -1:
            body.global_pos = body.local_pos
            body.global_vel = body.local_vel
        elif 0 <= body.parent_index < len(bodies):
            parent = bodies[body.parent_index]
            body.global_pos = vadd(body.local_pos, parent.global_pos)
            body.global_vel = vadd(body.local_vel, parent.global_vel)


# =============================================================================
# Velocity drift check
# =============================================================================

def check_velocity_drift(body, parent, parent_mass):
    """
    Check if local velocity has drifted from expected Keplerian velocity.
    If so, reconstruct orbital elements from current state.
    Matches C++ update_bodies_physics() drift check.
    """
    if parent is None:
        return

    _, expected_vel = orbital_to_cartesian(body.orbit, parent_mass)
    vel_diff = vmag(vsub(body.local_vel, expected_vel))
    if vel_diff > VEL_DRIFT_THRESHOLD:
        body.orbit = cartesian_to_orbital_elements(body.local_pos, body.local_vel, parent_mass)


# =============================================================================
# Body physics update
# =============================================================================

def update_body(bodies, body_index, dt):
    """
    Update a single body: drift check, propagation.
    Matches C++ update_bodies_physics() per-body logic (without SOI).
    """
    body = bodies[body_index]

    if body.parent_index == -1:
        return  # Root body doesn't propagate

    if 0 <= body.parent_index < len(bodies):
        parent = bodies[body.parent_index]
        check_velocity_drift(body, parent, parent.mass)
        body.orbit = propagate(body.orbit, dt, parent.mass)
        body.local_pos, body.local_vel = orbital_to_cartesian(body.orbit, parent.mass)


# =============================================================================
# TOML config loader
# =============================================================================

def load_config(config_path):
    """Load a TOML 1.0 config file and return parsed data."""
    with open(config_path, "rb") as f:
        return tomllib.load(f)


def bodies_from_config(config):
    """
    Create Body objects from TOML config.
    Parent references are resolved by name, then by index.
    """
    bodies = []
    name_to_idx = {}

    # First pass: create bodies without positions
    for body_cfg in config.get("bodies", []):
        orbit_cfg = body_cfg.get("orbit", {})
        elements = OrbitalElements(
            a=orbit_cfg.get("semi_major_axis", 0.0),
            e=orbit_cfg.get("eccentricity", 0.0),
            nu=orbit_cfg.get("true_anomaly", 0.0),
            inc=orbit_cfg.get("inclination", 0.0),
            Omega=orbit_cfg.get("longitude_of_ascending_node", 0.0),
            omega=orbit_cfg.get("argument_of_periapsis", 0.0),
        )

        parent_ref = body_cfg.get("parent_index", -1)
        if isinstance(parent_ref, str):
            # Resolve by name
            if parent_ref in name_to_idx:
                parent_index = name_to_idx[parent_ref]
            elif parent_ref == "Sun" or parent_ref == "root" or parent_ref == "-1":
                parent_index = -1
            else:
                raise ValueError(f"Unknown parent name: {parent_ref}")
        else:
            parent_index = int(parent_ref)

        body = Body(
            name=body_cfg.get("name", f"Body_{len(bodies)}"),
            mass=body_cfg.get("mass", 0.0),
            radius=body_cfg.get("radius", 0.0),
            parent_index=parent_index,
            orbit=elements,
        )
        bodies.append(body)
        name_to_idx[body.name] = len(bodies) - 1

    return bodies


# =============================================================================
# Initialization
# =============================================================================

def initialize_bodies(bodies):
    """
    Initialize orbital objects from orbital elements.
    Matches C++ initialize_orbital_objects() exactly (without SOI).
    """
    for i, body in enumerate(bodies):
        if body.parent_index >= 0 and body.parent_index < len(bodies):
            parent = bodies[body.parent_index]
            local_pos, local_vel = orbital_to_cartesian(body.orbit, parent.mass)
            body.local_pos = local_pos
            body.local_vel = local_vel
            body.global_pos = vadd(parent.global_pos, local_pos)
            body.global_vel = vadd(parent.global_vel, local_vel)
        else:
            body.local_pos = (0.0, 0.0, 0.0)
            body.local_vel = (0.0, 0.0, 0.0)
            body.global_pos = (0.0, 0.0, 0.0)
            body.global_vel = (0.0, 0.0, 0.0)


# =============================================================================
# Simulator — public API
# =============================================================================

class Simulator:
    """
    Generic orbital mechanics simulator.

    Usage:
        sim = Simulator("config.toml", dt=60.0)
        sim.run(steps=1000)

        # Access results
        for event in sim.events:
            print(event)

        # Access final state
        for body in sim.bodies:
            print(f"{body.name}: r={vmag(body.global_pos):.0f} m")
    """

    def __init__(self, config_path, dt=60.0):
        self.dt = dt
        self.time = 0.0
        self.events = []
        self._body_count = 0

        config = load_config(config_path)
        self.bodies = bodies_from_config(config)
        initialize_bodies(self.bodies)
        self._body_count = len(self.bodies)

    def run(self, steps):
        """Run simulation for the given number of timesteps."""
        for _ in range(steps):
            self._step()

    def _step(self):
        """Single simulation step. Matches C++ update_simulation() order."""
        # 1. Update body physics (drift, propagation)
        for i in range(self._body_count):
            update_body(self.bodies, i, self.dt)

        # 2. Compute global coordinates
        compute_global_coordinates(self.bodies)

        self.time += self.dt

    def record_state(self, label=""):
        """Record current simulation state as an event."""
        state = {}
        for body in self.bodies:
            r = vmag(body.global_pos)
            state[body.name] = {
                "r": r,
                "nu": body.orbit.nu,
                "a": body.orbit.a,
                "e": body.orbit.e,
                "parent": body.parent_index,
                "parent_name": self.bodies[body.parent_index].name if body.parent_index >= 0 else "root",
            }
        self.events.append(Event(kind="state", time=self.time, data={"label": label, "state": state}))

    def get_body(self, name_or_index):
        """Get a body by name or index."""
        if isinstance(name_or_index, int):
            return self.bodies[name_or_index]
        for body in self.bodies:
            if body.name == name_or_index:
                return body
        raise KeyError(f"Body not found: {name_or_index}")

    def print_summary(self):
        """Print a summary of all recorded state events."""
        for event in self.events:
            label = event.data.get("label", "")
            if label:
                print(f"\n*** {label} (t={event.time:.1f}s) ***")
            for name, info in event.data.get("state", {}).items():
                print(f"  {name}: r={info['r']:.0f} m, "
                      f"nu={math.degrees(info['nu']):.1f}°, "
                      f"a={info['a']:.0f}, e={info['e']:.6f}, "
                      f"parent={info['parent_name']}")