models.py

from enum import Enum
from typing import Dict, List, Optional

from pydantic import BaseModel, Field


VEHICLE_TYPES = ["cars", "bikes", "autos", "buses", "trucks"]
LANES = ["N", "S", "E", "W"]


class TrafficAction(str, Enum):
    NS_GREEN = "NS_GREEN"
    EW_GREEN = "EW_GREEN"
    LEFT_PRIORITY = "LEFT_PRIORITY"
    PEDESTRIAN_CROSS = "PEDESTRIAN_CROSS"
    EXTEND_GREEN = "EXTEND_GREEN"
    EMERGENCY_OVERRIDE = "EMERGENCY_OVERRIDE"
    ALL_RED = "ALL_RED"


class LaneQueues(BaseModel):
    cars: int = Field(ge=0)
    bikes: int = Field(ge=0)
    autos: int = Field(ge=0)
    buses: int = Field(ge=0)
    trucks: int = Field(ge=0)

    @property
    def total(self) -> int:
        return self.cars + self.bikes + self.autos + self.buses + self.trucks


class EmergencyVehicle(BaseModel):
    present: bool
    lane: Optional[str] = Field(default=None)
    type: Optional[str] = Field(default=None)
    wait_time: float = Field(default=0.0, ge=0.0)


class TrafficState(BaseModel):
    tick: int = Field(ge=0)
    lane_queues: Dict[str, LaneQueues]
    lane_waiting_time: Dict[str, float]
    current_signal_phase: TrafficAction
    time_since_last_phase_switch: int = Field(ge=0)
    pedestrian_count: int = Field(ge=0)
    pedestrian_wait_time: float = Field(ge=0.0)
    emergency_vehicle: EmergencyVehicle
    rain_level: float = Field(ge=0.0, le=1.0)
    random_traffic_inflow: Dict[str, LaneQueues]


class StepRequest(BaseModel):
    action: TrafficAction


class ResetRequest(BaseModel):
    seed: Optional[int] = 42
    task_id: str = "single_intersection"


class StepResult(BaseModel):
    observation: TrafficState
    reward: float = Field(ge=0.0, le=1.0)
    done: bool
    info: Dict[str, object]


class TaskSpec(BaseModel):
    id: str
    name: str
    difficulty: str
    description: str
    constraints: Dict[str, object]
    reward_weights: Dict[str, float]
    termination: Dict[str, object]


class GraderRequest(BaseModel):
    task_id: str = "single_intersection"
    seed: int = 42
    actions: Optional[List[TrafficAction]] = None
    max_steps: Optional[int] = None


class GraderOutput(BaseModel):
    score: float = Field(ge=0.0, le=1.0)
    average_waiting_time: float
    max_queue_length: int
    total_vehicles_cleared: int
    emergency_handling_efficiency: float = Field(ge=0.0, le=1.0)
    details: Dict[str, object]


class BaselineOutput(BaseModel):
    task_id: str
    seed: int
    score: float = Field(ge=0.0, le=1.0)
    total_reward: float
    steps: int
    grader: GraderOutput