tasks.py

"""
tasks.py — Difficulty Configurations for TrafficEnv
=====================================================

Three pre-defined task configurations:

  EASY_CONFIG   – Stable, balanced traffic; good for initial training.
  MEDIUM_CONFIG – Random bursts, moderate congestion; standard benchmark.
  HARD_CONFIG   – High intensity, frequent emergencies, strict fairness.

Each config is a plain dict consumed by TrafficEnv.__init__().
"""

from __future__ import annotations
from typing import Any, Dict


# ---------------------------------------------------------------------------
# Easy
# ---------------------------------------------------------------------------

EASY_CONFIG: Dict[str, Any] = {
    # Traffic flow
    "arrival_rate":       (0, 1),    # 0–1 cars per lane per step
    "discharge_rate":     (4, 5),    # 4–5 cars discharged per green lane per step
    "max_queue":          15,        # queue cap per lane
    "max_steps":          50,

    # Emergencies — rare
    "emergency_prob":     0.01,

    # Bursts — none
    "burst_prob":         0.0,
    "burst_multiplier":   1.0,

    # Reward knobs
    "switch_penalty":         0.10,
    "starvation_threshold":   20,
    "r_efficiency_scale":     0.20,
    "p_congestion_scale":     0.30,
    "p_max_q_scale":          0.10,
    "p_starvation_scale":     0.10,
    "r_fairness_bonus":       0.05,
    "r_improvement_bonus":    0.15,
    "p_emergency_scale":      0.30,
    "r_ev_bonus_scale":       0.20,
    
    # Logic thresholds
    "ev_golden_window":       8,     # Easy: very generous window
    "ev_max_delay":           20,
}

# ---------------------------------------------------------------------------
# Medium
# ---------------------------------------------------------------------------

MEDIUM_CONFIG: Dict[str, Any] = {
    # Traffic flow
    "arrival_rate":       (1, 3),    # moderate, variable arrivals
    "discharge_rate":     (3, 5),    # standard discharge
    "max_queue":          25,
    "max_steps":          100,

    # Emergencies — occasional
    "emergency_prob":     0.05,

    # Random bursts — 10% chance, 1.5× arrivals
    "burst_prob":         0.10,
    "burst_multiplier":   1.5,

    # Reward knobs
    "switch_penalty":         0.20,
    "starvation_threshold":   15,
    "r_efficiency_scale":     0.20,
    "p_congestion_scale":     0.40,
    "p_max_q_scale":          0.15,
    "p_starvation_scale":     0.15,
    "r_fairness_bonus":       0.10,
    "r_improvement_bonus":    0.20,
    "p_emergency_scale":      0.40,
    "r_ev_bonus_scale":       0.25,

    # Logic thresholds
    "ev_golden_window":       5,     # Medium: standard window
    "ev_max_delay":           15,
}

# ---------------------------------------------------------------------------
# Hard
# ---------------------------------------------------------------------------

HARD_CONFIG: Dict[str, Any] = {
    # Traffic flow — high intensity
    "arrival_rate":       (2, 5),    # heavy, bursty arrivals
    "discharge_rate":     (2, 4),    # reduced discharge (lane friction)
    "max_queue":          40,
    "max_steps":          200,

    # Emergencies — frequent
    "emergency_prob":     0.15,

    # Frequent aggressive bursts
    "burst_prob":         0.20,
    "burst_multiplier":   2.0,

    # Reward knobs — stricter penalties
    "switch_penalty":         0.30,
    "starvation_threshold":   10,    # stricter fairness
    "r_efficiency_scale":     0.25,
    "p_congestion_scale":     0.50,
    "p_max_q_scale":          0.20,
    "p_starvation_scale":     0.20,
    "r_fairness_bonus":       0.15,
    "r_improvement_bonus":    0.25,
    "p_emergency_scale":      0.60,  # amplified emergency penalty
    "r_ev_bonus_scale":       0.30,

    # Logic thresholds
    "ev_golden_window":       3,     # Hard: must clear immediately
    "ev_max_delay":           10,
}


# ---------------------------------------------------------------------------
# Accessor
# ---------------------------------------------------------------------------

_CONFIGS = {
    "easy":   EASY_CONFIG,
    "medium": MEDIUM_CONFIG,
    "hard":   HARD_CONFIG,
}


def get_config(mode: str) -> Dict[str, Any]:
    """
    Return the config dict for the requested difficulty mode.

    Parameters
    ----------
    mode : str
        One of "easy", "medium", "hard" (case-insensitive).

    Returns
    -------
    dict
        Configuration dictionary suitable for ``TrafficEnv(config)``.

    Raises
    ------
    ValueError
        If an unknown mode is requested.
    """
    key = mode.strip().lower()
    if key not in _CONFIGS:
        raise ValueError(
            f"Unknown difficulty mode '{mode}'. "
            f"Choose one of: {list(_CONFIGS)}"
        )
    # Return a copy so callers can mutate without side-effects
    return dict(_CONFIGS[key])