scheduling_rl.py

"""
scheduling_rl.py — PPO-based replacement scheduling for the HP Metal Jet S100.

Environment
-----------
    State  : [h(9), X_t(9), budget_remaining/W(1), t_hours(1)]  →  R^20
    Action : Discrete(512) — one of 2^9 joint component-replacement combinations
    Reward : +1.0 per survived hour (dt=1); 0.0 on the terminal step
    Done   : any h_i < HEALTH_THRESHOLD  OR  replacement cost exceeds budget

Algorithm
---------
    PPO (Stable-Baselines3) with a 512-output Softmax actor.
    The Softmax over the full 2^9 action space captures all joint replacement
    correlations without requiring an autoregressive architecture.

    Actor:  Linear(20→64) → Tanh → Linear(64→64) → Tanh → Linear(64→512)
    Critic: Linear(20→64) → Tanh → Linear(64→64) → Tanh → Linear(64→1)

    SB3 default orthogonal init + 0.01 final-layer scale → initial policy
    is near-uniform over all 512 actions (max-entropy start).

Component costs (€)
--------------------
    Sourced from analogous industrial parts; HP does not publish official prices.
    Sources listed in DEFAULT_COSTS docstring below.
"""

from __future__ import annotations

import numpy as np
import torch as th
import gymnasium as gym
from gymnasium import spaces
from stable_baselines3 import PPO

from model import DegradationModel

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

COMPONENT_NAMES: list[str] = [
    "recoater_blade",       # 0
    "nozzle_plate",         # 1
    "heating_elements",     # 2
    "temperature_sensors",  # 3
    "insulation_panels",    # 4
    "firing_resistors",     # 5
    "cleaning_interface",   # 6
    "recoater_motor",       # 7
    "linear_rail",          # 8
]

# Replacement costs in euros per component.
# Sources:
#   recoater_blade      — WINK3D EOS M290 ceramic blade ~$160, scaled to HP Metal Jet tier
#                         https://winking3d.com/product/ceramic-recoater-blade-3/
#   nozzle_plate        — 6× HP TIJ printheads; HP industrial OEM pricing estimate
#                         https://3dprintingindustry.com/news/hp-launches-new-metal-jet-s100...
#   heating_elements    — Industrial resistive heating elements for 180°C chamber
#                         https://www.sentrotech.com/heating-elements/
#   temperature_sensors — Industrial PT100/thermocouple bundle (Omega, multiple sensors)
#                         https://www.omega.co.uk/pptst/T3PROBES.html
#   insulation_panels   — Ceramic fibre panel section for printer-sized chamber
#                         https://www.sentrotech.com/ceramic-fiber-insulation/
#   firing_resistors    — Printhead array embedded resistors, proprietary HP part
#   cleaning_interface  — Wiper blade + solvent delivery module
#                         https://digiprint-usa.com/blogs/printhead-guides-tips-digiprint-usa/...
#   recoater_motor      — 500 W–1 kW industrial servo motor
#                         https://teknic.com/products/clearpath-brushless-dc-servo-motors/
#   linear_rail         — THK/HIWIN precision guide rail assembly
DEFAULT_COSTS = np.array([
     350.0,   # recoater_blade
    2500.0,   # nozzle_plate
    1500.0,   # heating_elements
     200.0,   # temperature_sensors
     800.0,   # insulation_panels
     400.0,   # firing_resistors
     600.0,   # cleaning_interface
     900.0,   # recoater_motor
     600.0,   # linear_rail
], dtype=np.float64)

HEALTH_THRESHOLD: float = 0.1   # printer fails when any component drops below this
N_COMPONENTS: int = 9
N_ACTIONS: int = 2 ** N_COMPONENTS  # 512

# Precompute action → binary replacement vector table once at import time.
# _ACTION_TABLE[i] is a (9,) float32 array: bit j = 1 means replace component j.
_ACTION_TABLE: np.ndarray = np.array(
    [[int(b) for b in format(i, f"0{N_COMPONENTS}b")] for i in range(N_ACTIONS)],
    dtype=np.float32,
)


# ---------------------------------------------------------------------------
# Environment
# ---------------------------------------------------------------------------

class PrinterEnv(gym.Env):
    """
    Gymnasium environment wrapping DegradationModel for replacement scheduling.

    Parameters
    ----------
    model       : Fitted DegradationModel with N=9 components and C inputs.
    X_series    : (T_max, C) array of predicted operating-condition vectors,
                  one row per simulated hour. If the episode outlasts X_series
                  the last row is repeated.
    W           : Total budget in euros available for the episode.
    costs       : (9,) replacement cost per component in euros.
                  Defaults to DEFAULT_COSTS when None.
    dt          : Simulation time-step in hours (default 1.0).
    stochastic  : Include Poisson shock terms from model.lambda_rates (default True).
    seed        : RNG seed for reproducibility.
    """

    metadata: dict = {"render_modes": []}

    def __init__(
        self,
        model: DegradationModel,
        X_series: np.ndarray,
        W: float,
        costs: np.ndarray | None = None,
        dt: float = 1.0,
        stochastic: bool = True,
        seed: int | None = None,
    ) -> None:
        """
        Initialise the environment and validate inputs.

        Builds the observation and action spaces, stores a reference to the
        fitted DegradationModel, and sets the episode state to its initial
        values (all components fully healthy, full budget, t=0).

        Raises ValueError if model.N != 9.
        """
        super().__init__()

        if model.N != N_COMPONENTS:
            raise ValueError(f"DegradationModel must have N={N_COMPONENTS}, got {model.N}")

        self.model = model
        self.X_series = np.asarray(X_series, dtype=np.float64)  # (T_max, C)
        self.W = float(W)
        self.costs = (
            np.asarray(costs, dtype=np.float64) if costs is not None else DEFAULT_COSTS
        )
        self.dt = float(dt)
        self.stochastic = stochastic
        self._rng = np.random.default_rng(seed)

        # Observation: h(9) + X_t(C) + budget/W(1) + t_hours(1)
        obs_dim = N_COMPONENTS + model.C + 2
        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf, shape=(obs_dim,), dtype=np.float32
        )
        self.action_space = spaces.Discrete(N_ACTIONS)

        # Episode state — initialised properly in reset()
        self._h = np.ones(N_COMPONENTS, dtype=np.float64)
        self._t: float = 0.0
        self._budget: float = self.W
        self._step_idx: int = 0

    # ------------------------------------------------------------------
    # Core Gymnasium interface
    # ------------------------------------------------------------------

    def reset(
        self,
        *,
        seed: int | None = None,
        options: dict | None = None,
    ) -> tuple[np.ndarray, dict]:
        """
        Reset the environment to the start of a new episode.

        All components are restored to full health (h=1), the budget is
        refilled to W, and the simulation clock is set to t=0. If a seed
        is supplied the internal RNG is re-seeded for reproducibility.

        Returns
        -------
        obs : (obs_dim,) float32 observation vector.
        info : empty dict (required by the Gymnasium API).
        """
        if seed is not None:
            self._rng = np.random.default_rng(seed)

        self._h = np.ones(N_COMPONENTS, dtype=np.float64)
        self._t = 0.0
        self._budget = self.W
        self._step_idx = 0

        return self._obs(), {}

    def step(
        self, action: int
    ) -> tuple[np.ndarray, float, bool, bool, dict]:
        """
        Advance the simulation by one hour and apply the chosen replacements.

        The action integer is decoded to a 9-bit binary replacement vector via
        _ACTION_TABLE. Components whose bit is 1 are replaced (health reset to
        1.0) before the ODE step is taken, so new components benefit from the
        full degradation rate of a healthy part.

        Budget is checked before applying replacements. If the total cost of
        the selected action exceeds the remaining budget the episode terminates
        immediately with reward 0 and no state change.

        After replacements the health vector is updated with one Euler step:

            h_{t+1} = h_t + f(h_t, X_t) * dt  [- Q * H  if stochastic]

        where Q ~ N(0,1)² and H ~ Poisson(λ_i * dt) are independent per
        component (matching the stochastic extension in model.py §5).

        Parameters
        ----------
        action : int in [0, 511].

        Returns
        -------
        obs         : (obs_dim,) float32 observation after the step.
        reward      : dt (1.0) if the printer survived; 0.0 on failure.
        terminated  : True when any h_i < HEALTH_THRESHOLD or budget exceeded.
        truncated   : always False (no fixed time limit).
        info        : dict with keys t_hours, budget_remaining, min_health
                      (or termination on budget failure).
        """
        bits = _ACTION_TABLE[int(action)]          # (9,) float32 binary vector
        replacement_cost = float(np.dot(bits, self.costs))

        # Hard budget constraint: action is unaffordable → episode ends
        if replacement_cost > self._budget:
            return self._obs(), 0.0, True, False, {
                "termination": "budget_exceeded",
                "t_hours": self._t,
                "budget_remaining": self._budget,
            }

        # Apply replacements: reset selected components to full health
        self._budget -= replacement_cost
        self._h[bits.astype(bool)] = 1.0

        # Euler step: dh = f(h, X_t) * dt  [+ optional stochastic shock]
        X_t = self._current_X()
        dh = self.model.f(self._h, X_t) * self.dt

        if self.stochastic:
            Q = self._rng.standard_normal(N_COMPONENTS) ** 2          # N(0,1)² — event intensity
            H = self._rng.poisson(self.model.lambda_rates * self.dt)  # Poisson(λ_i · dt)
            dh -= Q * H

        self._h = self._h + dh
        self._t += self.dt
        self._step_idx += 1

        # Failure: any component below threshold
        failed = bool(np.any(self._h < HEALTH_THRESHOLD))
        reward = 0.0 if failed else self.dt

        return self._obs(), reward, failed, False, {
            "t_hours": self._t,
            "budget_remaining": self._budget,
            "min_health": float(self._h.min()),
        }

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _obs(self) -> np.ndarray:
        """
        Build the current observation vector.

        Layout (obs_dim = 9 + C + 2):
            h[0:9]              — component health values (float, may be negative
                                  after a stochastic shock representing catastrophic
                                  failure before the done flag is raised)
            X_t[9:9+C]         — current operating-condition inputs from X_series
            budget/W [9+C]     — remaining budget as a fraction of the initial
                                  budget W; always in [0, 1] during a valid episode
            t_hours  [9+C+1]   — elapsed simulation time in hours

        Returns float32 to match observation_space.dtype.
        """
        X_t = self._current_X()
        return np.concatenate([
            self._h,                        # 9  — component healths (may be < 0 stochastic)
            X_t,                            # C  — current operating conditions
            [self._budget / self.W],        # 1  — remaining budget fraction ∈ [0, 1]
            [self._t],                      # 1  — elapsed hours
        ]).astype(np.float32)

    def _current_X(self) -> np.ndarray:
        """
        Return the operating-condition vector for the current simulation step.

        Clamps the index to the last row of X_series if the episode outlasts
        the provided forecast horizon, so the environment never raises an
        IndexError regardless of episode length.
        """
        idx = min(self._step_idx, len(self.X_series) - 1)
        return self.X_series[idx]


# ---------------------------------------------------------------------------
# Training
# ---------------------------------------------------------------------------

def train(
    env: PrinterEnv,
    *,
    total_timesteps: int = 1_000_000,
    save_path: str = "scheduler_ppo",
    learning_rate: float = 3e-4,
    n_steps: int = 2048,
    batch_size: int = 64,
    n_epochs: int = 10,
    gamma: float = 0.99,
    ent_coef: float = 0.01,
    verbose: int = 1,
) -> PPO:
    """
    Train a PPO agent on PrinterEnv and save the result to disk.

    Architecture
    ------------
    Both the actor and critic share the same MLP topology
    (two hidden layers of 64 units with Tanh activations) but have
    separate weights, as is standard in Actor-Critic methods:

        Actor  (π_θ): obs(20) → 64 → 64 → logits(512) → Softmax
        Critic (V_ϕ): obs(20) → 64 → 64 → scalar(1)

    The 512 Softmax output over the full 2^9 action space implicitly models
    the joint probability distribution over all replacement combinations,
    capturing correlations (e.g. replacing component 0 should reduce the
    probability of also replacing component 1 if they interact positively)
    without requiring an autoregressive sampling pass.

    Initialisation
    --------------
    SB3 applies orthogonal initialisation to all layers (scale √2 for hidden,
    scale 0.01 for the final policy layer). With scale 0.01 the logits start
    near zero, so Softmax(~0) ≈ 1/512 — effectively maximum entropy over the
    action space, ensuring the agent explores broadly before committing.

    Hyperparameters
    ---------------
    learning_rate : Adam step size (3e-4 is a reliable PPO default).
    n_steps       : Rollout length before each PPO update (2048 steps ≈
                    one to several full episodes depending on episode length).
    batch_size    : Mini-batch size for gradient updates.
    n_epochs      : Number of gradient passes over each collected rollout.
    gamma         : Discount factor; 0.99 weights future rewards heavily,
                    encouraging the agent to maximise long-term lifespan.
    ent_coef      : Entropy bonus coefficient; keeps action probabilities from
                    collapsing to a single action too early in training.

    Parameters
    ----------
    env             : Configured PrinterEnv instance.
    total_timesteps : Total environment steps to train for.
    save_path       : File path (without .zip) for the saved model.
    learning_rate   : Adam learning rate.
    n_steps         : Rollout buffer size (steps per PPO update).
    batch_size      : SGD mini-batch size.
    n_epochs        : PPO epochs per rollout.
    gamma           : Discount factor γ.
    ent_coef        : Entropy regularisation coefficient.
    verbose         : SB3 verbosity level (0=silent, 1=info, 2=debug).

    Returns
    -------
    PPO : The trained Stable-Baselines3 PPO model, ready for evaluate().
    """
    policy_kwargs = dict(
        net_arch=dict(pi=[64, 64], vf=[64, 64]),
        activation_fn=th.nn.Tanh,
    )

    ppo = PPO(
        policy="MlpPolicy",
        env=env,
        learning_rate=learning_rate,
        n_steps=n_steps,
        batch_size=batch_size,
        n_epochs=n_epochs,
        gamma=gamma,
        ent_coef=ent_coef,
        policy_kwargs=policy_kwargs,
        verbose=verbose,
    )

    ppo.learn(total_timesteps=total_timesteps)
    ppo.save(save_path)
    if verbose:
        print(f"[train] Model saved to {save_path}.zip")
    return ppo


# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------

def evaluate(
    ppo: PPO,
    env: PrinterEnv,
    n_episodes: int = 10,
) -> dict:
    """
    Run the trained policy deterministically and return summary statistics.

    Each episode starts from a fresh reset() call. The policy is queried with
    deterministic=True, meaning the action with the highest probability under
    the current Softmax distribution is always selected (no sampling noise).
    This gives a reproducible, greedy estimate of the policy's performance.

    Parameters
    ----------
    ppo         : Trained PPO model returned by train().
    env         : PrinterEnv instance (can be the same env used for training).
    n_episodes  : Number of evaluation episodes to average over.

    Returns
    -------
    dict with keys:
        mean_hours        — average hours survived across all episodes.
        std_hours         — standard deviation of hours survived.
        mean_replacements — average total number of component replacements
                            performed per episode (summed over all steps).
        mean_budget_spent — average euros spent on replacements per episode.
        episodes          — list of per-episode dicts, each containing:
                            episode, hours_survived, budget_spent, replacements.
    """
    hours, n_replacements, budget_spent = [], [], []
    episodes = []

    for ep in range(n_episodes):
        obs, _ = env.reset()
        done = False
        ep_replacements = 0

        while not done:
            action, _ = ppo.predict(obs, deterministic=True)
            obs, _, terminated, truncated, _ = env.step(int(action))
            done = terminated or truncated
            ep_replacements += int(_ACTION_TABLE[int(action)].sum())

        hours.append(env._t)
        n_replacements.append(ep_replacements)
        budget_spent.append(env.W - env._budget)
        episodes.append({
            "episode":        ep,
            "hours_survived": env._t,
            "budget_spent":   env.W - env._budget,
            "replacements":   ep_replacements,
        })

    return {
        "mean_hours":        float(np.mean(hours)),
        "std_hours":         float(np.std(hours)),
        "mean_replacements": float(np.mean(n_replacements)),
        "mean_budget_spent": float(np.mean(budget_spent)),
        "episodes":          episodes,
    }


# ---------------------------------------------------------------------------
# Sanity demo
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    from process_inputs import process_inputs

    deg_model = DegradationModel.load("model.npz")
    print(f"Loaded model.npz  N={deg_model.N}  C={deg_model.C}")

    # Representative operating-condition scenarios in natural units (matches phase2.py SCENARIOS).
    # Each row: [ambient_temp_c, chamber_temp_c, humidity_pct, contamination_aqi,
    #            print_hours, build_volume_cm3, recoating_speed_mm_s, recoating_cycles, maintenance_level]
    _SCENARIO_CONDITIONS = np.array([
        [ 22.0, 180.0, 40.0,   20.0,  56.0,  4500.0, 100.0,  5000.0, 0.10],  # nominal
        [ 22.0, 180.0, 45.0,  150.0,  56.0,  4500.0, 100.0,  5000.0, 0.20],  # high contamination
        [ 28.0, 185.0, 50.0,   80.0, 168.0, 13500.0, 120.0, 15000.0, 0.20],  # 24/7 heavy use
        [ 20.0, 175.0, 30.0,    5.0, 300.0,  3000.0,  90.0,  3500.0, 0.05],  # optimal lab
        [ 40.0, 200.0, 55.0,   60.0,  56.0,  4500.0, 100.0,  5000.0, 0.20],  # hot environment
        [ 22.0, 180.0, 85.0,   30.0,  56.0,  4500.0, 100.0,  5000.0, 0.15],  # high humidity
        [ 25.0, 182.0, 45.0,   40.0,  56.0,  4500.0, 100.0,  5000.0, 0.90],  # neglected maintenance
        [ 38.0, 195.0, 20.0,  200.0,  80.0,  6500.0, 110.0,  7000.0, 0.40],  # desert factory
    ], dtype=np.float64)

    # Build an 8 000-step X_series by cycling through scenarios with small noise
    rng = np.random.default_rng(42)
    T = 8_000
    base_rows = _SCENARIO_CONDITIONS[np.arange(T) % len(_SCENARIO_CONDITIONS)]
    noise = rng.normal(0.0, 0.02, size=base_rows.shape) * base_rows  # 2% relative noise
    X_natural = np.clip(base_rows + noise, 0.0, None)
    X_series = np.stack([process_inputs(row) for row in X_natural])  # (T, C) normalised

    env = PrinterEnv(
        model=deg_model,
        X_series=X_series,
        W=10_000.0,
        dt=1.0,
        stochastic=True,
        seed=42,
    )

    print("Observation space:", env.observation_space.shape)
    print("Action space:     ", env.action_space.n, "discrete actions")
    print()

    # --- Random-policy baseline ---
    obs, _ = env.reset(seed=0)
    done = False
    while not done:
        action = env.action_space.sample()
        obs, _, terminated, truncated, info = env.step(action)
        done = terminated or truncated
    print(f"Random policy  → {env._t:.0f} h survived | €{env.W - env._budget:.0f} spent")

    # --- PPO training (short demo: 200k steps) ---
    print("\nTraining PPO (200 000 timesteps) ...")
    trained = train(
        env,
        total_timesteps=200_000,
        save_path="scheduler_ppo",
        verbose=1,
    )

    # --- Evaluation ---
    results = evaluate(trained, env, n_episodes=10)
    print(f"\n{'─'*45}")
    print(f"  Mean hours survived : {results['mean_hours']:>8.1f} h")
    print(f"  Std  hours          : {results['std_hours']:>8.1f} h")
    print(f"  Mean replacements   : {results['mean_replacements']:>8.1f}")
    print(f"  Mean budget spent   : €{results['mean_budget_spent']:>7.0f}")
    print(f"{'─'*45}")