"""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, }