Blame for server/store_environment.py

"""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)