"""Store OpenEnv simulation.""" from __future__ import annotations import csv import math import random from copy import deepcopy from pathlib import Path from typing import Any from uuid import uuid4 try: from .._compat import Environment, State from ..models import ( AvailableActionSpec, MarketingAction, MarketOffer, OrderRequest, PendingOrder, PriceUpdate, SellerProfile, StockItem, StoreAction, StoreObservation, StoreSnapshot, ) except ImportError: # pragma: no cover - supports direct module imports from _compat import Environment, State from models import ( AvailableActionSpec, MarketingAction, MarketOffer, OrderRequest, PendingOrder, PriceUpdate, SellerProfile, StockItem, StoreAction, StoreObservation, StoreSnapshot, ) DATA_DIR = Path(__file__).resolve().parents[1] / "data" SELLER_PRICE_COLUMN = { "Seller A": "seller_a_price_inr", "Seller B": "seller_b_price_inr", "Seller C": "seller_c_price_inr", "Seller D": "seller_d_price_inr", } SELLER_ETD_COLUMN = { "Seller A": "seller_a_etd_days", "Seller B": "seller_b_etd_days", "Seller C": "seller_c_etd_days", "Seller D": "seller_d_etd_days", } NEW_PRODUCT_MARKUP = 1.18 class StoreEnvironment(Environment): """Simulation for a neighborhood grocery store.""" SUPPORTS_CONCURRENT_SESSIONS: bool = True EPISODE_HORIZON_DAYS = 45 COMMIT_WINDOW_DAYS = 3 REWARD_HORIZON_DAYS = 6 def __init__(self) -> None: self._state = State(episode_id=str(uuid4()), step_count=0) self._reset_count = 0 self._rng = random.Random(0) self._runtime = self._build_initial_runtime() def reset(self) -> StoreObservation: """Reset the environment back to the seed CSV state.""" self._state = State(episode_id=str(uuid4()), step_count=0) self._reset_count += 1 self._rng = random.Random(self._reset_count) self._runtime = self._build_initial_runtime() return self._build_observation( self._runtime, reward=0.0, done=False, summary=( "Day 1 store snapshot loaded. Submit one bundled action for the next " "3-day decision window; reward looks ahead across 6 simulated days." ), metadata={ "reset_count": self._reset_count, "simulation": { "commit_window_days": self.COMMIT_WINDOW_DAYS, "reward_horizon_days": self.REWARD_HORIZON_DAYS, }, }, ) def step(self, action: StoreAction) -> StoreObservation: # type: ignore[override] """Execute one bundled decision window.""" self._validate_action(action) live_runtime = self._runtime start_day = live_runtime["store"]["day"] starting_capital = self._capital(live_runtime) score_before = self._total_product_selling_score(live_runtime["stock"]) prepared_bundle = self._prepare_bundle(action, live_runtime) commit_branch = deepcopy(live_runtime) reward_branch = deepcopy(live_runtime) self._apply_bundle(commit_branch, prepared_bundle) self._apply_bundle(reward_branch, prepared_bundle) adjusted_satisfaction = live_runtime["store"]["customer_satisfaction"] + prepared_bundle[ "customer_satisfaction_delta" ] score_after = self._total_product_selling_score(commit_branch["stock"]) next_satisfaction = self._updated_customer_satisfaction( adjusted_satisfaction, score_before, score_after, ) commit_branch["store"]["customer_satisfaction"] = next_satisfaction reward_branch["store"]["customer_satisfaction"] = next_satisfaction commit_metrics = self._simulate_days(commit_branch, self.COMMIT_WINDOW_DAYS) reward_metrics = self._simulate_days(reward_branch, self.REWARD_HORIZON_DAYS) ending_capital = self._capital(reward_branch) difference = ( ending_capital - starting_capital - reward_metrics["recurring_costs_paid_inr"] ) normalized_difference = 0.0 if starting_capital <= 0 else difference / starting_capital reward = 0.5 + normalized_difference self._runtime = commit_branch self._state.step_count += 1 done_reason = "" if self._runtime["store"]["bank_balance_inr"] < 0: done_reason = "bank_balance_below_zero" elif self._runtime["store"]["day"] > self.EPISODE_HORIZON_DAYS: done_reason = "episode_horizon_reached" done = bool(done_reason) summary = ( f"Committed days {start_day}-{commit_branch['store']['day'] - 1}. " f"Reward forecast extends through day {reward_branch['store']['day'] - 1}." ) metadata = { "reward_breakdown": { "starting_capital_inr": self._round(starting_capital), "ending_capital_inr": self._round(ending_capital), "difference_inr": self._round(difference), "normalized_difference": self._round(normalized_difference, digits=6), "reward_horizon_days": self.REWARD_HORIZON_DAYS, "commit_window_days": self.COMMIT_WINDOW_DAYS, "recurring_costs_paid_inr": self._round(reward_metrics["recurring_costs_paid_inr"]), }, "commit_metrics": self._rounded_metrics(commit_metrics), "reward_metrics": self._rounded_metrics(reward_metrics), "done_reason": done_reason, } return self._build_observation( self._runtime, reward=reward, done=done, summary=summary, metadata=metadata, ) @property def state(self) -> State: """Return the OpenEnv transport state.""" return self._state def _build_initial_runtime(self) -> dict[str, Any]: store_rows = self._read_csv("store_state.csv") stock_rows = self._read_csv("stock.csv") pending_order_rows = self._read_csv("pending_orders.csv") seller_rows = self._read_csv("seller_profiles.csv") market_rows = self._read_csv("market_offers.csv") marketing_rows = self._read_csv("marketing_actions.csv") if len(store_rows) != 1: raise ValueError("store_state.csv must contain exactly one row") store_row = store_rows[0] pending_orders = [ { "order_id": row["order_id"], "product_name": row["product_name"], "quantity": int(row["quantity"]), "seller_name": row["seller_name"], "buy_price_inr": float(row["buy_price_inr"]), "days_remaining": int(row["days_remaining"]), "status": row["status"], "created_day": int(row["created_day"]), } for row in pending_order_rows ] return { "store": { "day": int(store_row["day"]), "current_daily_expense_inr": float(store_row["current_daily_expense_inr"]), "bank_balance_inr": float(store_row["bank_balance_inr"]), "customer_satisfaction": float(store_row["customer_satisfaction"]), "order_processing_capacity": float(store_row["order_processing_capacity"]), "cumulative_profit_inr": float(store_row["cumulative_profit_inr"]), }, "stock": [ { "product_name": row["product_name"], "unit": row["unit"], "buying_price_inr": float(row["buying_price_inr"]), "selling_price_inr": float(row["selling_price_inr"]), "demand_score": float(row["demand_score"]), "quantity_on_hand": float(row["quantity_on_hand"]), "stock_value_inr": float(row["stock_value_inr"]), } for row in stock_rows ], "pending_orders": pending_orders, "market_offers": [ { "product_name": row["product_name"], "category": row["category"], "unit": row["unit"], "demand_score": float(row["demand_score"]), "seller_a_price_inr": int(row["seller_a_price_inr"]), "seller_a_etd_days": int(row["seller_a_etd_days"]), "seller_b_price_inr": int(row["seller_b_price_inr"]), "seller_b_etd_days": int(row["seller_b_etd_days"]), "seller_c_price_inr": int(row["seller_c_price_inr"]), "seller_c_etd_days": int(row["seller_c_etd_days"]), "seller_d_price_inr": int(row["seller_d_price_inr"]), "seller_d_etd_days": int(row["seller_d_etd_days"]), } for row in market_rows ], "seller_profiles": [ { "seller_name": row["seller_name"], "timeline_behavior": row["timeline_behavior"], "reliability_note": row["reliability_note"], "default_etd_days": int(row["default_etd_days"]), } for row in seller_rows ], "marketing_actions_catalog": [ { "action_id": int(row["action_id"]), "description": row["description"], "customer_satisfaction_delta": int(row["customer_satisfaction_delta"]), "cost_inr": int(row["cost_inr"]), "recurring_cost_inr": int(row["recurring_cost_inr"]), "order_processing_capacity_delta": int(row["order_processing_capacity_delta"]), } for row in marketing_rows ], "active_action_ids": set(), "next_order_id": len(pending_orders) + 1, } def _read_csv(self, name: str) -> list[dict[str, str]]: with (DATA_DIR / name).open(newline="", encoding="utf-8") as handle: return list(csv.DictReader(handle)) def _validate_action(self, action: StoreAction) -> None: market_by_product = { row["product_name"]: row for row in self._runtime["market_offers"] } stock_by_product = { row["product_name"]: row for row in self._runtime["stock"] } marketing_ids = { row["action_id"] for row in self._runtime["marketing_actions_catalog"] } seen_price_products: set[str] = set() for update in action.prices or []: if update.product not in stock_by_product: raise ValueError( f"Cannot set price for {update.product!r} because it is not currently stocked." ) if update.product in seen_price_products: raise ValueError(f"Duplicate price update for {update.product!r}.") if update.price <= 0: raise ValueError("Price updates must be positive.") seen_price_products.add(update.product) seen_action_ids: set[int] = set() for action_id in action.action_ids: if action_id not in marketing_ids: raise ValueError(f"Unknown marketing action id {action_id}.") if action_id in seen_action_ids: raise ValueError(f"Duplicate marketing action id {action_id}.") if action_id in self._runtime["active_action_ids"]: raise ValueError(f"Marketing action {action_id} is already active.") seen_action_ids.add(action_id) for order in action.orders: if order.product not in market_by_product: raise ValueError(f"Unknown market product {order.product!r}.") if order.seller not in SELLER_PRICE_COLUMN: raise ValueError(f"Unknown seller {order.seller!r}.") if order.quantity <= 0: raise ValueError("Order quantity must be positive.") def _prepare_bundle(self, action: StoreAction, runtime: dict[str, Any]) -> dict[str, Any]: market_by_product = { row["product_name"]: row for row in runtime["market_offers"] } marketing_by_id = { row["action_id"]: row for row in runtime["marketing_actions_catalog"] } next_order_id = runtime["next_order_id"] prepared_orders = [] for index, request in enumerate(action.orders): market_row = market_by_product[request.product] price_key = SELLER_PRICE_COLUMN[request.seller] etd_key = SELLER_ETD_COLUMN[request.seller] quoted_eta = int(market_row[etd_key]) actual_eta = self._sample_actual_eta(request.seller, quoted_eta) prepared_orders.append( { "order_id": f"ORD-{next_order_id + index:04d}", "product_name": request.product, "quantity": request.quantity, "seller_name": request.seller, "buy_price_inr": float(market_row[price_key]), "days_remaining": actual_eta, "status": "pending", "created_day": int(runtime["store"]["day"]), } ) prepared_actions = [ deepcopy(marketing_by_id[action_id]) for action_id in action.action_ids ] return { "price_updates": { update.product: float(update.price) for update in action.prices or [] }, "new_orders": prepared_orders, "new_action_ids": list(action.action_ids), "selected_actions": prepared_actions, "customer_satisfaction_delta": sum( row["customer_satisfaction_delta"] for row in prepared_actions ), } def _apply_bundle(self, runtime: dict[str, Any], bundle: dict[str, Any]) -> None: stock_by_product = { row["product_name"]: row for row in runtime["stock"] } for product, price in bundle["price_updates"].items(): stock_by_product[product]["selling_price_inr"] = price self._refresh_stock_value(stock_by_product[product]) for marketing_action in bundle["selected_actions"]: runtime["store"]["bank_balance_inr"] -= marketing_action["cost_inr"] runtime["store"]["cumulative_profit_inr"] -= marketing_action["cost_inr"] runtime["store"]["order_processing_capacity"] += marketing_action[ "order_processing_capacity_delta" ] runtime["active_action_ids"].add(marketing_action["action_id"]) for order in bundle["new_orders"]: runtime["store"]["bank_balance_inr"] -= order["buy_price_inr"] * order["quantity"] runtime["pending_orders"].append(deepcopy(order)) runtime["next_order_id"] += len(bundle["new_orders"]) def _simulate_days(self, runtime: dict[str, Any], days: int) -> dict[str, float]: recurring_costs_paid = 0.0 total_revenue = 0.0 total_cogs = 0.0 total_units_sold = 0.0 for _ in range(days): recurring_cost = self._active_recurring_cost(runtime) revenue, cogs, units_sold = self._run_daily_sales(runtime) daily_expense = runtime["store"]["current_daily_expense_inr"] + recurring_cost runtime["store"]["bank_balance_inr"] += revenue - daily_expense runtime["store"]["cumulative_profit_inr"] += revenue - cogs - daily_expense recurring_costs_paid += recurring_cost total_revenue += revenue total_cogs += cogs total_units_sold += units_sold self._advance_pending_orders(runtime) runtime["store"]["day"] += 1 return { "days_simulated": float(days), "revenue_inr": total_revenue, "cogs_inr": total_cogs, "units_sold": total_units_sold, "recurring_costs_paid_inr": recurring_costs_paid, } def _run_daily_sales(self, runtime: dict[str, Any]) -> tuple[float, float, float]: effective_daily_orders = runtime["store"]["order_processing_capacity"] * ( runtime["store"]["customer_satisfaction"] / 100.0 ) if effective_daily_orders <= 0: return 0.0, 0.0, 0.0 sale_candidates: list[tuple[dict[str, Any], float]] = [] total_score = 0.0 for row in runtime["stock"]: if row["quantity_on_hand"] <= 0: continue score = self._product_selling_score(row) if score <= 0: continue sale_candidates.append((row, score)) total_score += score if total_score <= 0: return 0.0, 0.0, 0.0 revenue = 0.0 cogs = 0.0 units_sold = 0.0 for row, score in sale_candidates: desired_sales = effective_daily_orders * score / total_score actual_sales = min(row["quantity_on_hand"], desired_sales) if actual_sales <= 0: continue row["quantity_on_hand"] -= actual_sales revenue += actual_sales * row["selling_price_inr"] cogs += actual_sales * row["buying_price_inr"] units_sold += actual_sales self._refresh_stock_value(row) return revenue, cogs, units_sold def _advance_pending_orders(self, runtime: dict[str, Any]) -> None: remaining_orders = [] for order in runtime["pending_orders"]: order["days_remaining"] -= 1 if order["days_remaining"] <= 0: self._receive_order(runtime, order) else: remaining_orders.append(order) runtime["pending_orders"] = remaining_orders def _receive_order(self, runtime: dict[str, Any], order: dict[str, Any]) -> None: stock_by_product = { row["product_name"]: row for row in runtime["stock"] } market_by_product = { row["product_name"]: row for row in runtime["market_offers"] } if order["product_name"] in stock_by_product: stock_row = stock_by_product[order["product_name"]] existing_quantity = stock_row["quantity_on_hand"] new_quantity = existing_quantity + order["quantity"] if new_quantity <= 0: return stock_row["buying_price_inr"] = ( (existing_quantity * stock_row["buying_price_inr"]) + (order["quantity"] * order["buy_price_inr"]) ) / new_quantity stock_row["quantity_on_hand"] = new_quantity self._refresh_stock_value(stock_row) return market_row = market_by_product[order["product_name"]] buying_price = order["buy_price_inr"] runtime["stock"].append( { "product_name": order["product_name"], "unit": market_row["unit"], "buying_price_inr": buying_price, "selling_price_inr": max( buying_price + 1, math.ceil(buying_price * NEW_PRODUCT_MARKUP), ), "demand_score": market_row["demand_score"], "quantity_on_hand": float(order["quantity"]), "stock_value_inr": buying_price * order["quantity"], } ) def _product_selling_score(self, stock_row: dict[str, Any]) -> float: selling_price = stock_row["selling_price_inr"] if selling_price <= 0: return 0.0 margin_percent = ( (selling_price - stock_row["buying_price_inr"]) / selling_price ) * 100.0 return max(0.0, stock_row["demand_score"] - (margin_percent * 3.0)) def _total_product_selling_score(self, stock_rows: list[dict[str, Any]]) -> float: return sum( self._product_selling_score(row) for row in stock_rows if row["quantity_on_hand"] > 0 ) def _updated_customer_satisfaction( self, base_satisfaction: float, before_score: float, after_score: float, ) -> float: if after_score <= 0: updated = base_satisfaction else: updated = (1.0 + ((after_score - before_score) / after_score)) * base_satisfaction return min(100.0, max(1.0, updated)) def _active_recurring_cost(self, runtime: dict[str, Any]) -> float: marketing_by_id = { row["action_id"]: row for row in runtime["marketing_actions_catalog"] } return float( sum( marketing_by_id[action_id]["recurring_cost_inr"] for action_id in runtime["active_action_ids"] ) ) def _sample_actual_eta(self, seller_name: str, quoted_eta: int) -> int: if seller_name == "Seller A": offset = self._rng.choices([0, 1, 2], weights=[2, 4, 5], k=1)[0] elif seller_name == "Seller B": offset = 0 elif seller_name == "Seller C": offset = self._rng.choices([-2, -1, 0], weights=[5, 4, 2], k=1)[0] else: offset = self._rng.choices([-1, 0, 1, 2], weights=[2, 4, 3, 2], k=1)[0] return max(1, quoted_eta + offset) def _capital(self, runtime: dict[str, Any]) -> float: return runtime["store"]["bank_balance_inr"] + self._inventory_value(runtime["stock"]) def _inventory_value(self, stock_rows: list[dict[str, Any]]) -> float: return sum(row["stock_value_inr"] for row in stock_rows) def _refresh_stock_value(self, stock_row: dict[str, Any]) -> None: stock_row["stock_value_inr"] = stock_row["buying_price_inr"] * stock_row["quantity_on_hand"] def _build_observation( self, runtime: dict[str, Any], *, reward: float, done: bool, summary: str, metadata: dict[str, Any], ) -> StoreObservation: active_action_ids = sorted(runtime["active_action_ids"]) marketing_actions = [ MarketingAction( **row, is_active=row["action_id"] in runtime["active_action_ids"], ) for row in runtime["marketing_actions_catalog"] ] return StoreObservation( store=StoreSnapshot( day=int(runtime["store"]["day"]), current_daily_expense_inr=self._round(runtime["store"]["current_daily_expense_inr"]), bank_balance_inr=self._round(runtime["store"]["bank_balance_inr"]), customer_satisfaction=self._round(runtime["store"]["customer_satisfaction"]), order_processing_capacity=self._round(runtime["store"]["order_processing_capacity"]), cumulative_profit_inr=self._round(runtime["store"]["cumulative_profit_inr"]), active_action_ids=active_action_ids, active_recurring_cost_inr=self._round(self._active_recurring_cost(runtime)), episode_horizon_days=self.EPISODE_HORIZON_DAYS, commit_window_days=self.COMMIT_WINDOW_DAYS, reward_horizon_days=self.REWARD_HORIZON_DAYS, ), stock=[ StockItem( product_name=row["product_name"], unit=row["unit"], buying_price_inr=self._round(row["buying_price_inr"]), selling_price_inr=self._round(row["selling_price_inr"]), demand_score=self._round(row["demand_score"]), quantity_on_hand=self._round(row["quantity_on_hand"]), stock_value_inr=self._round(row["stock_value_inr"]), ) for row in runtime["stock"] ], pending_orders=[ PendingOrder( order_id=row["order_id"], product_name=row["product_name"], quantity=int(row["quantity"]), seller_name=row["seller_name"], buy_price_inr=self._round(row["buy_price_inr"]), days_remaining=int(row["days_remaining"]), status=row["status"], created_day=int(row["created_day"]), ) for row in runtime["pending_orders"] ], market=[MarketOffer(**row) for row in runtime["market_offers"]], seller_profiles=[SellerProfile(**row) for row in runtime["seller_profiles"]], marketing_actions=marketing_actions, available_action=self._available_action_spec(), summary=summary, reward=reward, done=done, metadata=metadata, ) def _available_action_spec(self) -> AvailableActionSpec: return AvailableActionSpec( name="action", description=( "Submit one bundled action for the next 3 committed days. The same " "bundle is used to forecast reward across 6 simulated days." ), pass_action=StoreAction(), example_action=StoreAction( orders=[ OrderRequest(product="Rice", seller="Seller A", quantity=25), OrderRequest(product="Bread", seller="Seller C", quantity=12), ], prices=[PriceUpdate(product="Milk", price=33)], action_ids=[7], ), notes=[ "Use an empty action to pass.", "Orders can target any product in the market catalog.", "Price changes only apply to products currently in stock.", "Marketing actions are durable and cannot be activated twice.", ], ) def _rounded_metrics(self, metrics: dict[str, float]) -> dict[str, float]: return { key: self._round(value) if key != "days_simulated" else float(value) for key, value in metrics.items() } def _round(self, value: float, *, digits: int = 2) -> float: return round(float(value), digits)