factory.py

from __future__ import annotations

import math
import random
import uuid
from typing import Any, Dict, List, Optional, Tuple

from .models import (
    DeliveryWindow,
    ManufacturingAction,
    ManufacturingObservation,
    ManufacturingReward,
    PendingOrder,
    PlatformState,
)

# ─── Tuneable constants ────────────────────────────────────────────────────────
PRODUCE_MATERIAL_COST = 15.0
PRODUCE_ENERGY_COST = 10.0
PRODUCE_COMPONENT_GAIN = 12.0

ASSEMBLE_COMPONENT_COST = 10.0
ASSEMBLE_ENERGY_COST = 12.0

DELIVER_ENERGY_COST = 8.0

RECHARGE_BASE = 20.0

ENERGY_MAX = 100.0
MATERIAL_MAX = 100.0
COMPONENT_MAX = 100.0
PRODUCT_MAX = 10

MATERIAL_REFILL_RATE = 5.0  # materials added each step passively

PRODUCE_MATERIAL_THRESHOLD = PRODUCE_MATERIAL_COST
ASSEMBLE_COMPONENT_THRESHOLD = ASSEMBLE_COMPONENT_COST

# ─── Product types ─────────────────────────────────────────────────────────────
SIMPLE_PRODUCTS = ["alloy_panel", "circuit_board", "fuel_cell"]
ASSEMBLED_PRODUCTS = ["solar_array", "thruster_module", "habitat_unit"]
ALL_PRODUCTS = SIMPLE_PRODUCTS + ASSEMBLED_PRODUCTS


class ManufacturingFactoryEnv:
    """Core simulation engine for the Space Manufacturing RL environment."""

    def __init__(
        self,
        num_platforms: int,
        max_steps: int,
        seed: int,
        pending_orders: List[PendingOrder],
        delivery_windows: List[DeliveryWindow],
        solar_zones: Optional[Dict[str, float]] = None,
    ) -> None:
        self.num_platforms = num_platforms
        self.max_steps = max_steps
        self.seed = seed
        self._initial_orders = list(pending_orders)
        self._initial_windows = list(delivery_windows)
        self._solar_zones: Dict[str, float] = solar_zones or {"zone_a": 0.8}

        # mutable state (initialised in reset)
        self.platforms: List[PlatformState] = []
        self.pending_orders: List[PendingOrder] = []
        self.delivery_windows: List[DeliveryWindow] = []
        self.step_count: int = 0
        self.done: bool = False
        self.total_reward: float = 0.0
        self.episode_id: str = ""
        self.metrics: Dict[str, Any] = {}
        self._rng = random.Random(seed)

    # ── Public API ─────────────────────────────────────────────────────────────

    def reset(self) -> ManufacturingObservation:
        self._rng = random.Random(self.seed)
        self.step_count = 0
        self.done = False
        self.total_reward = 0.0
        self.episode_id = str(uuid.uuid4())
        self.metrics = {
            "production_runs": 0,
            "assemblies_completed": 0,
            "deliveries_completed": 0,
            "on_time_deliveries": 0,
            "invalid_actions": 0,
            "total_steps": 0,
        }

        self.platforms = [self._init_platform(i) for i in range(self.num_platforms)]
        self.pending_orders = [o.model_copy() for o in self._initial_orders]
        self.delivery_windows = [w.model_copy() for w in self._initial_windows]

        return self._build_observation(reward=0.0)

    def step(
        self, action: ManufacturingAction
    ) -> Tuple[ManufacturingObservation, ManufacturingReward, bool, Dict[str, Any]]:
        if self.done:
            obs = self._build_observation(reward=0.0)
            return obs, ManufacturingReward(value=0.0), True, {}

        step_reward = 0.0
        components: Dict[str, float] = {}

        for pid, platform in enumerate(self.platforms):
            act = action.platform_actions.get(pid, "recharge")
            r, c = self._apply_action(platform, act)
            step_reward += r
            for k, v in c.items():
                components[k] = components.get(k, 0.0) + v

        # Passive material refill
        for p in self.platforms:
            p.material_stock = min(MATERIAL_MAX, p.material_stock + MATERIAL_REFILL_RATE)

        # Energy penalty for critically low
        for p in self.platforms:
            if p.energy < 10.0:
                step_reward -= 1.5
                components["energy_critical"] = components.get("energy_critical", 0.0) - 1.5

        # Expire overdue delivery windows (soft penalty)
        still_open: List[DeliveryWindow] = []
        for w in self.delivery_windows:
            if w.deadline >= self.step_count:
                still_open.append(w)
        self.delivery_windows = still_open

        self.step_count += 1
        self.metrics["total_steps"] = self.step_count
        self.total_reward += step_reward

        if self.step_count >= self.max_steps:
            self.done = True
        elif not self.pending_orders:
            self.done = True  # all orders delivered — stop early

        reward_obj = ManufacturingReward(value=step_reward, components=components)
        obs = self._build_observation(reward=step_reward)
        return obs, reward_obj, self.done, {"metrics": dict(self.metrics)}

    # ── Internal helpers ───────────────────────────────────────────────────────

    def _init_platform(self, pid: int) -> PlatformState:
        angle = (2 * math.pi / self.num_platforms) * pid
        pos = [round(math.cos(angle) * 400, 2), round(math.sin(angle) * 400, 2), 0.0]
        return PlatformState(
            id=pid,
            position=pos,
            energy=self._rng.uniform(60.0, 90.0),
            material_stock=self._rng.uniform(30.0, 60.0),
            component_stock=0.0,
            product_stock=0,
            last_action=None,
        )

    def _apply_action(
        self, platform: PlatformState, action: str
    ) -> Tuple[float, Dict[str, float]]:
        reward = 0.0
        components: Dict[str, float] = {}
        platform.last_action = action

        if action == "produce":
            if platform.material_stock < PRODUCE_MATERIAL_THRESHOLD:
                reward -= 2.0
                components["invalid_action"] = components.get("invalid_action", 0.0) - 2.0
                self.metrics["invalid_actions"] += 1
            else:
                platform.material_stock -= PRODUCE_MATERIAL_COST
                platform.component_stock = min(
                    COMPONENT_MAX, platform.component_stock + PRODUCE_COMPONENT_GAIN
                )
                platform.energy = max(0.0, platform.energy - PRODUCE_ENERGY_COST)
                reward += 3.0
                components["production"] = components.get("production", 0.0) + 3.0
                self.metrics["production_runs"] += 1

        elif action == "assemble":
            if platform.component_stock < ASSEMBLE_COMPONENT_THRESHOLD:
                reward -= 2.0
                components["invalid_action"] = components.get("invalid_action", 0.0) - 2.0
                self.metrics["invalid_actions"] += 1
            else:
                platform.component_stock -= ASSEMBLE_COMPONENT_COST
                platform.product_stock = min(PRODUCT_MAX, platform.product_stock + 1)
                platform.energy = max(0.0, platform.energy - ASSEMBLE_ENERGY_COST)
                reward += 5.0
                components["assembly"] = components.get("assembly", 0.0) + 5.0
                self.metrics["assemblies_completed"] += 1

        elif action == "deliver":
            if platform.product_stock < 1:
                reward -= 2.0
                components["invalid_action"] = components.get("invalid_action", 0.0) - 2.0
                self.metrics["invalid_actions"] += 1
            else:
                platform.product_stock -= 1
                platform.energy = max(0.0, platform.energy - DELIVER_ENERGY_COST)
                self.metrics["deliveries_completed"] += 1
                # Check if any window is still open → on-time bonus
                on_time = False
                for w in self.delivery_windows:
                    if w.deadline >= self.step_count:
                        on_time = True
                        self.metrics["on_time_deliveries"] += 1
                        self.delivery_windows.remove(w)
                        break
                if on_time:
                    reward += 8.0
                    components["on_time_delivery"] = (
                        components.get("on_time_delivery", 0.0) + 8.0
                    )
                else:
                    reward += 3.0
                    components["late_delivery"] = components.get("late_delivery", 0.0) + 3.0

                # Consume a pending order if matched
                if self.pending_orders:
                    self.pending_orders.pop(0)

        elif action == "recharge":
            zone = self._zone_for_platform(platform)
            irradiance = self._solar_zones.get(zone, 0.5)
            gain = RECHARGE_BASE * irradiance
            was_low = platform.energy < 25.0
            platform.energy = min(ENERGY_MAX, platform.energy + gain)
            if was_low:
                reward += 1.0
                components["timely_recharge"] = components.get("timely_recharge", 0.0) + 1.0

        else:
            # Unknown action treated as invalid
            reward -= 2.0
            components["invalid_action"] = components.get("invalid_action", 0.0) - 2.0
            self.metrics["invalid_actions"] += 1

        # Idle penalty — recharge when actionable work exists
        if action == "recharge" and platform.energy > 50.0 and (
            platform.material_stock >= PRODUCE_MATERIAL_THRESHOLD
            or platform.component_stock >= ASSEMBLE_COMPONENT_THRESHOLD
            or platform.product_stock > 0
        ):
            reward -= 0.3
            components["idle_penalty"] = components.get("idle_penalty", 0.0) - 0.3

        return reward, components

    def _zone_for_platform(self, platform: PlatformState) -> str:
        zones = list(self._solar_zones.keys())
        idx = platform.id % len(zones)
        return zones[idx]

    def _build_observation(self, reward: float) -> ManufacturingObservation:
        return ManufacturingObservation(
            platforms=[p.model_copy() for p in self.platforms],
            time_step=self.step_count,
            delivery_windows=list(self.delivery_windows),
            solar_conditions=dict(self._solar_zones),
            pending_orders=list(self.pending_orders),
            total_reward=self.total_reward,
            reward=reward,
            done=self.done,
            metadata=dict(self.metrics),
        )