server/traffic_environment.py

"""
server/traffic_environment.py — Core simulation engine.

This file contains the actual "world" — the physics of the simulation.
It is completely independent of HTTP/FastAPI. You can import and use it
directly in Python without running any server.

KEY FIX vs original: every episode now uses an ISOLATED random.Random
instance seeded from reset(seed=N). This means the same seed ALWAYS
produces the exact same episode, no matter what else is running.
That is what "reproducible scores" means in the requirements.

Three task levels:
  easy   — 1 intersection,  4 lanes, steady arrivals,    100 steps
  medium — 3 intersections, 4 lanes, rush-hour spike,    200 steps
  hard   — 9 intersections, 4 lanes, incidents + spikes, 300 steps
"""

import uuid
import random as _random_module   # alias so we don't shadow the local rng variable
import math
from typing import List, Optional, Tuple

# -- Dual-import pattern (required by OpenEnv) ----------
# Relative imports work when running inside the repo (PYTHONPATH=src:envs)
# Bare imports work inside Docker (PYTHONPATH=/app)
# Both paths must be tried or one environment will always fail.
try:
    from ..models import (
        TrafficAction, TrafficObservation, TrafficState,
        LaneState, IntersectionState, ActionType,
    )
except ImportError:
    from models import (
        TrafficAction, TrafficObservation, TrafficState,
        LaneState, IntersectionState, ActionType,
    )


# SIMULATION CONSTANTS
# Change these to make the world harder or easier.

PHASE_COUNT        = 4     # How many signal phases exist (N green, N yellow, S green, S yellow)
DEFAULT_PHASE_TIME = 30.0  # Seconds before auto-advancing to next phase (static baseline)
EXTEND_SECONDS     = 5.0   # How many seconds extend_green adds
STEP_DURATION      = 5.0   # Each env step = 5 simulated seconds
MAX_QUEUE          = 50    # Maximum cars that can queue in one lane
ARRIVAL_BASE       = 2.0   # Mean cars arriving per lane per step (Poisson mean)
DISCHARGE_RATE     = 3.5   # Max cars that clear a green lane per step



# LANE SIMULATION
# Models one lane of traffic (e.g. "North lane at intersection 0")

class LaneSim:
    """
    Internal mutable state for one lane.
    The agent never sees this object directly — it only sees LaneState
    (the Pydantic model built from this in to_model()).
    """

    def __init__(self, lane_id: str, arrival_mean: float, rng: _random_module.Random):
        self.lane_id      = lane_id
        self.arrival_mean = arrival_mean
        self._rng         = rng          # ISOLATED rng — not the global random module
        # Start with a small random queue so the episode isn't trivially easy at step 0
        self.queue:      int   = self._rng.randint(0, 5)
        self.total_wait: float = 0.0     # Cumulative vehicle-seconds of waiting

    def arrive(self, multiplier: float = 1.0) -> int:
        """
        Simulate cars arriving this step (Poisson-distributed).
        multiplier > 1.0 = rush hour or incident — more cars than usual.
        """
        mean = max(self.arrival_mean * multiplier, 0.1)
        # expovariate(1/mean) gives exponential distribution; rounding gives Poisson approx
        n = int(self._rng.expovariate(1.0 / mean) + 0.5)
        n = max(0, min(n, MAX_QUEUE - self.queue))
        self.queue = min(self.queue + n, MAX_QUEUE)
        return n

    def discharge(self, is_green: bool) -> int:
        """Cars that cross the stop-line this step. Only happens on green."""
        if not is_green or self.queue == 0:
            return 0
        discharged = min(self.queue, int(DISCHARGE_RATE))
        self.queue -= discharged
        return discharged

    def accumulate_wait(self):
        """Every step, every waiting car adds STEP_DURATION seconds to total wait."""
        self.total_wait += self.queue * STEP_DURATION

    @property
    def avg_wait(self) -> float:
        """Average wait time per car currently in the queue."""
        return round(self.total_wait / max(self.queue, 1), 2)

    @property
    def flow_rate(self) -> float:
        """Theoretical max vehicles per minute if this lane were always green."""
        return round(DISCHARGE_RATE * (60.0 / STEP_DURATION), 2)



# INTERSECTION SIMULATION
# Models one full intersection (4 lanes + signal controller)


class IntersectionSim:
    """
    Internal mutable state for one intersection.
    Contains 4 LaneSim objects and manages the signal phase.
    """

    # Which lane indices are green in each phase
    # Phase 0 = N+S green, Phase 1 = N only (yellow transition)
    # Phase 2 = E+W green, Phase 3 = E only (yellow transition)
    PHASE_MAP = {0: [0, 1], 1: [0], 2: [2, 3], 3: [2]}

    def __init__(self, intersection_id: int, n_lanes: int,
                 arrival_mean: float, rng: _random_module.Random):
        self.intersection_id = intersection_id
        self.current_phase   = 0
        self.phase_elapsed   = 0.0
        self._rng            = rng
        self.lanes: List[LaneSim] = [
            LaneSim(f"I{intersection_id}_L{i}", arrival_mean, rng)
            for i in range(n_lanes)
        ]

    def _green_lanes(self) -> List[int]:
        return self.PHASE_MAP.get(self.current_phase % PHASE_COUNT, [0])

    def apply_action(self, action: TrafficAction):
        """Apply agent's chosen action to this intersection's signal."""
        if action.action_type == ActionType.EXTEND_GREEN:
            # Push back the auto-advance clock by EXTEND_SECONDS
            self.phase_elapsed = max(0.0, self.phase_elapsed - EXTEND_SECONDS)
        elif action.action_type == ActionType.NEXT_PHASE:
            # Force immediate phase switch
            self.current_phase = (self.current_phase + 1) % PHASE_COUNT
            self.phase_elapsed = 0.0

    def step_sim(self, arrival_multiplier: float = 1.0) -> Tuple[int, float]:
        """
        Advance the simulation by one 5-second step.
        Returns (throughput, total_wait_added) for reward calculation.
        """
        green = self._green_lanes()
        throughput = 0

        for i, lane in enumerate(self.lanes):
            lane.arrive(arrival_multiplier)          # new cars arrive
            if i in green:
                throughput += lane.discharge(True)   # green lanes discharge
            lane.accumulate_wait()                   # all waiting cars accumulate wait

        # Advance phase clock; auto-switch if time exceeded
        self.phase_elapsed += STEP_DURATION
        if self.phase_elapsed >= DEFAULT_PHASE_TIME:
            self.current_phase = (self.current_phase + 1) % PHASE_COUNT
            self.phase_elapsed = 0.0

        total_wait = sum(l.queue * STEP_DURATION for l in self.lanes)
        return throughput, total_wait

    def to_model(self) -> IntersectionState:
        """Convert internal state → Pydantic model for the agent to observe."""
        return IntersectionState(
            intersection_id=self.intersection_id,
            current_phase=self.current_phase,
            phase_elapsed=round(self.phase_elapsed, 1),
            lanes=[
                LaneState(
                    lane_id=l.lane_id,
                    queue_length=l.queue,
                    avg_wait_time=l.avg_wait,
                    flow_rate=l.flow_rate,
                )
                for l in self.lanes
            ],
        )



# TRAFFIC ENVIRONMENT  (the OpenEnv interface)
# This is the class that implements reset() / step() / state()

class TrafficEnvironment:
    """
    OpenEnv-compatible environment for adaptive traffic signal control.

    Usage:
        env = TrafficEnvironment(task_level="easy")
        obs = env.reset(seed=42)
        while not obs.done:
            action = agent.decide(obs)
            obs = env.step(action)
        print(env.state.cumulative_reward)
    """

    TASK_CONFIG = {
        "easy":   {"n_intersections": 1, "n_lanes": 4, "max_steps": 100,
                   "arrival_mean": ARRIVAL_BASE},
        "medium": {"n_intersections": 3, "n_lanes": 4, "max_steps": 200,
                   "arrival_mean": ARRIVAL_BASE * 1.5},
        "hard":   {"n_intersections": 9, "n_lanes": 4, "max_steps": 300,
                   "arrival_mean": ARRIVAL_BASE * 2.0},
    }

    def __init__(self, task_level: str = "easy"):
        assert task_level in self.TASK_CONFIG, \
            f"task_level must be one of {list(self.TASK_CONFIG)}"
        self.task_level          = task_level
        self._cfg                = self.TASK_CONFIG[task_level]
        self._intersections:     List[IntersectionSim] = []
        self._state              = TrafficState(task_level=task_level)
        self._step               = 0
        self._cumulative_reward  = 0.0
        self._cumulative_throughput = 0
        self._cumulative_wait    = 0.0
        self._rng                = _random_module.Random()  # will be re-seeded in reset()

    # -- OpenEnv required method 1 -------

    def reset(self, seed: Optional[int] = None) -> TrafficObservation:
        """
        Start a fresh episode.
        seed=N guarantees identical episode every time — required for reproducibility.
        """
        # Create a NEW isolated RNG for this episode
        # This means global random state is never touched → fully reproducible
        self._rng = _random_module.Random(seed)

        cfg = self._cfg
        self._intersections = [
            IntersectionSim(i, cfg["n_lanes"], cfg["arrival_mean"], self._rng)
            for i in range(cfg["n_intersections"])
        ]
        self._step                  = 0
        self._cumulative_reward     = 0.0
        self._cumulative_throughput = 0
        self._cumulative_wait       = 0.0
        self._state = TrafficState(
            episode_id=str(uuid.uuid4()),
            step_count=0,
            task_level=self.task_level,
        )
        return self._build_obs(total_throughput=0, total_wait=0.0, reward=0.0, done=False)

    #  OpenEnv required method --------

    def step(self, action: TrafficAction) -> TrafficObservation:
        """
        Apply action, advance simulation by 5 seconds, return new observation + reward.
        """
        self._step += 1
        arrival_mult = self._arrival_multiplier()

        # Apply action to the targeted intersection
        if 0 <= action.intersection_id < len(self._intersections):
            self._intersections[action.intersection_id].apply_action(action)

        # Advance ALL intersections (even uncontrolled ones keep ticking)
        total_throughput = 0
        total_wait       = 0.0
        for inter in self._intersections:
            tp, wt = inter.step_sim(arrival_mult)
            total_throughput += tp
            total_wait       += wt

        reward = self._compute_reward(total_throughput, total_wait)
        self._cumulative_reward     += reward
        self._cumulative_throughput += total_throughput
        self._cumulative_wait       += total_wait

        done = (self._step >= self._cfg["max_steps"])

        # Update episode metadata
        self._state = TrafficState(
            episode_id=self._state.episode_id,
            step_count=self._step,
            task_level=self.task_level,
            cumulative_reward=round(self._cumulative_reward, 4),
            cumulative_throughput=self._cumulative_throughput,
            cumulative_wait=round(self._cumulative_wait, 2),
        )
        return self._build_obs(total_throughput, total_wait, reward, done)

    # ---- OpenEnv required method 3 --------

    @property
    def state(self) -> TrafficState:
        """Episode-level metadata (step count, cumulative reward, etc.)"""
        return self._state

    # ---- Internal helpers -----------

    def _arrival_multiplier(self) -> float:
        """
        Rush-hour spike: arrival rate peaks mid-episode for medium/hard.
        Hard additionally has random incidents (5% chance per step).
        """
        if self.task_level == "easy":
            return 1.0
        peak  = 50 if self.task_level == "medium" else 80
        spike = 1.0 + 1.5 * math.exp(-((self._step - peak) ** 2) / (2 * 20 ** 2))
        if self.task_level == "hard" and self._rng.random() < 0.05:
            spike += self._rng.uniform(0.5, 2.0)
        return round(spike, 3)

    def _compute_reward(self, throughput: int, total_wait: float) -> float:
        """
        Dense reward fired every single step (not just at the end).

        Formula:
            reward = 0.6 × throughput_bonus + 0.4 × wait_penalty

        throughput_bonus ∈ [0, 1]  — reward clearing cars
        wait_penalty     ∈ [-1, 0] — penalise idle waiting

        Both terms are normalised by the theoretical maximum so the
        reward stays in roughly [-1, 1] regardless of task size.
        """
        n = len(self._intersections) * self._cfg["n_lanes"]
        throughput_bonus = throughput / (DISCHARGE_RATE * n + 1e-9)
        wait_penalty     = -total_wait / (MAX_QUEUE * STEP_DURATION * n + 1e-9)
        return round(0.6 * throughput_bonus + 0.4 * wait_penalty, 6)

    def _build_obs(self, total_throughput: int, total_wait: float,
                   reward: float, done: bool) -> TrafficObservation:
        """Assemble the Pydantic observation object from current simulation state."""
        total_waiting = sum(l.queue for inter in self._intersections for l in inter.lanes)
        avg_wait = total_wait / max(total_waiting * STEP_DURATION, 1e-9)
        return TrafficObservation(
            intersections=[i.to_model() for i in self._intersections],
            total_waiting_vehicles=total_waiting,
            total_avg_wait=round(avg_wait, 2),
            throughput_last_step=total_throughput,
            reward=reward,
            done=done,
            info={
                "step":               float(self._step),
                "arrival_multiplier": self._arrival_multiplier(),
                "cumulative_reward":  round(self._cumulative_reward, 4),
            },
        )