server/task_registry.py

from __future__ import annotations

from server.models import DisasterType, RewardWeights, TaskSpec, TaskSpecPublic


TASKS: dict[str, TaskSpec] = {
    "task_1_fire_easy": TaskSpec(
        task_id="task_1_fire_easy",
        name="Single Fire Evacuation",
        difficulty="easy",
        disaster_type=DisasterType.fire,
        building_profile="small_3floor",
        success_criteria="Route all 15 civilians to ground exits",
        goal="ground_exit",
        max_steps=30,
        evaluation_seeds=[42, 123, 456],
        description="Fire starts in one room on floor 2. Spreads slowly. 3 floors, 2 stairwells, 2 ground exits. 15 civilians, no injured.",
        expected_score_range=[0.95, 0.999],
        reward_weights=RewardWeights(),
    ),
    "task_2_flood_medium": TaskSpec(
        task_id="task_2_flood_medium",
        name="Flood Rising Rooftop Evacuation",
        difficulty="medium",
        disaster_type=DisasterType.flood,
        building_profile="medium_5floor",
        success_criteria="Route civilians to rooftop before flood cuts stairwells",
        goal="rooftop",
        max_steps=40,
        evaluation_seeds=[42, 123, 456],
        description="Flood rising from ground floor. 5 floors, 3 stairwells (1 blocked), 1 rooftop exit. 30 civilians, 4 injured.",
        expected_score_range=[0.9, 0.949],
        reward_weights=RewardWeights(),
    ),
    "task_3_earthquake_hard": TaskSpec(
        task_id="task_3_earthquake_hard",
        name="Post-Earthquake Structural Evacuation",
        difficulty="hard",
        disaster_type=DisasterType.structural,
        building_profile="complex_5floor",
        success_criteria="Route civilians to the nearest safe exit before collapses cut routes",
        goal="nearest_exit",
        max_steps=50,
        evaluation_seeds=[42, 123, 456],
        description="Earthquake damages a 5-floor building and triggers progressive collapses. 4 stairwells plus same-floor fire-escape egress on selected floors, 50 civilians, 10 injured, 3 mobility-impaired.",
        expected_score_range=[0.8, 0.899],
        reward_weights=RewardWeights(),
    ),
    "task_4_cascade_hard": TaskSpec(
        task_id="task_4_cascade_hard",
        name="Multi-Hazard Cascade",
        difficulty="expert",
        disaster_type=DisasterType.multi_cascade,
        building_profile="complex_5floor_full",
        success_criteria="Evacuate maximum civilians across all exit types",
        goal="maximum_survival",
        max_steps=60,
        evaluation_seeds=[42, 123, 456],
        description="Fire on floor 1, gas rupture on floor 3 at step 10, stairwell collapse at step 15. 5 floors, full complexity, external fire-escape egress on selected floors. 60 civilians, mixed mobility, panic mechanics.",
        expected_score_range=[0.75, 0.899],
        reward_weights=RewardWeights(),
    ),
}


def get_task(task_id: str) -> TaskSpec:
    if task_id not in TASKS:
        raise ValueError(f"Unknown task: {task_id}. Available: {list(TASKS.keys())}")
    return TASKS[task_id]


def get_all_tasks() -> list[TaskSpec]:
    return list(TASKS.values())


def get_tasks_public() -> list[TaskSpecPublic]:
    """Return public-facing task info (no internal weights)."""
    return [
        TaskSpecPublic(
            task_id=task.task_id,
            name=task.name,
            difficulty=task.difficulty,
            disaster_type=task.disaster_type,
            goal=task.goal,
            description=task.description,
            max_steps=task.max_steps,
            expected_score_range=task.expected_score_range,
        )
        for task in TASKS.values()
    ]