env/core.py

import random
import uuid
from openenv.core.env_server import Environment

from env.schemas import (
    SupplyChainAction,
    SupplyChainObservation,
    SupplyChainState,
    ProductStatus,
)

# ---------------------------------------------------------------------------
# Task configuration constants
# ---------------------------------------------------------------------------
_PRODUCT_CATALOG = {
    "SKU-LAPTOP": {
        "margin": 400.0,
        "holding_cost": 2.0,
        "penalty": 100.0,
        "order_cost": 800.0,
    }
}

_TASK_CONFIGS = {
    "task_01_easy": {
        "cash_balance": 30000.0,
        "initial_inventory": {"SKU-LAPTOP": 50},
        "market_signal": "Demand is stable at exactly 10 units per day.",
    },
    "task_02_medium": {
        "cash_balance": 50000.0,
        "initial_inventory": {"SKU-LAPTOP": 20},
        "market_signal": (
            "WARNING: Black Friday sale begins on Day 10. "
            "Demand will spike from 10 to 40 units per day."
        ),
    },
    "task_03_hard": {
        "cash_balance": 30000.0,
        "initial_inventory": {"SKU-LAPTOP": 30},
        "market_signal": (
            "WARNING: Global shipping crisis. Standard 2-day shipping is delayed to 4 days. "
            "Expedited 1-day shipping costs $100 extra per unit."
        ),
    },
}

# Shipping / cost constants
STANDARD_LEAD_TIME = 2
EXPEDITED_LEAD_TIME = 1
CRISIS_LEAD_TIME = 4
EXPEDITE_SURCHARGE = 100  # extra cost per unit for expedited shipping


class SupplyChainEnv(Environment):
    """
    OpenEnv-compliant supply chain inventory management environment.

    Inherits from openenv.core.env_server.Environment.
    The agent acts as a Supply Chain Manager, placing daily purchase orders
    to maximize profit across a 30-day simulation.
    """

    SUPPORTS_CONCURRENT_SESSIONS = False

    def __init__(self):
        super().__init__()
        self.current_task_id: str = ""
        self.max_days: int = 30
        self.current_day: int = 0
        self.cash_balance: float = 0.0
        self.inventory: dict = {}
        self.shipment_pipeline: dict = {}
        self.history: dict = {"sales_yesterday": {}, "lost_sales_yesterday": {}}
        self.catalog: dict = {}
        self.market_signal: str = ""
        self.total_reward: float = 0.0  # cumulative sum of per-step newsvendor rewards
        self.last_accepted_qty: int = 0
        self.last_rejected_qty: int = 0
        # OpenEnv State container (tracks episode metadata)
        self._state = SupplyChainState(
            episode_id=str(uuid.uuid4()),
            step_count=0,
        )

    def reset(self, task_id: str = "task_01_easy", seed: int | None = None) -> SupplyChainObservation:
        """Reset the environment and load a specific task. Returns initial observation.

        Args:
            task_id: One of 'task_01_easy', 'task_02_medium', 'task_03_hard'.
            seed: Optional RNG seed for reproducible demand simulation.
        """
        if task_id not in _TASK_CONFIGS:
            raise ValueError(f"Unknown task_id '{task_id}'. Valid: {sorted(_TASK_CONFIGS)}")

        if seed is not None:
            random.seed(seed)

        cfg = _TASK_CONFIGS[task_id]
        self.current_task_id = task_id
        self.current_day = 1
        self.total_reward = 0.0
        self.last_accepted_qty = 0
        self.last_rejected_qty = 0

        self.catalog = {pid: dict(props) for pid, props in _PRODUCT_CATALOG.items()}
        self.cash_balance = cfg["cash_balance"]
        self.inventory = dict(cfg["initial_inventory"])
        self.shipment_pipeline = {pid: {1: 0, 2: 0, 3: 0, 4: 0} for pid in self.catalog}
        self.history = {
            "sales_yesterday": {pid: 0 for pid in self.catalog},
            "lost_sales_yesterday": {pid: 0 for pid in self.catalog},
        }
        self.market_signal = cfg["market_signal"]

        # Reset the OpenEnv State container
        self._state = SupplyChainState(
            episode_id=str(uuid.uuid4()),
            step_count=0,
            current_task_id=task_id,
        )

        return self._build_observation(reward=0.0, done=False)

    def step(self, action: SupplyChainAction) -> SupplyChainObservation:
        """Execute one day of the simulation. Returns observation with reward and done."""
        # 1. Process Arrivals
        for pid in self.catalog.keys():
            arriving_today = self.shipment_pipeline[pid].get(1, 0)
            self.inventory[pid] += arriving_today
            new_pipeline = {}
            for day, qty in self.shipment_pipeline[pid].items():
                if day > 1:
                    new_pipeline[day - 1] = qty
            self.shipment_pipeline[pid] = new_pipeline

        # 2. Process New Orders (track accepted vs rejected)
        self.last_accepted_qty = 0
        self.last_rejected_qty = 0
        for order in action.orders:
            pid = order.product_id
            if pid not in self.catalog or order.quantity <= 0:
                continue

            cost = order.quantity * self.catalog[pid]["order_cost"]

            # 4-day standard shipping during crisis, 2-day otherwise
            if self.current_task_id == "task_03_hard" and not order.expedite_shipping:
                delivery_days = CRISIS_LEAD_TIME
            else:
                delivery_days = EXPEDITED_LEAD_TIME if order.expedite_shipping else STANDARD_LEAD_TIME

            if order.expedite_shipping:
                cost += order.quantity * EXPEDITE_SURCHARGE

            if self.cash_balance >= cost:
                self.cash_balance -= cost
                current_queued = self.shipment_pipeline[pid].get(delivery_days, 0)
                self.shipment_pipeline[pid][delivery_days] = current_queued + order.quantity
                self.last_accepted_qty += order.quantity
            else:
                self.last_rejected_qty += order.quantity

        # 3. Simulate Daily Demand (stochastic)
        if self.current_task_id == "task_01_easy":
            demand_qty = random.randint(8, 12)
        elif self.current_task_id == "task_02_medium":
            demand_qty = random.randint(35, 45) if 10 <= self.current_day <= 17 else random.randint(8, 12)
        elif self.current_task_id == "task_03_hard":
            demand_qty = random.randint(5, 20)
        else:
            demand_qty = 10

        actual_demand = {"SKU-LAPTOP": demand_qty}

        # 4. Fulfill Demand, update cash, compute per-SKU newsvendor reward
        product_rewards = []
        for pid, demand in actual_demand.items():
            stock_before = self.inventory[pid]       # total stock before selling
            sold = min(stock_before, demand)
            missed = demand - sold
            remaining = stock_before - sold           # stock left after selling

            # Update inventory and history
            self.inventory[pid] = remaining
            self.history["sales_yesterday"][pid] = sold
            self.history["lost_sales_yesterday"][pid] = missed

            # Cash flow: revenue from sales (still drives budget for future orders)
            revenue = sold * (self.catalog[pid]["order_cost"] + self.catalog[pid]["margin"])
            self.cash_balance += revenue

            # ------------------------------------------------------------------
            # Newsvendor step reward — bounded [0, 1]
            #
            # Scenario 1 — Overage (all demand met, stock remains):
            #   reward = 1 - (remaining / stock_before)
            #   → rewards efficient use of inventory; perfect sell-through = 1.0
            #
            # Scenario 2 — Underage (stockout, unmet demand):
            #   reward = 1 - (missed / demand)
            #   → rewards high service level; zero missed orders = 1.0
            # ------------------------------------------------------------------
            if missed == 0:
                # Overage or perfect match
                product_reward = 1.0 - (remaining / stock_before) if stock_before > 0 else 1.0
            else:
                # Underage (demand > 0 guaranteed since missed > 0)
                product_reward = 1.0 - (missed / demand)

            product_rewards.append(product_reward)

        # Average across SKUs (scales cleanly to multi-product tasks)
        step_reward = sum(product_rewards) / len(product_rewards) if product_rewards else 0.0
        self.total_reward += step_reward

        # 5. Advance Time
        self.current_day += 1
        done = self.current_day > self.max_days

        # Sync OpenEnv state container
        self._state.step_count += 1
        self._state.current_task_id = self.current_task_id
        self._state.total_reward = self.total_reward
        self._state.grader_score = self.get_grader_score()

        return self._build_observation(reward=step_reward, done=done)

    @property
    def state(self) -> SupplyChainState:
        """OpenEnv-required property: returns current episode metadata."""
        self._state.current_task_id = self.current_task_id
        self._state.total_reward = self.total_reward
        self._state.grader_score = self.get_grader_score()
        return self._state

    def _build_observation(self, reward: float = 0.0, done: bool = False) -> SupplyChainObservation:
        """Build the current observation snapshot."""
        warehouse = []
        for pid in self.catalog.keys():
            status = ProductStatus(
                product_id=pid,
                current_stock=self.inventory.get(pid, 0),
                incoming_shipments=self.shipment_pipeline.get(pid, {}),
                sales_yesterday=self.history["sales_yesterday"].get(pid, 0),
                lost_sales_yesterday=self.history["lost_sales_yesterday"].get(pid, 0),
                holding_cost_per_unit=self.catalog[pid]["holding_cost"],
                stockout_penalty_per_unit=self.catalog[pid]["penalty"],
                margin_per_unit=self.catalog[pid]["margin"],
            )
            warehouse.append(status)

        return SupplyChainObservation(
            current_day=self.current_day,
            total_days=self.max_days,
            cash_balance=self.cash_balance,
            warehouse_status=warehouse,
            market_trend_signal=self.market_signal,
            last_order_accepted=self.last_accepted_qty,
            last_order_rejected=self.last_rejected_qty,
            reward=reward,
            done=done,
        )

    # ==========================================
    # OPENENV REQUIRED GRADER (0.0 to 1.0)
    # ==========================================
    def get_grader_score(self) -> float:
        """
        Returns a normalized score in [0.0, 1.0].
        total_reward = sum of per-step newsvendor rewards / max_days
        Each step reward is already in [0, 1], so dividing by max_days keeps
        the episode score in [0, 1] with partial-progress credit throughout.
        """
        return max(0.0, min(1.0, self.total_reward / self.max_days))