Add DispatchSimulation engine, geometry helpers, caller text templates, and observation renderer

scenario_loader.py

@@ -0,0 +1,36 @@
+"""Load task scenarios from YAML files bundled with the package."""
+
+from __future__ import annotations
+
+import os
+from typing import List
+
+import yaml
+
+# Resolve the tasks directory relative to this file so it works regardless of CWD
+_TASKS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tasks")
+
+VALID_TASKS = ("easy", "medium", "hard")
+
+
+def load_scenario(task_name: str) -> dict:
+    """Load and return the scenario dict for the given task.
+
+    Raises:
+        ValueError: if task_name is not one of {easy, medium, hard}.
+        FileNotFoundError: if the YAML file is missing.
+    """
+    if task_name not in VALID_TASKS:
+        raise ValueError(
+            f"Unknown task '{task_name}'. Valid tasks: {', '.join(VALID_TASKS)}"
+        )
+    path = os.path.join(_TASKS_DIR, f"{task_name}.yaml")
+    with open(path, "r", encoding="utf-8") as f:
+        scenario = yaml.safe_load(f)
+    if "name" not in scenario:
+        scenario["name"] = task_name
+    return scenario
+
+
+def list_tasks() -> List[str]:
+    return list(VALID_TASKS)

simulation.py

@@ -0,0 +1,405 @@
+"""DispatchSimulation engine. Pure Python, deterministic, seedable."""
+
+from __future__ import annotations
+
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+
+from models import (
+    EmergencyCall,
+    EmergencyType,
+    EmergencyUnit,
+    Hospital,
+    Position,
+    Severity,
+    UnitStatus,
+    UnitType,
+    WorldConfig,
+)
+from reward import calculate_call_outcome, get_effectiveness
+from utils import (
+    calculate_distance,
+    calculate_eta,
+    generate_caller_text,
+    get_capable_units,
+    get_optimal_unit,
+)
+
+
+def _parse_severity(value) -> Severity:
+    if isinstance(value, Severity):
+        return value
+    return Severity(int(value))
+
+
+def generate_call_schedule(scenario: dict, seed: int) -> List[EmergencyCall]:
+    """Build a deterministic list of EmergencyCall objects from a scenario dict."""
+    rng = np.random.RandomState(seed)
+    calls: List[EmergencyCall] = []
+    grid_size = scenario.get("grid_size", 10.0)
+    inaccuracy = float(scenario.get("caller_inaccuracy", 0.0))
+
+    for idx, call_cfg in enumerate(scenario["calls"], start=1):
+        true_type = EmergencyType(call_cfg["type"])
+        true_severity = _parse_severity(call_cfg["severity"])
+
+        if inaccuracy > 0 and rng.random() < inaccuracy:
+            other_types = [t for t in EmergencyType if t != true_type]
+            reported_type = EmergencyType(str(rng.choice([t.value for t in other_types])))
+            shifted = max(1, min(5, true_severity.value + int(rng.randint(-1, 2))))
+            reported_severity = Severity(shifted)
+        else:
+            reported_type = true_type
+            reported_severity = true_severity
+
+        location = Position(
+            x=round(float(rng.uniform(0.5, grid_size - 0.5)), 1),
+            y=round(float(rng.uniform(0.5, grid_size - 0.5)), 1),
+        )
+        caller_text = generate_caller_text(true_type, reported_type, rng)
+
+        calls.append(
+            EmergencyCall(
+                call_id=f"CALL-{idx:03d}",
+                timestamp=int(call_cfg["arrival_minute"]),
+                caller_description=caller_text,
+                location=location,
+                true_type=true_type,
+                true_severity=true_severity,
+                reported_type=reported_type,
+                reported_severity=reported_severity,
+                requires_unit_types=get_capable_units(true_type),
+                optimal_unit_type=get_optimal_unit(true_type),
+            )
+        )
+
+    calls.sort(key=lambda c: c.timestamp)
+    return calls
+
+
+# Scene-time table: how long a unit stays on scene treating a call
+SCENE_TIME_MINUTES = {
+    EmergencyType.CARDIAC_ARREST: 20,
+    EmergencyType.TRAUMA: 25,
+    EmergencyType.STROKE: 15,
+    EmergencyType.FIRE: 30,
+    EmergencyType.BREATHING: 15,
+    EmergencyType.MINOR_INJURY: 10,
+    EmergencyType.MENTAL_HEALTH: 20,
+}
+
+
+class DispatchSimulation:
+    """Discrete-time simulation of an emergency dispatch episode."""
+
+    def __init__(self, scenario: dict, seed: int = 42) -> None:
+        self.scenario_name: str = scenario.get("name", "unnamed")
+        self.scenario: dict = scenario
+        self.seed: int = seed
+        self.rng = np.random.RandomState(seed)
+
+        world_cfg = scenario.get("world_config", {})
+        self.config = WorldConfig(**world_cfg)
+
+        self.current_time: int = 0
+        self.episode_done: bool = False
+
+        self.all_calls: List[EmergencyCall] = generate_call_schedule(scenario, seed)
+        self.active_calls: List[EmergencyCall] = []
+        self.completed_calls: List[dict] = []
+        self.timed_out_calls: List[dict] = []
+        self.dispatches: List[dict] = []
+
+        self.units: Dict[str, EmergencyUnit] = {}
+        for unit_cfg in scenario["units"]:
+            unit = EmergencyUnit(**unit_cfg)
+            self.units[unit.unit_id] = unit
+
+        self.hospitals: Dict[str, Hospital] = {}
+        for hosp_cfg in scenario["hospitals"]:
+            hosp = Hospital(**hosp_cfg)
+            self.hospitals[hosp.hospital_id] = hosp
+
+        self.call_index: int = 0
+        # Release any calls scheduled for time 0
+        self._release_due_calls()
+
+    # ------------------------------------------------------------------
+    # Time advancement
+    # ------------------------------------------------------------------
+
+    def _release_due_calls(self) -> None:
+        """Move calls whose arrival time has passed into the active queue."""
+        while (
+            self.call_index < len(self.all_calls)
+            and self.all_calls[self.call_index].timestamp <= self.current_time
+        ):
+            call = self.all_calls[self.call_index]
+            call.active = True
+            self.active_calls.append(call)
+            self.call_index += 1
+
+    def advance_time(self, minutes: int = 1) -> None:
+        """Step the simulation forward by ``minutes`` discrete minutes."""
+        if self.episode_done:
+            return
+        minutes = max(1, int(minutes))
+        for _ in range(minutes):
+            self.current_time += 1
+            self._tick_once()
+            if self.episode_done:
+                break
+
+    def _tick_once(self) -> None:
+        """Advance simulation by exactly one minute, updating units & calls."""
+        # 1. Move units according to their status
+        for unit in self.units.values():
+            if unit.status == UnitStatus.EN_ROUTE:
+                self._move_unit_toward_call(unit)
+            elif unit.status == UnitStatus.ON_SCENE:
+                if unit.busy_until is not None and self.current_time >= unit.busy_until:
+                    unit.status = UnitStatus.RETURNING
+                    unit.assigned_call_id = None
+                    unit.assigned_hospital_id = None
+            elif unit.status == UnitStatus.RETURNING:
+                self._move_unit_toward_base(unit)
+
+        # 2. Time-out any active call that has waited too long
+        for call in list(self.active_calls):
+            if call.dispatched_unit_id is None:
+                wait = self.current_time - call.timestamp
+                if wait >= self.config.call_timeout_minutes:
+                    call.active = False
+                    self.active_calls.remove(call)
+                    self.timed_out_calls.append(
+                        {
+                            "call_id": call.call_id,
+                            "true_type": call.true_type.value,
+                            "true_severity": call.true_severity.value,
+                            "outcome_score": 0.0,
+                            "reason": "timed_out",
+                        }
+                    )
+
+        # 3. Release new calls
+        self._release_due_calls()
+
+        # 4. Episode end conditions
+        if self.current_time >= self.config.time_limit_minutes:
+            self._finalize_episode("time_limit")
+            return
+
+        no_more_incoming = self.call_index >= len(self.all_calls)
+        no_pending = all(c.dispatched_unit_id is not None for c in self.active_calls)
+        all_units_idle = all(u.status == UnitStatus.AVAILABLE for u in self.units.values())
+        if no_more_incoming and no_pending and all_units_idle and not self.active_calls:
+            self._finalize_episode("all_resolved")
+
+    def _finalize_episode(self, reason: str) -> None:
+        """Mark episode done.
+
+        Any remaining call (whether un-dispatched OR dispatched but the unit
+        never actually arrived on scene) is recorded as a timeout — the agent
+        failed to deliver care in time, so the patient outcome is 0.0.
+        """
+        self.episode_done = True
+        for call in list(self.active_calls):
+            self.timed_out_calls.append(
+                {
+                    "call_id": call.call_id,
+                    "true_type": call.true_type.value,
+                    "true_severity": call.true_severity.value,
+                    "outcome_score": 0.0,
+                    "reason": reason
+                    if call.dispatched_unit_id is None
+                    else f"{reason}_in_transit",
+                }
+            )
+        self.active_calls.clear()
+
+    # ------------------------------------------------------------------
+    # Unit movement
+    # ------------------------------------------------------------------
+
+    def _move_unit_toward_call(self, unit: EmergencyUnit) -> None:
+        call = self._get_call_by_id(unit.assigned_call_id) if unit.assigned_call_id else None
+        if call is None:
+            unit.status = UnitStatus.AVAILABLE
+            unit.assigned_call_id = None
+            return
+
+        distance_per_step = (unit.speed_kmh / 60.0) * self.config.step_duration_minutes
+        dist = calculate_distance(unit.position, call.location)
+
+        if dist <= distance_per_step:
+            unit.position = Position(x=call.location.x, y=call.location.y)
+            unit.status = UnitStatus.ON_SCENE
+            response_time = float(self.current_time - call.timestamp)
+            call.response_time = response_time
+
+            hospital = (
+                self.hospitals.get(unit.assigned_hospital_id)
+                if unit.assigned_hospital_id
+                else None
+            )
+
+            outcome = calculate_call_outcome(call, unit, response_time, hospital)
+            call.outcome_score = outcome
+            call.active = False
+            if call in self.active_calls:
+                self.active_calls.remove(call)
+
+            if hospital is not None and not hospital.on_diversion and hospital.available_beds > 0:
+                hospital.available_beds = max(0, hospital.available_beds - 1)
+                call.delivered_hospital_id = hospital.hospital_id
+
+            self.completed_calls.append(
+                {
+                    "call_id": call.call_id,
+                    "true_type": call.true_type.value,
+                    "true_severity": call.true_severity.value,
+                    "response_time": response_time,
+                    "outcome_score": outcome,
+                    "unit_id": unit.unit_id,
+                    "unit_type": unit.unit_type.value,
+                    "effectiveness": get_effectiveness(unit.unit_type, call.true_type),
+                    "hospital_id": call.delivered_hospital_id,
+                }
+            )
+
+            scene_time = SCENE_TIME_MINUTES.get(call.true_type, 15)
+            unit.busy_until = self.current_time + scene_time
+        else:
+            ratio = distance_per_step / dist
+            unit.position = Position(
+                x=round(unit.position.x + (call.location.x - unit.position.x) * ratio, 3),
+                y=round(unit.position.y + (call.location.y - unit.position.y) * ratio, 3),
+            )
+
+    def _move_unit_toward_base(self, unit: EmergencyUnit) -> None:
+        distance_per_step = (unit.speed_kmh / 60.0) * self.config.step_duration_minutes
+        dist = calculate_distance(unit.position, unit.base_position)
+        if dist <= distance_per_step:
+            unit.position = Position(x=unit.base_position.x, y=unit.base_position.y)
+            unit.status = UnitStatus.AVAILABLE
+            unit.busy_until = None
+        else:
+            ratio = distance_per_step / dist
+            unit.position = Position(
+                x=round(unit.position.x + (unit.base_position.x - unit.position.x) * ratio, 3),
+                y=round(unit.position.y + (unit.base_position.y - unit.position.y) * ratio, 3),
+            )
+
+    # ------------------------------------------------------------------
+    # Action handlers (called from the MCP environment)
+    # ------------------------------------------------------------------
+
+    def dispatch(
+        self, call_id: str, unit_id: str, hospital_id: Optional[str] = None
+    ) -> Tuple[float, str]:
+        """Dispatch a unit to a call (optionally pre-assigning a destination hospital)."""
+        call = self._get_active_undispatched_call(call_id)
+        if call is None:
+            return -0.05, f"Call {call_id} not found in pending queue."
+
+        unit = self.units.get(unit_id)
+        if unit is None:
+            return -0.05, f"Unit {unit_id} not found."
+        if unit.status != UnitStatus.AVAILABLE:
+            return -0.05, f"Unit {unit_id} is {unit.status.value}, cannot dispatch."
+
+        # Treat empty string / whitespace as "no hospital chosen"
+        if isinstance(hospital_id, str):
+            hospital_id = hospital_id.strip() or None
+
+        chosen_hospital = None
+        if hospital_id is not None:
+            chosen_hospital = self.hospitals.get(hospital_id)
+            if chosen_hospital is None:
+                return -0.02, f"Hospital '{hospital_id}' not found."
+
+        unit.status = UnitStatus.EN_ROUTE
+        unit.assigned_call_id = call.call_id
+        unit.assigned_hospital_id = hospital_id
+        call.dispatched_unit_id = unit.unit_id
+
+        eta = calculate_eta(unit, call.location)
+        effectiveness = get_effectiveness(unit.unit_type, call.true_type)
+
+        self.dispatches.append(
+            {
+                "call_id": call.call_id,
+                "unit_id": unit.unit_id,
+                "unit_type": unit.unit_type.value,
+                "true_type": call.true_type.value,
+                "true_severity": call.true_severity.value,
+                "arrival_time": call.timestamp,
+                "dispatch_time": self.current_time,
+                "timeout_window": self.config.call_timeout_minutes,
+                "eta": eta,
+                "effectiveness": effectiveness,
+                "hospital_id": hospital_id,
+            }
+        )
+
+        msg = (
+            f"Dispatched {unit.unit_id} to {call.call_id}. "
+            f"ETA {eta:.1f} min. Unit effectiveness for {call.true_type.value}: "
+            f"{effectiveness:.0%}."
+        )
+        if hospital_id is not None and chosen_hospital is not None:
+            msg += f" Destination hospital: {chosen_hospital.name}."
+        return 0.02 * effectiveness, msg
+
+    def classify(self, call_id: str, severity: int) -> Tuple[float, str]:
+        call = self._get_active_undispatched_call(call_id)
+        if call is None:
+            return -0.02, f"Call {call_id} not in pending queue."
+        try:
+            new_sev = Severity(int(severity))
+        except ValueError:
+            return -0.02, f"Invalid severity {severity}; must be 1-5."
+        old = call.reported_severity
+        call.reported_severity = new_sev
+        return 0.01, f"Reclassified {call_id} severity from {old} to {new_sev.value}."
+
+    def callback(self, call_id: str, question: str) -> Tuple[float, str]:
+        call = self._get_active_undispatched_call(call_id)
+        if call is None:
+            return -0.02, f"Call {call_id} not in pending queue."
+        # 70% chance the caller clarifies; 30% they're too distressed
+        if self.rng.random() < 0.70:
+            call.reported_type = call.true_type
+            call.reported_severity = call.true_severity
+            return (
+                0.02,
+                f"Caller for {call.call_id} confirms: this is a {call.true_type.value}, "
+                f"severity {call.true_severity.value}.",
+            )
+        return 0.0, f"Caller for {call.call_id} is too distressed to give clear info."
+
+    # ------------------------------------------------------------------
+    # Lookups
+    # ------------------------------------------------------------------
+
+    def _get_call_by_id(self, call_id: str) -> Optional[EmergencyCall]:
+        for c in self.all_calls:
+            if c.call_id == call_id:
+                return c
+        return None
+
+    def _get_active_undispatched_call(self, call_id: str) -> Optional[EmergencyCall]:
+        for c in self.active_calls:
+            if c.call_id == call_id and c.dispatched_unit_id is None:
+                return c
+        return None
+
+    def get_pending_calls(self) -> List[EmergencyCall]:
+        return [c for c in self.active_calls if c.dispatched_unit_id is None]
+
+    def get_available_units(self) -> List[EmergencyUnit]:
+        return [u for u in self.units.values() if u.status == UnitStatus.AVAILABLE]
+
+    def total_calls(self) -> int:
+        return len(self.all_calls)

text_view.py

@@ -0,0 +1,142 @@
+"""Render a DispatchSimulation as a human-readable text view for the LLM agent."""
+
+from __future__ import annotations
+
+from typing import List
+
+from models import EmergencyCall, EmergencyUnit, Hospital, UnitStatus
+from simulation import DispatchSimulation
+from utils import calculate_distance, calculate_eta
+
+# Maximum number of calls / units / outcomes to show in the text view.
+# Truncation prevents context blow-up on the hard task (30 calls).
+MAX_PENDING_CALLS = 8
+MAX_BUSY_UNITS = 8
+MAX_RECENT_OUTCOMES = 3
+
+
+def _format_call(call: EmergencyCall, sim: DispatchSimulation) -> List[str]:
+    wait = sim.current_time - call.timestamp
+    rt = call.reported_type.value if call.reported_type else "unknown"
+    rs = call.reported_severity.value if call.reported_severity else "?"
+    return [
+        f'  {call.call_id} [t={call.timestamp}min] "{call.caller_description}"',
+        (
+            f"    location=({call.location.x}, {call.location.y})  "
+            f"reported={rt}/sev{rs}  waiting={wait}min"
+        ),
+    ]
+
+
+def _format_unit(unit: EmergencyUnit, sim: DispatchSimulation, pending: list) -> str:
+    base = (
+        f"  {unit.unit_id:7s} | {unit.unit_type.value:14s} | "
+        f"pos=({unit.position.x:.1f}, {unit.position.y:.1f})"
+    )
+    if pending:
+        closest = min(pending, key=lambda c: calculate_distance(unit.position, c.location))
+        eta = calculate_eta(unit, closest.location)
+        base += f"  closest_call_eta={eta:.1f}min ({closest.call_id})"
+    return base
+
+
+def _format_busy_unit(unit: EmergencyUnit) -> str:
+    detail = unit.status.value
+    if unit.assigned_call_id:
+        detail += f" -> {unit.assigned_call_id}"
+    if unit.busy_until is not None and unit.status == UnitStatus.ON_SCENE:
+        detail += f" (free at t={unit.busy_until}min)"
+    return f"  {unit.unit_id:7s} | {unit.unit_type.value:14s} | {detail}"
+
+
+def _format_hospital(hosp: Hospital) -> str:
+    specs = []
+    if hosp.has_trauma_center:
+        specs.append("trauma")
+    if hosp.has_cardiac_unit:
+        specs.append("cardiac")
+    if hosp.has_stroke_unit:
+        specs.append("stroke")
+    status = "DIVERSION" if hosp.on_diversion else "open"
+    return (
+        f"  {hosp.hospital_id} {hosp.name} ({hosp.position.x},{hosp.position.y})  "
+        f"beds={hosp.available_beds}/{hosp.capacity}  "
+        f"specialties=[{','.join(specs) or 'none'}]  status={status}"
+    )
+
+
+def render_dispatch_center(sim: DispatchSimulation, task_name: str) -> str:
+    """Pretty-print the current state for the LLM agent."""
+    lines: List[str] = []
+    lines.append("=== DISPATCHPULSE DISPATCH CENTER ===")
+    lines.append(
+        f"task={task_name}  time={sim.current_time}min/"
+        f"{sim.config.time_limit_minutes}min  "
+        f"scenario={sim.scenario_name}"
+    )
+    lines.append("")
+
+    # Pending calls (sorted by reported severity, then arrival time)
+    pending = sim.get_pending_calls()
+    pending_sorted = sorted(
+        pending,
+        key=lambda c: (
+            c.reported_severity.value if c.reported_severity else 5,
+            c.timestamp,
+        ),
+    )
+    lines.append(f"PENDING CALLS ({len(pending_sorted)} total):")
+    if not pending_sorted:
+        lines.append("  (none)")
+    for call in pending_sorted[:MAX_PENDING_CALLS]:
+        lines.extend(_format_call(call, sim))
+    if len(pending_sorted) > MAX_PENDING_CALLS:
+        hidden = len(pending_sorted) - MAX_PENDING_CALLS
+        lines.append(f"  ... and {hidden} more lower-priority calls")
+    lines.append("")
+
+    # Available units
+    available = sim.get_available_units()
+    lines.append(f"AVAILABLE UNITS ({len(available)} total):")
+    if not available:
+        lines.append("  (none — all units busy)")
+    for unit in available:
+        lines.append(_format_unit(unit, sim, pending_sorted[:MAX_PENDING_CALLS]))
+    lines.append("")
+
+    # Busy units
+    busy = [u for u in sim.units.values() if u.status != UnitStatus.AVAILABLE]
+    if busy:
+        lines.append(f"BUSY UNITS ({len(busy)} total):")
+        for unit in busy[:MAX_BUSY_UNITS]:
+            lines.append(_format_busy_unit(unit))
+        if len(busy) > MAX_BUSY_UNITS:
+            lines.append(f"  ... and {len(busy) - MAX_BUSY_UNITS} more busy units")
+        lines.append("")
+
+    # Hospitals
+    lines.append("HOSPITALS:")
+    for hosp in sim.hospitals.values():
+        lines.append(_format_hospital(hosp))
+    lines.append("")
+
+    # Recent outcomes
+    if sim.completed_calls:
+        recent = sim.completed_calls[-MAX_RECENT_OUTCOMES:]
+        lines.append("RECENT OUTCOMES:")
+        for r in recent:
+            mark = "OK " if r["outcome_score"] >= 0.5 else "BAD"
+            lines.append(
+                f"  [{mark}] {r['call_id']} {r['true_type']} sev{r['true_severity']} "
+                f"response={r['response_time']:.1f}min outcome={r['outcome_score']:.2f}"
+            )
+        lines.append("")
+
+    # Stats footer
+    lines.append(
+        f"STATS: total={sim.total_calls()}  completed={len(sim.completed_calls)}  "
+        f"timed_out={len(sim.timed_out_calls)}  pending={len(pending_sorted)}"
+    )
+    if sim.episode_done:
+        lines.append("EPISODE: DONE")
+    return "\n".join(lines)

utils.py

@@ -0,0 +1,152 @@
+"""Helper functions for distance, ETA, lookup tables, caller text generation."""
+
+from __future__ import annotations
+
+import math
+from typing import List
+
+import numpy as np
+
+from models import (
+    EmergencyCall,
+    EmergencyType,
+    EmergencyUnit,
+    Position,
+    Severity,
+    UnitType,
+)
+
+# ----------------------------------------------------------------------------
+# Geometry
+# ----------------------------------------------------------------------------
+
+
+def calculate_distance(a: Position, b: Position) -> float:
+    """Euclidean distance in km."""
+    return math.sqrt((a.x - b.x) ** 2 + (a.y - b.y) ** 2)
+
+
+def calculate_eta(unit: EmergencyUnit, destination: Position) -> float:
+    """ETA in minutes assuming constant speed."""
+    distance_km = calculate_distance(unit.position, destination)
+    return (distance_km / unit.speed_kmh) * 60.0
+
+
+# ----------------------------------------------------------------------------
+# Capability lookups
+# ----------------------------------------------------------------------------
+
+CAPABLE_UNITS = {
+    EmergencyType.CARDIAC_ARREST: [UnitType.ALS_AMBULANCE, UnitType.BLS_AMBULANCE],
+    EmergencyType.TRAUMA: [UnitType.ALS_AMBULANCE, UnitType.BLS_AMBULANCE],
+    EmergencyType.STROKE: [UnitType.ALS_AMBULANCE, UnitType.BLS_AMBULANCE],
+    EmergencyType.FIRE: [UnitType.FIRE_ENGINE],
+    EmergencyType.BREATHING: [UnitType.ALS_AMBULANCE, UnitType.BLS_AMBULANCE],
+    EmergencyType.MINOR_INJURY: [UnitType.BLS_AMBULANCE, UnitType.ALS_AMBULANCE],
+    EmergencyType.MENTAL_HEALTH: [UnitType.POLICE, UnitType.ALS_AMBULANCE],
+}
+
+OPTIMAL_UNIT = {
+    EmergencyType.CARDIAC_ARREST: UnitType.ALS_AMBULANCE,
+    EmergencyType.TRAUMA: UnitType.ALS_AMBULANCE,
+    EmergencyType.STROKE: UnitType.ALS_AMBULANCE,
+    EmergencyType.FIRE: UnitType.FIRE_ENGINE,
+    EmergencyType.BREATHING: UnitType.ALS_AMBULANCE,
+    EmergencyType.MINOR_INJURY: UnitType.BLS_AMBULANCE,
+    EmergencyType.MENTAL_HEALTH: UnitType.POLICE,
+}
+
+
+def get_capable_units(emergency_type: EmergencyType) -> List[UnitType]:
+    return CAPABLE_UNITS[emergency_type]
+
+
+def get_optimal_unit(emergency_type: EmergencyType) -> UnitType:
+    return OPTIMAL_UNIT[emergency_type]
+
+
+# ----------------------------------------------------------------------------
+# Caller description templates
+# ----------------------------------------------------------------------------
+
+CALLER_TEMPLATES = {
+    EmergencyType.CARDIAC_ARREST: [
+        "Man collapsed on the sidewalk, he's not breathing! Please hurry!",
+        "My husband just grabbed his chest and fell down, he's turning blue!",
+        "Someone collapsed at the bus stop, no pulse, people are trying CPR!",
+        "Elderly woman found unconscious in her apartment, not responding!",
+    ],
+    EmergencyType.TRAUMA: [
+        "Car accident on the main road, driver is bleeding from the head!",
+        "Construction worker fell from scaffolding, conscious but can't move his legs!",
+        "Two-wheeler hit a pedestrian, person on the road with a broken leg!",
+        "Fight broke out, one person stabbed, lots of blood!",
+    ],
+    EmergencyType.STROKE: [
+        "My mother's face is drooping on one side and she can't speak properly!",
+        "Colleague suddenly can't move his right arm, speech is slurred!",
+        "Old man at the temple suddenly confused, can't walk straight!",
+    ],
+    EmergencyType.FIRE: [
+        "Kitchen fire in our apartment, smoke everywhere, we're on the 3rd floor!",
+        "Warehouse on fire near the industrial area, flames visible from outside!",
+        "Electrical fire in the office building, fire alarm going off!",
+        "Grass fire spreading toward houses near the park!",
+    ],
+    EmergencyType.MINOR_INJURY: [
+        "I slipped on the stairs and twisted my ankle, it's swollen.",
+        "Small cut on my hand from a knife, bleeding a bit but not too bad.",
+        "Child fell off bicycle, scraped knee and elbow, crying but alert.",
+        "Bumped my head on a low beam, small bump, feeling a bit dizzy.",
+    ],
+    EmergencyType.BREATHING: [
+        "My asthma is really bad, can't catch my breath, inhaler isn't working!",
+        "Allergic reaction, my throat is swelling up, hard to breathe!",
+        "Elderly patient with COPD, oxygen levels dropping, very short of breath!",
+    ],
+    EmergencyType.MENTAL_HEALTH: [
+        "My neighbor is threatening to hurt himself, he's very distressed!",
+        "Person on the building ledge, seems very agitated, won't come down!",
+        "Family member having a severe panic attack, can't calm down, hyperventilating!",
+    ],
+}
+
+MISREPORT_TEMPLATES = {
+    (EmergencyType.MENTAL_HEALTH, EmergencyType.CARDIAC_ARREST): [
+        "Someone is clutching their chest and can't breathe! I think it's a heart attack!",
+    ],
+    (EmergencyType.MINOR_INJURY, EmergencyType.TRAUMA): [
+        "Bad accident! Person is hurt, there's some blood!",
+    ],
+    (EmergencyType.BREATHING, EmergencyType.STROKE): [
+        "My father can't talk properly and is struggling, please come fast!",
+    ],
+    (EmergencyType.CARDIAC_ARREST, EmergencyType.MENTAL_HEALTH): [
+        "He's holding his chest and crying, very upset, hyperventilating!",
+    ],
+    (EmergencyType.TRAUMA, EmergencyType.MINOR_INJURY): [
+        "Someone fell off a ladder, looks like a small bump but they're complaining a lot.",
+    ],
+}
+
+
+def generate_caller_text(
+    true_type: EmergencyType,
+    reported_type: EmergencyType,
+    rng: np.random.RandomState,
+) -> str:
+    """Generate a caller description string from templates."""
+    if true_type != reported_type:
+        key = (true_type, reported_type)
+        if key in MISREPORT_TEMPLATES:
+            return rng.choice(MISREPORT_TEMPLATES[key])
+    return rng.choice(CALLER_TEMPLATES[reported_type])
+
+
+# ----------------------------------------------------------------------------
+# Severity helpers
+# ----------------------------------------------------------------------------
+
+
+def severity_from_int(s: int) -> Severity:
+    return Severity(int(s))