forward_model.py

"""Shared demand model for agent-side planning and evaluation.

Provides the same demand math as BusinessEnvironment but parameterised by
the agent's *inferred* parameters rather than the hidden ground truth.
Used by both Greedy (1-step score) and MCTS (multi-step chained rollouts).

Key design property: deterministic (no Poisson sampling). Stochasticity in
planning comes from sampling different parameter sets across MCTS rollouts,
not from observation noise within a rollout (see plan decision D2).
"""

from dataclasses import dataclass, field
from typing import Optional

import numpy as np

from business_environment import Action, BusinessEnvironment, EnvConfig, Observation


# ---------------------------------------------------------------------------
# Observable state snapshot
# ---------------------------------------------------------------------------

@dataclass
class EnvState:
    """What the agent legitimately knows — no hidden parameters.

    Includes ``pending_orders`` so that MCTS rollouts can chain steps and
    observe the effect of a reorder arriving ``reorder_lag`` weeks later —
    exactly the phenomenon MCTS is supposed to exploit in Phase 3+.
    """
    week: int
    n_products: int
    unit_costs: np.ndarray
    season: Optional[int] = None
    inventory: Optional[np.ndarray] = None
    recent_observations: list[Observation] = field(default_factory=list)
    # List of (arrival_week, quantities) waiting to be delivered. Mirrors
    # BusinessEnvironment.pending_orders. Empty when inventory is disabled.
    pending_orders: list[tuple[int, np.ndarray]] = field(default_factory=list)


def extract_env_state(env: BusinessEnvironment,
                      recent_obs: list[Observation]) -> EnvState:
    """Build an EnvState snapshot from the environment's observable state."""
    cfg = env.config
    pending: list[tuple[int, np.ndarray]] = (
        [(w, q.copy()) for w, q in env.pending_orders]
        if cfg.enable_inventory else []
    )
    return EnvState(
        week=env.week,
        n_products=cfg.n_products,
        unit_costs=env.unit_costs.copy(),
        season=(env.week // 13) % 4 if cfg.enable_seasonality else None,
        inventory=env.inventory.copy() if cfg.enable_inventory else None,
        recent_observations=list(recent_obs),
        pending_orders=pending,
    )


# ---------------------------------------------------------------------------
# Demand computation
# ---------------------------------------------------------------------------

def compute_expected_demand(
    params: dict[str, np.ndarray],
    action: Action,
    config: EnvConfig,
    week: int,
    unit_costs: np.ndarray,
) -> np.ndarray:
    """Deterministic expected demand from inferred parameters.

    Same multiplicative constant-elasticity formula as
    BusinessEnvironment._compute_expected_demand:

      demand_i = base_demand_i
        * (reference_price_i / price_i) ^ elasticity_i
        * (1 + marketing_response_i * log(1 + marketing_i))      # Phase 2+
        * seasonal_multiplier[season, i]                          # Phase 3+
        * exp(cross_elasticities[i] @ log(prices / ref_prices))   # Phase 4

    Args:
        params: Inferred demand parameters. Required keys per phase:
            All:     "base_demands" (n,), "elasticities" (n,)
            Phase 2: + "marketing_responses" (n,)
            Phase 3: + "seasonal_multipliers" (4, n)
            Phase 4: + "cross_elasticities" (n, n)
        action: Prices and marketing spend for this step.
        config: Feature toggles and product count.
        week: Current week (for seasonality calculation).
        unit_costs: Known per-product costs (reference_price = cost * 1.5).

    Returns:
        Expected demand per product, shape (n_products,).
    """
    prices = np.asarray(action.prices, dtype=float)
    reference_prices = unit_costs * 1.5

    price_ratio = reference_prices / prices
    expected = params["base_demands"] * np.power(price_ratio, params["elasticities"])

    if config.enable_marketing:
        mkt = (np.asarray(action.marketing, dtype=float)
               if action.marketing is not None
               else np.zeros(config.n_products))
        expected = expected * (1.0 + params["marketing_responses"] * np.log1p(mkt))

    if config.enable_seasonality:
        season = (week // 13) % 4
        expected = expected * params["seasonal_multipliers"][season]

    if config.enable_cross_effects:
        log_price_ratio = np.log(prices / reference_prices)
        cross_shift = params["cross_elasticities"] @ log_price_ratio
        expected = expected * np.exp(cross_shift)

    return np.maximum(expected, 0.0)


# ---------------------------------------------------------------------------
# One-step simulation for planning
# ---------------------------------------------------------------------------

def simulate_step(
    params: dict[str, np.ndarray],
    env_state: EnvState,
    action: Action,
    config: EnvConfig,
) -> tuple[np.ndarray, float, EnvState]:
    """Simulate one planning step. Returns (expected_sold, profit, next_state).

    Same order of operations as ``BusinessEnvironment.step`` so MCTS can
    chain rollouts with correct inventory / reorder dynamics:

      1. Compute expected demand at the current week.
      2. Receive pending reorders whose arrival_week ≤ current week.
      3. Sell up to inventory (any excess demand is a stockout).
      4. Queue the new reorder (if any) for arrival at week + reorder_lag.

    Greedy can ignore ``next_state``; MCTS uses it to walk the tree.
    Deterministic: no Poisson sampling inside — see plan D2.

    Returns:
        expected_sold:  shape (n_products,), sales after inventory cap.
        profit:         scalar, revenue - COGS - marketing - holding - stockout.
        next_state:     EnvState for the next step (week advanced, inventory
                        and pending_orders updated). Unchanged fields:
                        n_products, unit_costs, recent_observations.
    """
    expected_demand = compute_expected_demand(
        params, action, config, env_state.week, env_state.unit_costs,
    )
    prices = np.asarray(action.prices, dtype=float)

    # -- Inventory, arrivals, sales, remaining stock --
    if config.enable_inventory and env_state.inventory is not None:
        inventory, pending = _receive_arrivals(
            env_state.inventory, env_state.pending_orders, env_state.week,
        )
        expected_sold = np.minimum(expected_demand, inventory)
        stockouts = expected_demand - expected_sold
        remaining_inventory = inventory - expected_sold
        if action.reorder_quantities is not None:
            pending = pending + [(
                env_state.week + config.reorder_lag,
                np.asarray(action.reorder_quantities, dtype=float),
            )]
    else:
        expected_sold = expected_demand
        stockouts = np.zeros(env_state.n_products)
        remaining_inventory = None
        pending = env_state.pending_orders

    # -- Profit = revenue − COGS − marketing − holding − stockout penalty --
    revenue = float(np.sum(expected_sold * prices))
    costs = float(np.sum(expected_sold * env_state.unit_costs))
    if config.enable_marketing and action.marketing is not None:
        costs += float(np.sum(action.marketing))
    if config.enable_inventory and remaining_inventory is not None:
        costs += float(np.sum(remaining_inventory * config.holding_cost))
        costs += float(np.sum(stockouts * config.stockout_penalty))

    next_week = env_state.week + 1
    next_state = EnvState(
        week=next_week,
        n_products=env_state.n_products,
        unit_costs=env_state.unit_costs,
        season=(next_week // 13) % 4 if config.enable_seasonality else None,
        inventory=remaining_inventory,
        recent_observations=env_state.recent_observations,
        pending_orders=pending,
    )
    return expected_sold, revenue - costs, next_state


def _receive_arrivals(
    inventory: np.ndarray,
    pending_orders: list[tuple[int, np.ndarray]],
    current_week: int,
) -> tuple[np.ndarray, list[tuple[int, np.ndarray]]]:
    """Add arrivals whose due week has passed; return (new_inventory, remaining)."""
    new_inventory = inventory.copy()
    remaining: list[tuple[int, np.ndarray]] = []
    for arrival_week, qty in pending_orders:
        if arrival_week <= current_week:
            new_inventory = new_inventory + qty
        else:
            remaining.append((arrival_week, qty))
    return new_inventory, remaining