tasks.py

from copy import deepcopy
from typing import Dict

from models import TaskSpec


TASKS: Dict[str, TaskSpec] = {
    "single_intersection": TaskSpec(
        id="single_intersection",
        name="Single Intersection Control",
        difficulty="easy",
        description="Optimize a mixed-traffic four-way Indian urban intersection.",
        constraints={
            "max_steps": 120,
            "min_green_time": 3,
            "max_queue_before_failure": 220,
            "initial_peak_multiplier": 1.0,
            "emergency_rate": 0.02,
            "rain_level": 0.15,
        },
        reward_weights={
            "vehicles_cleared": 3.0,
            "total_waiting_time": -0.2,
            "queue_length": -0.5,
            "pedestrian_wait_time": -0.3,
            "unsafe_switch_penalty": -5.0,
            "emergency_clear_bonus": 20.0,
        },
        termination={"target_cleared": 350, "allow_gridlock": False},
    ),
    "rush_hour": TaskSpec(
        id="rush_hour",
        name="Rush Hour Traffic Management",
        difficulty="medium",
        description="Handle asymmetric commuter inflow with rain and heavy two-wheeler traffic.",
        constraints={
            "max_steps": 160,
            "min_green_time": 4,
            "max_queue_before_failure": 300,
            "initial_peak_multiplier": 1.65,
            "emergency_rate": 0.015,
            "rain_level": 0.35,
        },
        reward_weights={
            "vehicles_cleared": 3.2,
            "total_waiting_time": -0.24,
            "queue_length": -0.55,
            "pedestrian_wait_time": -0.25,
            "unsafe_switch_penalty": -5.5,
            "emergency_clear_bonus": 18.0,
        },
        termination={"target_cleared": 520, "allow_gridlock": False},
    ),
    "emergency_priority": TaskSpec(
        id="emergency_priority",
        name="Emergency Vehicle Prioritization",
        difficulty="hard",
        description="Prioritize ambulances and fire trucks without letting the rest of the junction fail.",
        constraints={
            "max_steps": 140,
            "min_green_time": 3,
            "max_queue_before_failure": 260,
            "initial_peak_multiplier": 1.25,
            "emergency_rate": 0.08,
            "rain_level": 0.25,
        },
        reward_weights={
            "vehicles_cleared": 2.7,
            "total_waiting_time": -0.2,
            "queue_length": -0.45,
            "pedestrian_wait_time": -0.25,
            "unsafe_switch_penalty": -6.5,
            "emergency_clear_bonus": 28.0,
        },
        termination={"target_cleared": 420, "max_emergency_wait": 18, "allow_gridlock": False},
    ),
}


def get_task(task_id: str) -> TaskSpec:
    if task_id not in TASKS:
        raise KeyError(f"Unknown task_id '{task_id}'. Available tasks: {', '.join(TASKS)}")
    return deepcopy(TASKS[task_id])


def list_tasks():
    return [deepcopy(task) for task in TASKS.values()]