src/supplymind_env/generator.py

from __future__ import annotations

import random

from .models import DifficultyProfile, HiddenRecipe, OrderTemplate, WarehouseSpec


PROFILES: dict[str, DifficultyProfile] = {
    "v1_cooperative_restock": DifficultyProfile(
        task_id="v1_cooperative_restock",
        warehouse_count=4,
        total_rounds=25,
        transfer_cap=7,
        truck_capacity=8,
        depot_replenishment_cap=5,
        depot_trucks=1,
        depot_procurement_cap=5,
        depot_procurement_lead_time=2,
        starting_trust=0.72,
        stockout_penalty_multiplier=0.75,
        holding_cost_per_unit=0.04,
        waste_cost_per_unit=0.9,
        fairness_penalty_weight=0.18,
        coalition_bonus=1.5,
    ),
    "v1_scarcity_negotiation": DifficultyProfile(
        task_id="v1_scarcity_negotiation",
        warehouse_count=5,
        total_rounds=30,
        transfer_cap=8,
        truck_capacity=8,
        depot_replenishment_cap=6,
        depot_trucks=1,
        depot_procurement_cap=6,
        depot_procurement_lead_time=2,
        starting_trust=0.62,
        stockout_penalty_multiplier=0.95,
        holding_cost_per_unit=0.055,
        waste_cost_per_unit=1.15,
        fairness_penalty_weight=0.28,
        coalition_bonus=2.5,
    ),
    "v1_crisis_coalition": DifficultyProfile(
        task_id="v1_crisis_coalition",
        warehouse_count=7,
        total_rounds=40,
        transfer_cap=10,
        truck_capacity=10,
        depot_replenishment_cap=8,
        depot_trucks=2,
        depot_procurement_cap=8,
        depot_procurement_lead_time=3,
        starting_trust=0.55,
        stockout_penalty_multiplier=1.05,
        holding_cost_per_unit=0.065,
        waste_cost_per_unit=1.35,
        fairness_penalty_weight=0.36,
        coalition_bonus=4.0,
    ),
}


def generate_recipe(task_id: str, seed: int) -> HiddenRecipe:
    profile = PROFILES[task_id]
    rng = random.Random(f"{task_id}:{seed}")
    specs = _warehouse_specs(profile.warehouse_count, rng)
    initial_inventory = _initial_inventory(specs, task_id, rng)
    central_depot_inventory = _central_depot_inventory(task_id, rng)
    initial_drivers = {spec.warehouse_id: 2 if profile.warehouse_count == 3 else rng.choice([1, 2]) for spec in specs}
    private_forecasts = _private_forecasts(specs, task_id, rng)
    orders = _orders(specs, task_id, profile.total_rounds, rng)
    return HiddenRecipe(
        task_id=task_id,
        seed=seed,
        profile=profile,
        warehouse_specs=tuple(specs),
        initial_inventory=initial_inventory,
        central_depot_inventory=central_depot_inventory,
        initial_drivers=initial_drivers,
        private_forecasts=private_forecasts,
        orders=tuple(orders),
    )


def _warehouse_specs(count: int, rng: random.Random) -> list[WarehouseSpec]:
    personalities = ["cooperative", "risk_averse", "selfish", "opportunistic", "risk_averse", "cooperative", "selfish"]
    rng.shuffle(personalities)
    base = [
        ("north", "North", "uptown", {"uptown": 1.0, "downtown": 2.4, "suburb": 3.0, "industrial": 2.8}),
        ("east", "East", "suburb", {"uptown": 2.7, "downtown": 2.1, "suburb": 1.0, "industrial": 2.5}),
        ("south", "South", "downtown", {"uptown": 2.2, "downtown": 1.0, "suburb": 2.5, "industrial": 2.4}),
        ("west", "West", "industrial", {"uptown": 2.8, "downtown": 2.3, "suburb": 2.1, "industrial": 1.0}),
        ("central", "Central", "midtown", {"uptown": 1.7, "downtown": 1.6, "suburb": 2.0, "industrial": 1.8, "midtown": 1.0, "riverside": 2.2, "campus": 1.9}),
        ("riverside", "Riverside", "riverside", {"uptown": 3.0, "downtown": 2.6, "suburb": 2.2, "industrial": 2.4, "midtown": 2.2, "riverside": 1.0, "campus": 2.8}),
        ("campus", "Campus", "campus", {"uptown": 2.4, "downtown": 2.7, "suburb": 2.0, "industrial": 3.1, "midtown": 1.9, "riverside": 2.8, "campus": 1.0}),
    ]
    regions = [item[2] for item in base[:count]]
    specs: list[WarehouseSpec] = []
    for index, (warehouse_id, label, region, costs) in enumerate(base[:count]):
        scoped_costs = {key: float(costs.get(key, 2.8)) for key in regions}
        delivery_time = {key: max(1, round(value)) for key, value in scoped_costs.items()}
        specs.append(
            WarehouseSpec(
                warehouse_id=warehouse_id,
                label=label,
                region=region,
                personality=personalities[index],
                safety_stock={
                    "fresh_milk": rng.choice([2, 3]),
                    "rice_bag_5kg": rng.choice([2, 3]),
                    "insulin_pack": rng.choice([1, 2]),
                    "usb_c_charger": rng.choice([1, 2]),
                },
                delivery_cost_by_region=scoped_costs,
                delivery_time_by_region=delivery_time,
            )
        )
    return specs


def _initial_inventory(specs: list[WarehouseSpec], task_id: str, rng: random.Random) -> dict[str, dict[str, int]]:
    base_units = 7 if task_id == "v1_cooperative_restock" else 6
    inventory: dict[str, dict[str, int]] = {}
    for spec in specs:
        inventory[spec.warehouse_id] = {
            "fresh_milk": max(1, base_units + rng.randint(-2, 3)),
            "rice_bag_5kg": max(1, base_units + rng.randint(-2, 3)),
            "insulin_pack": max(1, base_units - 2 + rng.randint(-1, 3)),
            "usb_c_charger": max(1, base_units - 3 + rng.randint(-1, 3)),
        }
    hot = rng.choice(specs).warehouse_id
    cold = rng.choice([spec.warehouse_id for spec in specs if spec.warehouse_id != hot])
    inventory[hot]["fresh_milk"] = max(1, inventory[hot]["fresh_milk"] - 3)
    inventory[cold]["fresh_milk"] += 3
    return inventory


def _central_depot_inventory(task_id: str, rng: random.Random) -> dict[str, int]:
    base = 8 if task_id == "v1_cooperative_restock" else 7
    if task_id == "v1_crisis_coalition":
        base = 9
    return {
        "fresh_milk": base + rng.randint(-1, 2),
        "rice_bag_5kg": base + rng.randint(-1, 2),
        "insulin_pack": max(3, base - 2 + rng.randint(-1, 2)),
        "usb_c_charger": max(2, base - 3 + rng.randint(-1, 2)),
    }


def _private_forecasts(specs: list[WarehouseSpec], task_id: str, rng: random.Random) -> dict[str, dict[str, int]]:
    multiplier = 1 if task_id == "v1_cooperative_restock" else 2
    return {
        spec.warehouse_id: {
            "fresh_milk": rng.randint(1, 3 + multiplier),
            "rice_bag_5kg": rng.randint(1, 3 + multiplier),
            "insulin_pack": rng.randint(1, 2 + multiplier),
            "usb_c_charger": rng.randint(0, 2 + multiplier),
        }
        for spec in specs
    }


def _orders(
    specs: list[WarehouseSpec],
    task_id: str,
    total_rounds: int,
    rng: random.Random,
) -> list[OrderTemplate]:
    regions = [spec.region for spec in specs]
    if "downtown" not in regions:
        regions.append("downtown")
    count_by_task = {
        "v1_cooperative_restock": 42,
        "v1_scarcity_negotiation": 58,
        "v1_crisis_coalition": 84,
    }
    orders: list[OrderTemplate] = []
    for index in range(count_by_task[task_id]):
        created = min(total_rounds - 1, int(index * total_rounds / count_by_task[task_id]) + rng.choice([0, 0, 1]))
        sku = rng.choices(
            ["fresh_milk", "rice_bag_5kg", "insulin_pack", "usb_c_charger"],
            weights=[0.34, 0.30, 0.22, 0.14],
            k=1,
        )[0]
        units = rng.choice([1, 1, 2, 2, 3])
        if task_id == "v1_crisis_coalition" and index % 6 == 0:
            units = rng.choice([3, 4])
        value_by_sku = {"rice_bag_5kg": 11.0, "fresh_milk": 13.0, "insulin_pack": 22.0, "usb_c_charger": 24.0}
        value = value_by_sku[sku]
        if task_id != "v1_cooperative_restock" and created >= total_rounds // 2:
            value += rng.choice([2.0, 3.5, 4.0])
        deadline = min(total_rounds, created + rng.choice([2, 3, 3, 4]))
        orders.append(
            OrderTemplate(
                order_id=f"o{index + 1}",
                created_round=created,
                region=rng.choice(regions),
                sku=sku,
                units=units,
                value_per_unit=value,
                deadline_round=deadline,
                priority=2 if value >= 12 else 1,
            )
        )
    return sorted(orders, key=lambda order: (order.created_round, order.order_id))