server/ui_demo.py

"""UI demo data for the Space Manufacturing control experience."""

from __future__ import annotations

import math
import sys
from threading import Lock
from typing import Any, Dict, List

from SpaceFactory.models import ManufacturingAction
from SpaceFactory.env import ManufacturingTaskEnv
from SpaceFactory.graders import ManufacturingTaskGrader


# ── geometry helpers ───────────────────────────────────────────────────────────

def _xyz_to_geo(position: List[float]) -> Dict[str, float]:
    """Convert [x, y, z] orbital position (km) to lat/lng/altitude_km."""
    x, y, z = float(position[0]), float(position[1]), float(position[2])
    radius = math.sqrt(x * x + y * y + z * z) or 1.0
    latitude  = math.degrees(math.asin(max(-1.0, min(1.0, z / radius))))
    longitude = math.degrees(math.atan2(y, x))
    altitude_km = max(radius - 6371.0, 200.0)   # platforms are at ~400 km notional
    return {"latitude": latitude, "longitude": longitude, "altitude_km": altitude_km}


def _orbit_path(platform_id: int, num_platforms: int, steps: int = 72) -> List[Dict[str, float]]:
    """
    Build a full orbit ring for a platform using a simple circular orbit.
    Platforms are equally spaced around the ring; each has a slight inclination
    to make the globe visually interesting.
    """
    inclination = math.radians(28.5 + platform_id * 8.0)   # ISS-like, varies per platform
    ascending_node = math.radians((360.0 / num_platforms) * platform_id)
    radius = 6771.0  # Earth radius + ~400 km

    segment: List[Dict[str, float]] = []
    for i in range(steps + 1):
        phase = (2.0 * math.pi * i) / steps
        cos_p = math.cos(phase)
        sin_p = math.sin(phase)
        cos_i = math.cos(inclination)
        sin_i = math.sin(inclination)
        cos_n = math.cos(ascending_node)
        sin_n = math.sin(ascending_node)

        x = radius * (cos_n * cos_p - sin_n * sin_p * cos_i)
        y = radius * (sin_n * cos_p + cos_n * sin_p * cos_i)
        z = radius * (sin_p * sin_i)

        geo = _xyz_to_geo([x, y, z])
        segment.append({"latitude": geo["latitude"], "longitude": geo["longitude"]})
    return segment


def _platform_geo(platform_id: int, num_platforms: int, step_count: int) -> Dict[str, float]:
    """
    Return the current lat/lng/altitude_km for a platform by advancing its
    orbital phase by step_count steps.
    """
    inclination    = math.radians(28.5 + platform_id * 8.0)
    ascending_node = math.radians((360.0 / num_platforms) * platform_id)
    # initial phase offset so platforms don't stack at the same point
    phase_offset   = (2.0 * math.pi / num_platforms) * platform_id
    # orbit period ~90 min → advance ~4° per step
    phase = phase_offset + math.radians(4.0 * step_count)

    cos_p = math.cos(phase)
    sin_p = math.sin(phase)
    cos_i = math.cos(inclination)
    sin_i = math.sin(inclination)
    cos_n = math.cos(ascending_node)
    sin_n = math.sin(ascending_node)

    radius = 6771.0
    x = radius * (cos_n * cos_p - sin_n * sin_p * cos_i)
    y = radius * (sin_n * cos_p + cos_n * sin_p * cos_i)
    z = radius * (sin_p * sin_i)

    return _xyz_to_geo([x, y, z])


# ── demo class ─────────────────────────────────────────────────────────────────

class ManufacturingUIDemo:
    """Keeps a small, isolated environment instance for the UI experience."""

    def __init__(self) -> None:
        self._lock = Lock()
        self._task_name = "medium"
        self._env = ManufacturingTaskEnv(task_name=self._task_name)
        self._observation = self._env.reset()
        self._last_reward = 0.0
        self._reward_history: List[float] = []
        # pre-compute orbit paths (they don't change)
        self._orbit_paths: Dict[int, List[Dict[str, float]]] = {}
        self._rebuild_orbit_paths()

    # ── public API ─────────────────────────────────────────────────────────────

    def reset(self, task_name: str = "medium") -> Dict[str, Any]:
        with self._lock:
            self._task_name = task_name
            self._env = ManufacturingTaskEnv(task_name=task_name)
            self._observation = self._env.reset()
            self._last_reward = 0.0
            self._reward_history = []
            self._rebuild_orbit_paths()
            return self._snapshot()

    def step(self) -> Dict[str, Any]:
        with self._lock:
            action_map = self._choose_actions()
            self._observation, reward, _, _info = self._env.step(
                ManufacturingAction(platform_actions=action_map)
            )
            self._last_reward = reward.value
            self._reward_history.append(round(reward.value, 3))
            if len(self._reward_history) > 100:
                self._reward_history = self._reward_history[-100:]
            return self._snapshot()

    def snapshot(self) -> Dict[str, Any]:
        with self._lock:
            return self._snapshot()

    # ── heuristic policy for the live demo ────────────────────────────────────

    def _choose_actions(self) -> Dict[int, str]:
        action_map: Dict[int, str] = {}
        has_open_window = len(self._observation.delivery_windows) > 0

        for p in self._observation.platforms:
            if p.energy < 15.0:
                action = "recharge"
            elif p.product_stock > 0 and has_open_window:
                action = "deliver"
            elif p.component_stock >= 10.0 and p.product_stock < 5:
                action = "assemble"
            elif p.material_stock >= 15.0 and p.component_stock < 30.0:
                action = "produce"
            elif p.energy < 40.0:
                action = "recharge"
            elif p.material_stock >= 15.0:
                action = "produce"
            else:
                action = "recharge"
            action_map[p.id] = action

        return action_map

    # ── orbit helpers ──────────────────────────────────────────────────────────

    def _rebuild_orbit_paths(self) -> None:
        n = len(self._observation.platforms)
        self._orbit_paths = {
            p.id: _orbit_path(p.id, n)
            for p in self._observation.platforms
        }

    # ── mission score ─────────────────────────────────────────────────────────

    def _mission_score(self) -> float:
        try:
            grader = ManufacturingTaskGrader(self._task_name)
            env_state = self._env.state()
            return grader.grade(
                env_state.metrics,
                env_state.step_count,
                env_state.platforms,
            )
        except Exception:
            return 0.0

    # ── snapshot serialiser ───────────────────────────────────────────────────

    def _snapshot(self) -> Dict[str, Any]:
        env_state = self._env.state()
        n = len(self._observation.platforms)
        step = self._observation.time_step

        platforms: List[Dict[str, Any]] = []
        for p in self._observation.platforms:
            geo = _platform_geo(p.id, n, step)
            platforms.append({
                "id": p.id,
                "position": p.position,
                "latitude": geo["latitude"],
                "longitude": geo["longitude"],
                "altitude_km": geo["altitude_km"],
                "energy": round(p.energy, 1),
                "material_stock": round(p.material_stock, 1),
                "component_stock": round(p.component_stock, 1),
                "product_stock": p.product_stock,
                "last_action": p.last_action,
                "route": self._orbit_paths.get(p.id, []),
            })

        delivery_windows: List[Dict[str, Any]] = [
            {
                "order_id": w.order_id,
                "product_type": w.product_type,
                "deadline": w.deadline,
                "reward_value": w.reward_value,
            }
            for w in self._observation.delivery_windows
        ]

        pending_orders: List[Dict[str, Any]] = [
            {
                "order_id": o.order_id,
                "product_type": o.product_type,
                "requires_assembly": o.requires_assembly,
            }
            for o in self._observation.pending_orders
        ]

        metrics = {k: float(v) for k, v in env_state.metrics.items()}

        return {
            "task_name": self._task_name,
            "step_count": step,
            "max_steps": env_state.max_steps,
            "done": env_state.done,
            "reward": round(self._observation.reward, 3),
            "last_reward": round(self._last_reward, 3),
            "total_reward": round(self._observation.total_reward, 3),
            "mission_score": self._mission_score(),
            "reward_history": list(self._reward_history),
            "platforms": platforms,
            "delivery_windows": delivery_windows,
            "pending_orders": pending_orders,
            "solar_conditions": dict(self._observation.solar_conditions),
            "metrics": metrics,
        }