server/tasks.py

"""
Task configurations for the MedChain Env environment.

Four tasks of increasing difficulty:
  0. orientation_ward        — 2 days,  1 ward,  no events       (intro)
  1. single_ward_stable      — 3 days,  1 ward,  no events
  2. multi_ward_seasonal     — 6 days,  4 wards, flu surge + supplier delay
  3. hospital_network_crisis — 12 days, 4 sites, 5 overlapping crises
"""

from dataclasses import dataclass, field
from typing import Dict, List, Optional


@dataclass
class Product:
    product_id: str
    name: str
    shelf_life_days: Optional[int]  # None = non-perishable
    criticality: str                # "CRITICAL", "HIGH", "NORMAL"
    unit_cost: float
    locations: List[str]
    base_demand: float
    demand_std: float
    seasonal_amplitude: float       # 0.0 = no seasonality
    seasonal_period: int            # 0 = no cycle
    seasonal_phase: float           # phase offset in radians


@dataclass
class Supplier:
    supplier_id: str
    name: str
    base_lead_time: int
    lead_time_std: float            # 0.0 = deterministic
    cost_multiplier: float
    products: List[str]


@dataclass
class Location:
    location_id: str
    name: str
    capacity: Optional[int]         # None = unlimited


@dataclass(frozen=True)
class InboxMessageTemplate:
    priority: str
    sender: str
    subject: str
    body: str


@dataclass
class SimEvent:
    event_id: str
    event_type: str                 # "mci", "supplier_disruption", "product_recall",
                                    # "demand_surge", "budget_tighten", "cold_chain_breach"
    trigger_day: int
    duration_days: int              # 0 = instant/one-off
    params: Dict
    message: InboxMessageTemplate
    warning_message: Optional[InboxMessageTemplate]  # fires on trigger_day-1


@dataclass
class TaskConfig:
    name: str
    max_days: int
    actions_per_shift: int
    budget_limit: float
    products: List[Product]
    suppliers: List[Supplier]
    locations: List[Location]
    events: List[SimEvent]
    initial_stock_days: int
    benchmark_score: float


# ─── Task 0: Orientation Ward ────────────────────────────────────────────────

TASK0 = TaskConfig(
    name="orientation_ward",
    max_days=2,
    actions_per_shift=8,
    budget_limit=5_000.0,
    initial_stock_days=1,
    benchmark_score=0.88,
    locations=[
        Location("ward_general", "General Ward", capacity=None),
    ],
    products=[
        Product("GLOVE-001", "Surgical Gloves (box)", None, "NORMAL", 5.0, ["ward_general"], 20, 3, 0.0, 0, 0.0),
        Product("SYR-10",    "Syringes 10ml",         None, "NORMAL", 0.5, ["ward_general"], 30, 5, 0.0, 0, 0.0),
        Product("MASK-001",  "Surgical Masks",        None, "NORMAL", 1.0, ["ward_general"], 25, 4, 0.0, 0, 0.0),
    ],
    suppliers=[
        Supplier("MEDLINE", "MedLine Medical", base_lead_time=1, lead_time_std=0.0,
                 cost_multiplier=1.0,
                 products=["GLOVE-001", "SYR-10", "MASK-001"]),
    ],
    events=[],
)


# ─── Task 1: Single Ward Stable ──────────────────────────────────────────────

TASK1 = TaskConfig(
    name="single_ward_stable",
    max_days=3,
    actions_per_shift=10,
    budget_limit=20_000.0,
    initial_stock_days=2,
    benchmark_score=0.68,
    locations=[
        Location("ward_general", "General Ward", capacity=None),
    ],
    products=[
        Product("GLOVE-001", "Surgical Gloves (box)",  None, "NORMAL", 5.0,  ["ward_general"], 50, 8,  0.0, 0, 0.0),
        Product("IV-500",    "IV Bags 500ml",          540,  "NORMAL", 3.5,  ["ward_general"], 30, 5,  0.0, 0, 0.0),
        Product("SYR-10",    "Syringes 10ml",          None, "NORMAL", 0.5,  ["ward_general"], 80, 12, 0.0, 0, 0.0),
        Product("PARA-500",  "Paracetamol 500mg",      360,  "NORMAL", 0.1,  ["ward_general"], 40, 6,  0.0, 0, 0.0),
        Product("MASK-001",  "Surgical Masks",         None, "NORMAL", 1.0,  ["ward_general"], 60, 10, 0.0, 0, 0.0),
        Product("SAL-001",   "Saline Solution",        720,  "NORMAL", 2.5,  ["ward_general"], 25, 4,  0.0, 0, 0.0),
    ],
    suppliers=[
        Supplier("MEDLINE", "MedLine Medical", base_lead_time=2, lead_time_std=0.0,
                 cost_multiplier=1.0,
                 products=["GLOVE-001", "IV-500", "SYR-10", "PARA-500", "MASK-001", "SAL-001"]),
    ],
    events=[],
)


# ─── Task 2: Multi-Ward Seasonal ─────────────────────────────────────────────

TASK2 = TaskConfig(
    name="multi_ward_seasonal",
    max_days=6,
    actions_per_shift=14,
    budget_limit=50_000.0,
    initial_stock_days=3,
    benchmark_score=0.55,
    locations=[
        Location("central_pharmacy", "Central Pharmacy", capacity=15_000),
        Location("ward_icu",         "ICU",              capacity=None),
        Location("ward_emergency",   "Emergency",        capacity=None),
        Location("ward_general",     "General Medicine", capacity=None),
    ],
    products=[
        Product("IV-500",    "IV Bags 500ml",        540,  "NORMAL",   3.5, ["central_pharmacy", "ward_icu"], 25, 4, 0.0, 0, 0.0),
        Product("SAL-001",   "Saline Solution",      720,  "NORMAL",   2.5, ["central_pharmacy", "ward_icu", "ward_emergency"], 20, 3, 0.0, 0, 0.0),
        Product("SYR-10",    "Syringes 10ml",        None, "NORMAL",   0.5, ["central_pharmacy", "ward_icu", "ward_emergency", "ward_general"], 70, 10, 0.0, 0, 0.0),
        Product("ANTIVIR-01","Antiviral Medication",  365, "HIGH",    15.0, ["central_pharmacy", "ward_emergency"], 10, 2, 0.55, 365, 0.0),
        Product("MASK-N95",  "N95 Masks",            None, "NORMAL",   2.5, ["central_pharmacy", "ward_emergency", "ward_icu"], 30, 5, 0.4, 365, 0.0),
        Product("PARA-500",  "Paracetamol 500mg",    360,  "NORMAL",   0.1, ["central_pharmacy", "ward_general"], 40, 6, 0.3, 180, 0.0),
        Product("GLOVE-001", "Surgical Gloves (box)", None, "NORMAL",  5.0, ["central_pharmacy", "ward_icu", "ward_emergency", "ward_general"], 50, 8, 0.0, 0, 0.0),
        Product("SUTURE-01", "Suture Kit",           None, "NORMAL",  12.0, ["central_pharmacy", "ward_general"], 8,  2, 0.0, 0, 0.0),
        Product("DRAIN-01",  "Surgical Drain",       None, "NORMAL",  18.0, ["central_pharmacy", "ward_general"], 5,  1, 0.0, 0, 0.0),
        Product("DRESS-01",  "Wound Dressing Kit",   None, "NORMAL",   6.0, ["central_pharmacy", "ward_general", "ward_emergency"], 20, 4, 0.0, 0, 0.0),
    ],
    suppliers=[
        Supplier("FASTMED", "FastMed Express", base_lead_time=1, lead_time_std=0.0, cost_multiplier=1.4,
                 products=["IV-500", "SAL-001", "SYR-10", "ANTIVIR-01", "MASK-N95", "PARA-500", "GLOVE-001", "SUTURE-01", "DRAIN-01", "DRESS-01"]),
        Supplier("MEDLINE", "MedLine Medical", base_lead_time=4, lead_time_std=0.0, cost_multiplier=1.0,
                 products=["IV-500", "SAL-001", "SYR-10", "ANTIVIR-01", "MASK-N95", "PARA-500", "GLOVE-001", "SUTURE-01", "DRAIN-01", "DRESS-01"]),
    ],
    events=[
        SimEvent(
            event_id="flu_surge",
            event_type="demand_surge",
            trigger_day=3,
            duration_days=3,
            params={"products": ["ANTIVIR-01", "MASK-N95", "PARA-500"], "multiplier": 1.5},
            message=InboxMessageTemplate(
                priority="HIGH",
                sender="Regional Health Authority",
                subject="Influenza Activity Alert — Demand Surge Active",
                body=(
                    "Regional influenza activity has been elevated above seasonal baseline.\n"
                    "Emergency department visits up 40-55%. Surge in effect now.\n"
                    "Antiviral, N95, and analgesic stock levels require immediate review.\n"
                    "Prepare for increased ICU admissions over the next 3 days."
                ),
            ),
            warning_message=InboxMessageTemplate(
                priority="LOW",
                sender="Regional Health Authority",
                subject="Early Warning: Influenza Activity Rising",
                body=(
                    "Early indicators suggest influenza activity is rising above seasonal norms.\n"
                    "Consider reviewing antiviral and PPE stock levels as a precaution."
                ),
            ),
        ),
        SimEvent(
            event_id="medline_delay",
            event_type="supplier_disruption",
            trigger_day=4,
            duration_days=3,
            params={"supplier_id": "MEDLINE", "new_lead_time": 7, "reason": "industrial action"},
            message=InboxMessageTemplate(
                priority="HIGH",
                sender="MedLine Medical — Supply Chain",
                subject="Service Disruption Notice — Lead Time Extension",
                body=(
                    "Due to ongoing industrial action at our primary warehouse facility,\n"
                    "MedLine Medical lead times are currently extended from 4 to 7 days.\n"
                    "This affects all standard orders. FastMed Express remains unaffected.\n"
                    "We apologise for the inconvenience. Reference: SUPDIS-2026-0006."
                ),
            ),
            warning_message=None,
        ),
    ],
)


# ─── Task 3: Hospital Network Crisis ─────────────────────────────────────────

TASK3 = TaskConfig(
    name="hospital_network_crisis",
    max_days=12,
    actions_per_shift=18,
    budget_limit=150_000.0,
    initial_stock_days=4,
    benchmark_score=0.38,
    locations=[
        Location("regional_dc", "Regional Distribution Centre", capacity=None),
        Location("hospital_a",  "Hospital A",                   capacity=None),
        Location("hospital_b",  "Hospital B",                   capacity=None),
        Location("hospital_c",  "Hospital C",                   capacity=None),
    ],
    products=[
        Product("B-001",     "O-Negative Blood (RBC)",      42,   "CRITICAL", 350.0, ["regional_dc", "hospital_a", "hospital_b", "hospital_c"], 3, 1, 0.0, 0, 0.0),
        Product("B-002",     "Platelet Pack",               5,    "CRITICAL", 250.0, ["regional_dc", "hospital_a", "hospital_b", "hospital_c"], 2, 1, 0.0, 0, 0.0),
        Product("FFP-001",   "Fresh Frozen Plasma",         365,  "HIGH",     120.0, ["regional_dc", "hospital_a", "hospital_b", "hospital_c"], 4, 1, 0.0, 0, 0.0),
        Product("INS-001",   "Insulin (opened vial)",       28,   "HIGH",     45.0,  ["hospital_a", "hospital_b", "hospital_c"], 8,  2, 0.0, 0, 0.0),
        Product("CHEMO-01",  "Chemotherapy Agent (recon.)", 7,    "HIGH",     800.0, ["hospital_a", "hospital_b"],              1,  0, 0.0, 0, 0.0),
        Product("IV-SAL-500","IV Saline Solution 500ml",    720,  "NORMAL",   3.5,   ["regional_dc", "hospital_a", "hospital_b", "hospital_c"], 60, 10, 0.0, 0, 0.0),
        Product("IV-500",    "IV Bags 500ml",               540,  "NORMAL",   3.5,   ["regional_dc", "hospital_a", "hospital_b", "hospital_c"], 40, 7,  0.0, 0, 0.0),
        Product("SYR-10",    "Syringes 10ml",               None, "NORMAL",   0.5,   ["regional_dc", "hospital_a", "hospital_b", "hospital_c"], 100,15, 0.0, 0, 0.0),
        Product("GLOVE-001", "Surgical Gloves (box)",       None, "NORMAL",   5.0,   ["regional_dc", "hospital_a", "hospital_b", "hospital_c"], 40, 6,  0.0, 0, 0.0),
        Product("PARA-500",  "Paracetamol 500mg",           360,  "NORMAL",   0.1,   ["hospital_a", "hospital_b", "hospital_c"], 60, 8,  0.0, 0, 0.0),
        Product("B-003",     "AB-Positive Blood (RBC)",     42,   "HIGH",     320.0, ["regional_dc", "hospital_a", "hospital_b", "hospital_c"], 2, 1, 0.0, 0, 0.0),
        Product("MASK-001",  "Surgical Masks",              None, "NORMAL",   1.0,   ["regional_dc", "hospital_a", "hospital_b", "hospital_c"], 80, 12, 0.0, 0, 0.0),
        Product("DRAIN-01",  "Surgical Drain",              None, "NORMAL",   18.0,  ["hospital_a", "hospital_b", "hospital_c"], 4,  1,  0.0, 0, 0.0),
        Product("SAL-001",   "Saline Solution",             720,  "NORMAL",   2.5,   ["regional_dc", "hospital_a", "hospital_b", "hospital_c"], 30, 5,  0.0, 0, 0.0),
        Product("DRESS-01",  "Wound Dressing Kit",          None, "NORMAL",   6.0,   ["hospital_a", "hospital_b", "hospital_c"], 25, 4,  0.0, 0, 0.0),
    ],
    suppliers=[
        Supplier("BLOODBANK-A", "Regional Blood Bank", base_lead_time=1, lead_time_std=0.5, cost_multiplier=1.0,
                 products=["B-001", "B-002", "FFP-001", "B-003"]),
        Supplier("SUPPLIER-A",  "HealthCo Supplies",   base_lead_time=3, lead_time_std=0.5, cost_multiplier=1.0,
                 products=["IV-SAL-500", "IV-500", "SYR-10", "GLOVE-001", "PARA-500", "MASK-001", "DRAIN-01", "SAL-001", "DRESS-01"]),
        Supplier("SUPPLIER-B",  "MedFast Express",     base_lead_time=2, lead_time_std=0.0, cost_multiplier=1.4,
                 products=["IV-SAL-500", "IV-500", "SYR-10", "GLOVE-001", "PARA-500", "MASK-001", "DRAIN-01", "SAL-001", "DRESS-01"]),
        Supplier("PHARMA-X",    "PharmaCorp",          base_lead_time=3, lead_time_std=0.5, cost_multiplier=1.0,
                 products=["INS-001", "CHEMO-01", "PARA-500"]),
    ],
    events=[
        SimEvent(
            event_id="cold_chain_breach",
            event_type="cold_chain_breach",
            trigger_day=3,
            duration_days=0,
            params={
                "location_id": "regional_dc",
                "product_id": "B-002",
                "qty_affected": "all",
            },
            message=InboxMessageTemplate(
                priority="CRITICAL",
                sender="Regional DC Facilities Management",
                subject="URGENT: Cold Chain Breach — Platelet Inventory Compromised",
                body=(
                    "Cold chain monitoring detected temperature excursion in Fridge Unit 3.\n"
                    "Recorded +8°C for 4 hours (safe range: +20 to +24°C).\n"
                    "ALL platelet units at Regional DC are presumed compromised and auto-quarantined.\n"
                    "Affected product: Platelet Pack (B-002) — ALL lots at regional_dc.\n"
                    "Action required:\n"
                    "  1. Arrange emergency replacement order from BLOODBANK-A\n"
                    "  2. Notify clinical teams — platelet supply now limited to hospital stock only.\n"
                    "Ref: CCBR-2026-0003"
                ),
            ),
            warning_message=None,
        ),
        SimEvent(
            event_id="supplier_a_disruption",
            event_type="supplier_disruption",
            trigger_day=6,
            duration_days=9,
            params={"supplier_id": "SUPPLIER-A", "new_lead_time": 7, "reason": "force majeure — flu absenteeism"},
            message=InboxMessageTemplate(
                priority="HIGH",
                sender="HealthCo Supplies — Logistics",
                subject="Force Majeure Notice — Extended Lead Times",
                body=(
                    "HealthCo Supplies hereby provides formal notice of force majeure event.\n"
                    "Widespread workforce absenteeism due to influenza has impacted fulfilment operations.\n"
                    "Effective immediately, standard lead times are extended from 3 to 7 days.\n"
                    "All product lines affected. MedFast Express (SUPPLIER-B) is unaffected.\n"
                    "Ref: FM-2026-0006"
                ),
            ),
            warning_message=None,
        ),
        SimEvent(
            event_id="mci_warning",
            event_type="demand_surge",
            trigger_day=8,
            duration_days=0,
            params={},
            message=InboxMessageTemplate(
                priority="HIGH",
                sender="Emergency Management Coordination",
                subject="Mass Casualty Incident — STANDBY",
                body=(
                    "Multi-vehicle collision reported on Interstate highway.\n"
                    "Current status: STANDBY. Incident Command not yet activated.\n"
                    "Preliminary estimate: 15-25 critically injured.\n"
                    "All blood banks placed on AMBER alert.\n"
                    "Recommend pre-emptive review of O-neg and platelet stock levels NOW.\n"
                    "Further update to follow."
                ),
            ),
            warning_message=None,
        ),
        SimEvent(
            event_id="mci_activation",
            event_type="mci",
            trigger_day=9,
            duration_days=3,
            params={
                "products": ["B-001", "B-002", "FFP-001"],
                "demand_multiplier": 3.0,
                "locations": ["hospital_a", "hospital_b", "hospital_c"],
                "mci_tracking": True,
            },
            message=InboxMessageTemplate(
                priority="CRITICAL",
                sender="Incident Command System",
                subject="MCI ACTIVATION — Mass Casualty Event",
                body=(
                    "INCIDENT COMMAND ACTIVATED.\n"
                    "Multi-vehicle collision — confirmed 23 critically injured en route.\n"
                    "Hospital A ETA: 18 min | Hospital B ETA: 29 min | Hospital C ETA: 41 min.\n"
                    "Trauma surgery teams mobilised at all hospitals.\n"
                    "All blood banks placed on RED alert.\n"
                    "Est. blood product requirements: 60-90 units O-neg RBC, 30-40 platelet packs.\n"
                    "IMMEDIATE ACTION: Verify blood product inventory and initiate emergency procurement."
                ),
            ),
            warning_message=None,
        ),
        SimEvent(
            event_id="iv_saline_recall",
            event_type="product_recall",
            trigger_day=11,
            duration_days=0,
            params={
                "product_id": "IV-SAL-500",
                "lot_id": "RECALL-LOT-IV2026-9821",
                "locations_with_lot": ["hospital_a", "hospital_b", "hospital_c", "regional_dc"],
                "qty_per_location": 60,
            },
            message=InboxMessageTemplate(
                priority="CRITICAL",
                sender="Pharmacy Automated System",
                subject="MANDATORY RECALL — IV Saline Solution 500ml",
                body=(
                    "MANDATORY RECALL — Health Authority ref #HA-2026-0013\n"
                    "Product: IV Saline Solution 500ml (IV-SAL-500)\n"
                    "Affected lot: RECALL-LOT-IV2026-9821\n"
                    "Reason: Potential endotoxin contamination detected in batch.\n"
                    "ACTION REQUIRED:\n"
                    "  1. Query inventory at ALL locations for lot RECALL-LOT-IV2026-9821\n"
                    "  2. Quarantine ALL affected units immediately — do not use\n"
                    "  3. Submit replacement order\n"
                    "Supplier contact: SUPPLIER-A case #88291. Ref: RECALL-2026-0013."
                ),
            ),
            warning_message=None,
        ),
    ],
)


def get_task_config(task_name: str) -> TaskConfig:
    if task_name == "orientation_ward":
        return TASK0
    elif task_name == "single_ward_stable":
        return TASK1
    elif task_name == "multi_ward_seasonal":
        return TASK2
    elif task_name == "hospital_network_crisis":
        return TASK3
    else:
        raise ValueError(
            f"Unknown task: {task_name!r}. "
            "Choose from: orientation_ward, single_ward_stable, multi_ward_seasonal, hospital_network_crisis"
        )