env.py

import random
from dataclasses import dataclass, field
from typing import Dict, Iterable, Optional, Tuple

from models import EmergencyVehicle, LANES, VEHICLE_TYPES, LaneQueues, TrafficAction, TrafficState
from tasks import get_task


LANE_GROUPS = {
    TrafficAction.NS_GREEN: ("N", "S"),
    TrafficAction.EW_GREEN: ("E", "W"),
    TrafficAction.LEFT_PRIORITY: ("N", "E"),
}

VEHICLE_CLEARANCE_UNITS = {
    "cars": 1.0,
    "bikes": 0.45,
    "autos": 0.75,
    "buses": 2.2,
    "trucks": 2.5,
}


@dataclass
class TrafficMetrics:
    total_vehicles_cleared: int = 0
    total_wait_observations: float = 0.0
    wait_samples: int = 0
    max_queue_length: int = 0
    emergency_seen: int = 0
    emergency_cleared_fast: int = 0
    unsafe_switches: int = 0
    full_clearances: int = 0


@dataclass
class StepDiagnostics:
    vehicles_cleared: int = 0
    unsafe_switch_penalty: int = 0
    emergency_clear_bonus: int = 0
    raw_reward: float = 0.0
    normalized_reward: float = 0.0
    invalid_action: bool = False
    notes: list = field(default_factory=list)


class IndianTrafficEnv:
    """Deterministic mixed-traffic signal environment for one Indian urban intersection.

    The class is intentionally network-ready: lane state and flow calculations are kept
    per-lane, so a future MultiIntersectionEnv can compose several instances.
    """

    def __init__(self, task_id: str = "single_intersection"):
        self.task_id = task_id
        self.task = get_task(task_id)
        self.rng = random.Random()
        self.seed_value = 42
        self.tick = 0
        self.current_phase = TrafficAction.ALL_RED
        self.time_since_switch = 0
        self.queues: Dict[str, Dict[str, int]] = {}
        self.waiting_time: Dict[str, float] = {}
        self.last_inflow: Dict[str, Dict[str, int]] = {}
        self.pedestrian_count = 0
        self.pedestrian_wait_time = 0.0
        self.emergency = EmergencyVehicle(present=False)
        self.rain_level = float(self.task.constraints["rain_level"])
        self.driver_aggression = 0.4
        self.random_blockage_probability = 0.02
        self.peak_hour_multiplier = float(self.task.constraints["initial_peak_multiplier"])
        self.metrics = TrafficMetrics()
        self.done = False
        self.reset(self.seed_value, task_id)

    def reset(self, seed: Optional[int] = None, task_id: Optional[str] = None) -> TrafficState:
        if task_id is not None:
            self.task_id = task_id
            self.task = get_task(task_id)
        self.seed_value = 42 if seed is None else int(seed)
        self.rng = random.Random(self.seed_value)
        self.tick = 0
        self.current_phase = TrafficAction.ALL_RED
        self.time_since_switch = 0
        self.rain_level = float(self.task.constraints["rain_level"])
        self.driver_aggression = self.rng.uniform(0.22, 0.72)
        self.random_blockage_probability = self.rng.uniform(0.01, 0.05)
        self.peak_hour_multiplier = float(self.task.constraints["initial_peak_multiplier"])
        self.queues = {lane: self._sample_initial_lane(lane) for lane in LANES}
        self.waiting_time = {lane: 0.0 for lane in LANES}
        self.last_inflow = {lane: {vehicle: 0 for vehicle in VEHICLE_TYPES} for lane in LANES}
        self.pedestrian_count = self.rng.randint(2, 14)
        self.pedestrian_wait_time = 0.0
        self.emergency = EmergencyVehicle(present=False)
        self.metrics = TrafficMetrics()
        self.done = False
        return self.get_state()

    def step(self, action: TrafficAction) -> Tuple[TrafficState, float, bool, Dict[str, object]]:
        if self.done:
            return self.get_state(), 0.0, True, {"message": "Episode already complete."}

        diagnostics = StepDiagnostics()
        try:
            action = TrafficAction(action)
        except ValueError:
            action = TrafficAction.ALL_RED
            diagnostics.invalid_action = True
            diagnostics.notes.append("Invalid action converted to ALL_RED.")

        effective_action = self._apply_action_constraints(action, diagnostics)
        self._maybe_switch_phase(effective_action)

        self.last_inflow = self._generate_inflow()
        self._add_inflow(self.last_inflow)
        self._maybe_spawn_emergency()
        diagnostics.vehicles_cleared = self._clear_traffic(effective_action)
        emergency_cleared = self._update_emergency(effective_action)
        diagnostics.emergency_clear_bonus = 1 if emergency_cleared else 0
        self._update_pedestrians(effective_action)
        self._update_waiting_times()

        queue_length = self._total_queue_length()
        total_wait = sum(self.waiting_time.values())
        self.metrics.total_vehicles_cleared += diagnostics.vehicles_cleared
        self.metrics.total_wait_observations += total_wait
        self.metrics.wait_samples += 1
        self.metrics.max_queue_length = max(self.metrics.max_queue_length, queue_length)
        if queue_length == 0:
            self.metrics.full_clearances += 1

        reward = self._calculate_reward(
            vehicles_cleared=diagnostics.vehicles_cleared,
            total_waiting_time=total_wait,
            queue_length=queue_length,
            pedestrian_wait_time=self.pedestrian_wait_time,
            unsafe_switch_penalty=diagnostics.unsafe_switch_penalty,
            emergency_clear_bonus=diagnostics.emergency_clear_bonus,
            diagnostics=diagnostics,
        )

        self.tick += 1
        self.time_since_switch += 1
        self.done = self._is_done(queue_length)
        return self.get_state(), reward, self.done, self._info(diagnostics)

    def get_state(self) -> TrafficState:
        return TrafficState(
            tick=self.tick,
            lane_queues={lane: LaneQueues(**self.queues[lane]) for lane in LANES},
            lane_waiting_time={lane: round(self.waiting_time[lane], 3) for lane in LANES},
            current_signal_phase=self.current_phase,
            time_since_last_phase_switch=self.time_since_switch,
            pedestrian_count=self.pedestrian_count,
            pedestrian_wait_time=round(self.pedestrian_wait_time, 3),
            emergency_vehicle=self.emergency,
            rain_level=round(self.rain_level, 3),
            random_traffic_inflow={lane: LaneQueues(**self.last_inflow[lane]) for lane in LANES},
        )

    def _sample_initial_lane(self, lane: str) -> Dict[str, int]:
        arterial_bonus = 3 if lane in ("N", "S") else 1
        return {
            "cars": self.rng.randint(3 + arterial_bonus, 8 + arterial_bonus),
            "bikes": self.rng.randint(6 + arterial_bonus, 14 + arterial_bonus),
            "autos": self.rng.randint(2, 6),
            "buses": self.rng.randint(0, 2),
            "trucks": self.rng.randint(0, 2),
        }

    def _generate_inflow(self) -> Dict[str, Dict[str, int]]:
        inflow = {}
        rain_slowdown = 1.0 + self.rain_level * 0.25
        for lane in LANES:
            arterial = 1.25 if lane in ("N", "S") else 0.95
            base = self.peak_hour_multiplier * arterial * rain_slowdown
            inflow[lane] = {
                "cars": self._bounded_arrivals(base, 2),
                "bikes": self._bounded_arrivals(base * 1.8, 4),
                "autos": self._bounded_arrivals(base * 0.85, 2),
                "buses": 1 if self.rng.random() < 0.08 * base else 0,
                "trucks": 1 if self.rng.random() < 0.06 * base else 0,
            }
        if self.tick % 40 == 0 and self.tick > 0:
            self.peak_hour_multiplier = max(0.8, self.peak_hour_multiplier * 0.94)
        return inflow

    def _bounded_arrivals(self, intensity: float, cap: int) -> int:
        arrivals = int(intensity)
        fractional = intensity - arrivals
        if self.rng.random() < fractional:
            arrivals += 1
        if self.rng.random() < 0.12 * self.peak_hour_multiplier:
            arrivals += 1
        return max(0, min(cap, arrivals))

    def _add_inflow(self, inflow: Dict[str, Dict[str, int]]) -> None:
        for lane in LANES:
            for vehicle in VEHICLE_TYPES:
                self.queues[lane][vehicle] += inflow[lane][vehicle]

    def _apply_action_constraints(self, action: TrafficAction, diagnostics: StepDiagnostics) -> TrafficAction:
        min_green = int(self.task.constraints["min_green_time"])
        switching = action != self.current_phase and action not in (TrafficAction.EXTEND_GREEN, TrafficAction.EMERGENCY_OVERRIDE)
        green_to_green = self.current_phase in (TrafficAction.NS_GREEN, TrafficAction.EW_GREEN, TrafficAction.LEFT_PRIORITY)
        wants_green = action in (TrafficAction.NS_GREEN, TrafficAction.EW_GREEN, TrafficAction.LEFT_PRIORITY)
        if switching and green_to_green and wants_green and self.time_since_switch < min_green:
            diagnostics.unsafe_switch_penalty += 1
            diagnostics.notes.append("Minimum green-time violation.")
        if action == TrafficAction.PEDESTRIAN_CROSS and self.pedestrian_count == 0:
            diagnostics.unsafe_switch_penalty += 1
            diagnostics.notes.append("Pedestrian phase requested without demand.")
        if action == TrafficAction.EMERGENCY_OVERRIDE and not self.emergency.present:
            diagnostics.unsafe_switch_penalty += 1
            diagnostics.notes.append("Emergency override requested without emergency vehicle.")
        if action == TrafficAction.EXTEND_GREEN:
            if self.current_phase in (TrafficAction.NS_GREEN, TrafficAction.EW_GREEN, TrafficAction.LEFT_PRIORITY):
                return self.current_phase
            diagnostics.notes.append("EXTEND_GREEN from non-green phase converted to ALL_RED.")
            return TrafficAction.ALL_RED
        return action

    def _maybe_switch_phase(self, action: TrafficAction) -> None:
        if action != self.current_phase:
            self.current_phase = action
            self.time_since_switch = 0

    def _clear_traffic(self, action: TrafficAction) -> int:
        if action == TrafficAction.ALL_RED or action == TrafficAction.PEDESTRIAN_CROSS:
            return 0
        lanes = self._active_lanes(action)
        if not lanes:
            return 0

        cleared = 0
        rain_factor = 1.0 - self.rain_level * 0.32
        blockage_factor = 0.45 if self.rng.random() < self.random_blockage_probability else 1.0
        aggression_bonus = 1.0 + self.driver_aggression * 0.18
        capacity_units = 8.5 * rain_factor * blockage_factor * aggression_bonus
        if action == TrafficAction.LEFT_PRIORITY:
            capacity_units *= 0.65
        if action == TrafficAction.EMERGENCY_OVERRIDE:
            capacity_units *= 1.15

        for lane in lanes:
            remaining_units = capacity_units
            for vehicle in ("bikes", "autos", "cars", "buses", "trucks"):
                cleared_count, remaining_units = self._clear_vehicle_type(lane, vehicle, remaining_units)
                cleared += cleared_count
            if self.queues[lane] and sum(self.queues[lane].values()) == 0:
                self.waiting_time[lane] = 0.0
        return cleared

    def _clear_vehicle_type(self, lane: str, vehicle: str, remaining_units: float) -> Tuple[int, float]:
        unit = VEHICLE_CLEARANCE_UNITS[vehicle]
        possible = int(remaining_units // unit)
        count = min(self.queues[lane][vehicle], possible)
        self.queues[lane][vehicle] -= count
        return count, remaining_units - count * unit

    def _active_lanes(self, action: TrafficAction) -> Iterable[str]:
        if action == TrafficAction.EMERGENCY_OVERRIDE and self.emergency.present and self.emergency.lane:
            return (self.emergency.lane,)
        return LANE_GROUPS.get(action, ())

    def _maybe_spawn_emergency(self) -> None:
        if self.emergency.present:
            return
        if self.rng.random() < float(self.task.constraints["emergency_rate"]):
            self.emergency = EmergencyVehicle(
                present=True,
                lane=self.rng.choice(LANES),
                type=self.rng.choice(["ambulance", "fire_truck"]),
                wait_time=0.0,
            )
            self.metrics.emergency_seen += 1

    def _update_emergency(self, action: TrafficAction) -> bool:
        if not self.emergency.present or not self.emergency.lane:
            return False
        active = set(self._active_lanes(action))
        if self.emergency.lane in active:
            if action == TrafficAction.EMERGENCY_OVERRIDE or self.rng.random() < 0.55:
                if self.emergency.wait_time <= 8:
                    self.metrics.emergency_cleared_fast += 1
                self.emergency = EmergencyVehicle(present=False)
                return True
        self.emergency.wait_time += 1.0
        return False

    def _update_pedestrians(self, action: TrafficAction) -> None:
        arrivals = 1 if self.rng.random() < 0.35 else 0
        if self.rng.random() < 0.08 * self.peak_hour_multiplier:
            arrivals += self.rng.randint(1, 3)
        self.pedestrian_count += arrivals
        if action == TrafficAction.PEDESTRIAN_CROSS:
            crossed = min(self.pedestrian_count, 18)
            self.pedestrian_count -= crossed
            if self.pedestrian_count == 0:
                self.pedestrian_wait_time = 0.0
            return
        if self.pedestrian_count:
            unsafe_crossing_pressure = self.driver_aggression * self.pedestrian_wait_time / 60.0
            if self.rng.random() < unsafe_crossing_pressure:
                self.pedestrian_count = max(0, self.pedestrian_count - 1)
            self.pedestrian_wait_time += self.pedestrian_count * 0.55

    def _update_waiting_times(self) -> None:
        active = set(self._active_lanes(self.current_phase))
        for lane in LANES:
            queue = sum(self.queues[lane].values())
            if queue == 0:
                self.waiting_time[lane] = 0.0
                continue
            pressure = queue * (1.0 + self.rain_level * 0.3)
            if lane in active:
                pressure *= 0.35
            self.waiting_time[lane] += pressure

    def _calculate_reward(
        self,
        vehicles_cleared: int,
        total_waiting_time: float,
        queue_length: int,
        pedestrian_wait_time: float,
        unsafe_switch_penalty: int,
        emergency_clear_bonus: int,
        diagnostics: StepDiagnostics,
    ) -> float:
        weights = self.task.reward_weights
        raw = (
            vehicles_cleared * weights["vehicles_cleared"]
            + total_waiting_time * weights["total_waiting_time"]
            + queue_length * weights["queue_length"]
            + pedestrian_wait_time * weights["pedestrian_wait_time"]
            + unsafe_switch_penalty * weights["unsafe_switch_penalty"]
            + emergency_clear_bonus * weights["emergency_clear_bonus"]
        )
        if queue_length == 0:
            raw += 10.0
            diagnostics.notes.append("Full clearance milestone.")
        if diagnostics.invalid_action:
            raw -= 8.0

        diagnostics.raw_reward = raw
        normalized = max(0.0, min(1.0, (raw + 180.0) / 260.0))
        diagnostics.normalized_reward = normalized
        if unsafe_switch_penalty:
            self.metrics.unsafe_switches += unsafe_switch_penalty
        return normalized

    def _total_queue_length(self) -> int:
        return sum(sum(queue.values()) for queue in self.queues.values())

    def _is_done(self, queue_length: int) -> bool:
        if self.tick + 1 >= int(self.task.constraints["max_steps"]):
            return True
        if self.metrics.total_vehicles_cleared >= int(self.task.termination["target_cleared"]):
            return True
        if queue_length >= int(self.task.constraints["max_queue_before_failure"]):
            return True
        max_emergency_wait = self.task.termination.get("max_emergency_wait")
        if max_emergency_wait and self.emergency.present and self.emergency.wait_time > float(max_emergency_wait):
            return True
        return False

    def _info(self, diagnostics: StepDiagnostics) -> Dict[str, object]:
        total_wait = sum(self.waiting_time.values())
        return {
            "task_id": self.task_id,
            "vehicles_cleared": diagnostics.vehicles_cleared,
            "total_vehicles_cleared": self.metrics.total_vehicles_cleared,
            "total_waiting_time": round(total_wait, 3),
            "queue_length": self._total_queue_length(),
            "unsafe_switch_penalty": diagnostics.unsafe_switch_penalty,
            "emergency_clear_bonus": diagnostics.emergency_clear_bonus,
            "raw_reward": round(diagnostics.raw_reward, 3),
            "normalized_reward": round(diagnostics.normalized_reward, 3),
            "hidden_dynamics": {
                "driver_aggression": round(self.driver_aggression, 3),
                "random_blockage_probability": round(self.random_blockage_probability, 3),
                "peak_hour_multiplier": round(self.peak_hour_multiplier, 3),
            },
            "notes": diagnostics.notes,
        }