uncertainty_aware_trajectory_model

"""
SafePathAI: Advanced Uncertainty-Aware Trajectory Prediction for Safe Autonomous Driving
Complete implementation with advanced features for research and real-world applications
"""

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
import math
import pandas as pd
import random
from collections import defaultdict
from typing import Dict, List, Tuple, Optional, Union, Callable, Any, Set
import copy
import tqdm
import time
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.cluster import KMeans
from sklearn.metrics import r2_score
from scipy.stats import multivariate_normal
import networkx as nx

# For normalizing flows
try:
    import pyro
    import pyro.distributions as dist
    import pyro.distributions.transforms as T
    from pyro.nn import DenseNN
    PYRO_AVAILABLE = True
except ImportError:
    PYRO_AVAILABLE = False
    print("Warning: Pyro not available. Normalizing flows will not be available.")

# For multi-GPU training
try:
    import horovod.torch as hvd
    HOROVOD_AVAILABLE = True
except ImportError:
    HOROVOD_AVAILABLE = False
    print("Warning: Horovod not available. Multi-GPU training will not be available.")

# For tensor visualizations
try:
    from tensorboardX import SummaryWriter
    TB_AVAILABLE = True
except ImportError:
    TB_AVAILABLE = False
    print("Warning: TensorboardX not available. Visualizations will be limited.")

#######################
# CONFIG
#######################

class Config:
    """Advanced configuration for the SafePathAI model."""
    
    def __init__(self):
        # Data parameters
        self.input_seq_len = 20  # Length of input trajectory sequence
        self.pred_seq_len = 30   # Length of predicted trajectory sequence
        self.input_dim = 4       # Base features (x, y, vx, vy)
        self.output_dim = 2      # Position prediction (x, y)
        self.map_features = 64   # Number of map features
        self.agent_features = 32  # Features for other agents
        self.weather_features = 8  # Weather condition features
        self.time_features = 4    # Time of day features
        self.intention_classes = 6  # Number of intention classes
        
        # Dataset parameters
        self.dataset_type = "nuScenes"  # Options: "nuScenes", "Argoverse", "Waymo", "Lyft"
        self.data_path = "./data"
        self.use_map_data = True
        self.use_traffic_rules = True
        self.use_weather_data = True
        self.use_time_data = True
        self.random_rotation = True    # Apply random rotations for data augmentation
        self.random_noise = True       # Add random noise for data augmentation
        self.max_agents = 32           # Maximum number of agents to consider
        self.max_neighbor_dist = 50.0  # Maximum distance to consider agents as neighbors (meters)
        
        # Model parameters
        self.model_type = "transformer"  # Options: "lstm", "transformer", "gnn", "social_gan", "scene_transformer"
        self.hidden_size = 256
        self.embedding_size = 128
        self.num_layers = 4
        self.num_heads = 8
        self.dropout = 0.1
        self.mc_dropout_samples = 50
        self.ensemble_size = 5
        self.num_modes = 6       # Number of prediction modes (for multimodal prediction)
        self.use_attention = True
        self.use_scene_graph = True
        self.use_map_encoder = True
        self.use_social_encoder = True
        self.use_multimodal = True
        self.use_evidential = True
        self.use_normalizing_flows = PYRO_AVAILABLE
        
        # Training parameters
        self.batch_size = 64
        self.learning_rate = 0.001
        self.weight_decay = 1e-5
        self.num_epochs = 100
        self.lr_scheduler = "cosine"  # Options: "step", "plateau", "cosine"
        self.lr_step_size = 20
        self.lr_gamma = 0.5
        self.gradient_clip = 1.0
        self.early_stopping = True
        self.patience = 10
        self.validation_freq = 1
        self.save_freq = 5
        self.num_workers = 8
        self.seed = 42
        
        # Kalman Filter parameters
        self.kf_type = "IMM"  # Options: "KF", "EKF", "UKF", "IMM"
        self.num_motion_models = 3  # Number of motion models for IMM
        self.q_var = 0.01  # Process noise variance
        self.r_var = 0.1   # Measurement noise variance
        
        # Uncertainty thresholds
        self.uncertainty_threshold = 0.5  # Threshold for high uncertainty
        self.safety_factor = 2.0  # Safety factor for uncertainty inflation
        
        # Evaluation parameters
        self.eval_metrics = ["ade", "fde", "nll", "mr", "calibration", "collisions"]
        self.adversarial_test = True
        self.eval_k_predictions = [1, 5, 10]  # K values for minADE/minFDE
        
        # Safety parameters
        self.use_safety_envelope = True
        self.safety_envelope_size = 2.0  # Size of safety envelope (meters)
        self.collision_threshold = 1.5  # Collision threshold (meters)
        self.emergency_stopping_threshold = 0.8  # Emergency stopping threshold
        self.use_contingency_planning = True
        self.max_contingency_plans = 3
        
        # Visualization parameters
        self.visualize_predictions = True
        self.visualize_uncertainty = True
        self.visualize_attention = True
        self.visualize_scene_graph = True
        self.save_visualizations = True
        self.vis_path = "./visualizations"
        
        # Logging parameters
        self.log_dir = "./logs"
        self.checkpoint_dir = "./checkpoints"
        self.result_dir = "./results"
        self.use_tensorboard = TB_AVAILABLE
        self.log_freq = 10
        
        # Environment parameters
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.use_multi_gpu = HOROVOD_AVAILABLE and torch.cuda.device_count() > 1
        self.precision = "mixed"  # Options: "fp32", "fp16", "mixed"
        
        # Create directories
        os.makedirs(self.log_dir, exist_ok=True)
        os.makedirs(self.checkpoint_dir, exist_ok=True)
        os.makedirs(self.result_dir, exist_ok=True)
        os.makedirs(self.vis_path, exist_ok=True)
    
    def update(self, **kwargs):
        """Update config parameters based on a dictionary."""
        for key, value in kwargs.items():
            if hasattr(self, key):
                setattr(self, key, value)
            else:
                raise ValueError(f"Config has no attribute '{key}'")
        return self
    
    def to_dict(self):
        """Convert config to dictionary."""
        return {key: value for key, value in self.__dict__.items() 
                if not key.startswith('__') and not callable(value)}
    
    def __str__(self):
        """String representation of the config."""
        return "\n".join(f"{key}: {value}" for key, value in self.to_dict().items())


#######################
# DATA PROCESSING
#######################

class MapData:
    """
    Handle map data processing for trajectory prediction.
    
    Includes methods for loading, preprocessing, and extracting map features.
    """
    
    def __init__(self, config: Config):
        """
        Initialize the map data processor.
        
        Args:
            config: Configuration object
        """
        self.config = config
        self.maps = {}
        self.map_features = {}
        self.traffic_lights = {}
        self.traffic_signs = {}
        self.lane_connections = {}
        self.crosswalks = {}
        self.speed_limits = {}
        
        # Load map data if specified
        if self.config.use_map_data:
            self._load_map_data()
    
    def _load_map_data(self):
        """
        Load map data from the specified dataset.
        
        Supports multiple dataset formats (nuScenes, Argoverse, etc.).
        """
        print(f"Loading map data from {self.config.dataset_type}...")
        
        # Different loading procedures based on dataset type
        if self.config.dataset_type == "nuScenes":
            self._load_nuscenes_maps()
        elif self.config.dataset_type == "Argoverse":
            self._load_argoverse_maps()
        elif self.config.dataset_type == "Waymo":
            self._load_waymo_maps()
        else:
            # Example implementation for a custom dataset
            map_files = os.listdir(os.path.join(self.config.data_path, "maps"))
            for map_file in map_files:
                map_id = map_file.split(".")[0]
                self.maps[map_id] = self._load_map(os.path.join(self.config.data_path, "maps", map_file))
        
        print(f"Loaded {len(self.maps)} maps.")
    
    def _load_nuscenes_maps(self):
        """Load maps from nuScenes dataset."""
        # This is a placeholder - in a real implementation, use nuScenes API
        try:
            from nuscenes.map_expansion.map_api import NuScenesMap
            from nuscenes.nuscenes import NuScenes
            
            nusc = NuScenes(version='v1.0-mini', dataroot=self.config.data_path, verbose=True)
            
            # Load maps for each location
            for location in ['boston-seaport', 'singapore-hollandvillage', 
                            'singapore-onenorth', 'singapore-queenstown']:
                nusc_map = NuScenesMap(dataroot=self.config.data_path, map_name=location)
                self.maps[location] = nusc_map
                
                # Extract lanes, crosswalks, etc.
                lanes = nusc_map.lane_polygons
                ped_crossings = nusc_map.ped_crossing_polygons
                walkways = nusc_map.walkway_polygons
                stop_lines = nusc_map.stop_line_polygons
                
                # Store map features
                self.map_features[location] = {
                    'lanes': lanes,
                    'ped_crossings': ped_crossings,
                    'walkways': walkways,
                    'stop_lines': stop_lines
                }
                
                # Extract traffic lights and signs
                traffic_lights = []  # In a real implementation, extract from nuScenes
                traffic_signs = []   # In a real implementation, extract from nuScenes
                
                self.traffic_lights[location] = traffic_lights
                self.traffic_signs[location] = traffic_signs
                
                # Extract lane connections
                lane_connections = {}  # In a real implementation, extract lane connectivity
                self.lane_connections[location] = lane_connections
                
                # Extract speed limits
                speed_limits = {}  # In a real implementation, extract speed limits
                self.speed_limits[location] = speed_limits
                
        except ImportError:
            print("Warning: nuScenes API not available. Using placeholder map data.")
            # Create placeholder map data
            self.maps["placeholder"] = {"lanes": [], "crosswalks": []}
            self.map_features["placeholder"] = {"lanes": [], "crosswalks": []}
    
    def _load_argoverse_maps(self):
        """Load maps from Argoverse dataset."""
        # This is a placeholder - in a real implementation, use Argoverse API
        try:
            from argoverse.map_representation.map_api import ArgoverseMap
            
            argoverse_map = ArgoverseMap()
            self.maps["argoverse"] = argoverse_map
            
            # Extract map features
            for city in argoverse_map.city_name_to_city_id.keys():
                # Extract lanes, crosswalks, etc.
                lanes = []  # In a real implementation, extract lanes from Argoverse
                crosswalks = []  # In a real implementation, extract crosswalks
                
                self.map_features[city] = {
                    'lanes': lanes,
                    'crosswalks': crosswalks
                }
                
                # Extract traffic lights and signs
                self.traffic_lights[city] = []
                self.traffic_signs[city] = []
                
                # Extract lane connections
                self.lane_connections[city] = {}
                
                # Extract speed limits
                self.speed_limits[city] = {}
                
        except ImportError:
            print("Warning: Argoverse API not available. Using placeholder map data.")
            # Create placeholder map data
            self.maps["placeholder"] = {"lanes": [], "crosswalks": []}
            self.map_features["placeholder"] = {"lanes": [], "crosswalks": []}
    
    def _load_waymo_maps(self):
        """Load maps from Waymo dataset."""
        # This is a placeholder - in a real implementation, use Waymo API
        print("Warning: Waymo map loading not implemented. Using placeholder map data.")
        # Create placeholder map data
        self.maps["placeholder"] = {"lanes": [], "crosswalks": []}
        self.map_features["placeholder"] = {"lanes": [], "crosswalks": []}
    
    def _load_map(self, map_path: str):
        """Load a single map file."""
        # This is a placeholder - in a real implementation, load actual map data
        return {"lanes": [], "crosswalks": []}
    
    def extract_map_features(self, position: np.ndarray, 
                            map_id: str, 
                            radius: float = 50.0) -> np.ndarray:
        """
        Extract map features around a given position.
        
        Args:
            position: (x, y) position
            map_id: ID of the map to use
            radius: Radius around the position to extract features
            
        Returns:
            Map features as a numpy array
        """
        # Handle cases where map data isn't available
        if not self.config.use_map_data or map_id not in self.map_features:
            return np.zeros(self.config.map_features)
        
        # In a real implementation, extract actual map features
        # For this example, we'll return a random feature vector
        map_features = np.random.randn(self.config.map_features)
        
        # Normalize the features
        map_features = map_features / (np.linalg.norm(map_features) + 1e-8)
        
        return map_features
    
    def get_lane_direction(self, position: np.ndarray, map_id: str) -> np.ndarray:
        """
        Get the lane direction at a given position.
        
        Args:
            position: (x, y) position
            map_id: ID of the map to use
            
        Returns:
            Lane direction as a unit vector [dx, dy]
        """
        # In a real implementation, find the closest lane and get its direction
        # For this example, we'll return a random direction
        direction = np.random.randn(2)
        return direction / (np.linalg.norm(direction) + 1e-8)
    
    def get_nearest_lane_distance(self, position: np.ndarray, map_id: str) -> float:
        """
        Get the distance to the nearest lane.
        
        Args:
            position: (x, y) position
            map_id: ID of the map to use
            
        Returns:
            Distance to the nearest lane
        """
        # In a real implementation, find the closest lane and calculate distance
        # For this example, we'll return a random distance
        return np.random.uniform(0, 5)
    
    def get_speed_limit(self, position: np.ndarray, map_id: str) -> float:
        """
        Get the speed limit at a given position.
        
        Args:
            position: (x, y) position
            map_id: ID of the map to use
            
        Returns:
            Speed limit in m/s
        """
        # In a real implementation, find the speed limit for the current lane
        # For this example, we'll return a random speed limit
        return np.random.uniform(5, 25)  # 5-25 m/s (18-90 km/h)
    
    def get_traffic_light_state(self, position: np.ndarray, 
                              map_id: str, 
                              direction: np.ndarray) -> int:
        """
        Get the traffic light state in the direction of travel.
        
        Args:
            position: (x, y) position
            map_id: ID of the map to use
            direction: Direction of travel [dx, dy]
            
        Returns:
            Traffic light state (0=red, 1=yellow, 2=green, -1=none)
        """
        # In a real implementation, find the relevant traffic light
        # For this example, we'll return a random state
        return np.random.choice([-1, 0, 1, 2], p=[0.7, 0.1, 0.05, 0.15])
    
    def is_on_crosswalk(self, position: np.ndarray, map_id: str) -> bool:
        """
        Check if the position is on a crosswalk.
        
        Args:
            position: (x, y) position
            map_id: ID of the map to use
            
        Returns:
            True if on crosswalk, False otherwise
        """
        # In a real implementation, check if the position is inside any crosswalk polygon
        # For this example, we'll return a random boolean
        return np.random.random() < 0.05  # 5% chance of being on a crosswalk
    
    def render_map(self, ax, position: np.ndarray, map_id: str, radius: float = 50.0):
        """
        Render the map on a matplotlib axis.
        
        Args:
            ax: Matplotlib axis
            position: Center position [x, y]
            map_id: ID of the map to use
            radius: Radius around the position to render
        """
        if not self.config.use_map_data or map_id not in self.map_features:
            return
        
        # This is a placeholder - in a real implementation, render actual map features
        # For this example, we'll just add a grid
        ax.grid(True, linestyle='--', alpha=0.7)
        
        # Add some fake lanes
        for i in range(-3, 4):
            # Horizontal lanes
            ax.plot([position[0] - radius, position[0] + radius], 
                    [position[1] + i * 4, position[1] + i * 4], 
                    'g-', alpha=0.5)
            
            # Vertical lanes
            ax.plot([position[0] + i * 4, position[0] + i * 4], 
                    [position[1] - radius, position[1] + radius], 
                    'g-', alpha=0.5)
        
        # Add some fake crosswalks
        for i in range(-2, 3, 2):
            # Horizontal crosswalks
            ax.plot([position[0] - radius/4, position[0] + radius/4], 
                    [position[1] + i * 10, position[1] + i * 10], 
                    'y-', linewidth=3, alpha=0.7)
            
            # Vertical crosswalks
            ax.plot([position[0] + i * 10, position[0] + i * 10], 
                    [position[1] - radius/4, position[1] + radius/4], 
                    'y-', linewidth=3, alpha=0.7)


class WeatherData:
    """
    Handle weather data processing for trajectory prediction.
    
    Includes methods for loading, preprocessing, and extracting weather features.
    """
    
    def __init__(self, config: Config):
        """
        Initialize the weather data processor.
        
        Args:
            config: Configuration object
        """
        self.config = config
        self.weather_data = {}
        
        # Load weather data if specified
        if self.config.use_weather_data:
            self._load_weather_data()
    
    def _load_weather_data(self):
        """Load weather data from files or API."""
        print("Loading weather data...")
        
        # This is a placeholder - in a real implementation, load actual weather data
        # For this example, we'll create synthetic weather data
        weather_types = ["clear", "rain", "snow", "fog", "cloudy"]
        temperature_range = (-10, 40)  # Celsius
        precipitation_range = (0, 50)  # mm/h
        visibility_range = (50, 10000)  # meters
        wind_speed_range = (0, 30)  # m/s
        
        # Create synthetic weather data for a year
        dates = pd.date_range(start="2020-01-01", end="2020-12-31", freq="H")
        
        for date in dates:
            day_of_year = date.dayofyear / 365.0
            hour_of_day = date.hour / 24.0
            
            # Generate weather parameters with seasonal and daily variations
            season_factor = math.sin(day_of_year * 2 * math.pi)
            day_factor = math.sin(hour_of_day * 2 * math.pi)
            
            # Temperature follows seasonal pattern
            temperature = temperature_range[0] + (temperature_range[1] - temperature_range[0]) * (0.5 + 0.4 * season_factor + 0.1 * day_factor)
            
            # Precipitation more likely in certain seasons
            precip_probability = 0.2 + 0.2 * (1 + season_factor)
            precipitation = 0
            if random.random() < precip_probability:
                precipitation = random.uniform(0, precipitation_range[1])
            
            # Visibility affected by precipitation and time of day
            visibility_factor = 1.0
            if precipitation > 10:
                visibility_factor = 0.5
            elif precipitation > 0:
                visibility_factor = 0.8
            
            if hour_of_day < 0.25 or hour_of_day > 0.75:  # Night time
                visibility_factor *= 0.7
                
            visibility = visibility_range[0] + (visibility_range[1] - visibility_range[0]) * visibility_factor
            
            # Wind speed
            wind_speed = random.uniform(wind_speed_range[0], wind_speed_range[1])
            
            # Weather type
            if precipitation > 20 and temperature < 2:
                weather_type = "snow"
            elif precipitation > 5:
                weather_type = "rain"
            elif visibility < 1000:
                weather_type = "fog"
            elif random.random() < 0.3:
                weather_type = "cloudy"
            else:
                weather_type = "clear"
            
            # Store the weather data
            date_str = date.strftime("%Y-%m-%d %H:%M:%S")
            self.weather_data[date_str] = {
                "temperature": temperature,
                "precipitation": precipitation,
                "visibility": visibility,
                "wind_speed": wind_speed,
                "weather_type": weather_type
            }
        
        print(f"Generated weather data for {len(self.weather_data)} time points.")
    
    def get_weather_features(self, timestamp: str) -> np.ndarray:
        """
        Get weather features for a given timestamp.
        
        Args:
            timestamp: Timestamp in the format "YYYY-MM-DD HH:MM:SS"
            
        Returns:
            Weather features as a numpy array
        """
        if not self.config.use_weather_data or not self.weather_data:
            return np.zeros(self.config.weather_features)
        
        # Find the closest timestamp in our weather data
        # In a real implementation, use more sophisticated time matching
        weather_timestamps = list(self.weather_data.keys())
        closest_ts = min(weather_timestamps, key=lambda x: abs(pd.Timestamp(x) - pd.Timestamp(timestamp)))
        
        weather = self.weather_data[closest_ts]
        
        # Create a feature vector
        features = np.zeros(self.config.weather_features)
        
        # Normalize temperature to [-1, 1]
        features[0] = (weather["temperature"] - 15) / 25
        
        # Normalize precipitation to [0, 1]
        features[1] = min(1.0, weather["precipitation"] / 50)
        
        # Normalize visibility to [0, 1]
        features[2] = weather["visibility"] / 10000
        
        # Normalize wind speed to [0, 1]
        features[3] = weather["wind_speed"] / 30
        
        # One-hot encode weather type
        weather_types = ["clear", "rain", "snow", "fog", "cloudy"]
        weather_idx = weather_types.index(weather["weather_type"]) if weather["weather_type"] in weather_types else 0
        features[4 + weather_idx] = 1.0
        
        return features


class TrafficRules:
    """
    Handle traffic rules for trajectory prediction.
    
    Includes methods for loading and checking traffic rules.
    """
    
    def __init__(self, config: Config, map_data: MapData):
        """
        Initialize the traffic rules processor.
        
        Args:
            config: Configuration object
            map_data: Map data object
        """
        self.config = config
        self.map_data = map_data
        self.traffic_rules = {}
        
        # Load traffic rules if specified
        if self.config.use_traffic_rules:
            self._load_traffic_rules()
    
    def _load_traffic_rules(self):
        """Load traffic rules from files or database."""
        print("Loading traffic rules...")
        
        # This is a placeholder - in a real implementation, load actual traffic rules
        # For this example, we'll create synthetic traffic rules
        rule_types = ["speed_limit", "stop_sign", "yield", "no_turn", "one_way"]
        
        for rule_type in rule_types:
            self.traffic_rules[rule_type] = {}
    
    def check_traffic_rules(self, 
                           position: np.ndarray, 
                           direction: np.ndarray, 
                           map_id: str) -> Dict[str, Any]:
        """
        Check traffic rules at a given position and direction.
        
        Args:
            position: (x, y) position
            direction: Direction of travel [dx, dy]
            map_id: ID of the map to use
            
        Returns:
            Dictionary of applicable traffic rules
        """
        if not self.config.use_traffic_rules:
            return {}
        
        # In a real implementation, find applicable traffic rules
        # For this example, we'll return synthetic rules
        
        # Get speed limit from map data
        speed_limit = self.map_data.get_speed_limit(position, map_id)
        
        # Get traffic light state from map data
        traffic_light = self.map_data.get_traffic_light_state(position, map_id, direction)
        
        # Random rules
        stop_sign = random.random() < 0.05
        yield_sign = random.random() < 0.05 and not stop_sign
        no_turn = random.random() < 0.1
        one_way = random.random() < 0.3
        
        # Return combined rules
        return {
            "speed_limit": speed_limit,
            "traffic_light": traffic_light,
            "stop_sign": stop_sign,
            "yield_sign": yield_sign,
            "no_turn": no_turn,
            "one_way": one_way
        }


class AgentState:
    """
    Represents the state of an agent (vehicle, pedestrian, etc.) in the scene.
    """
    
    def __init__(self, 
                agent_id: str, 
                agent_type: str, 
                position: np.ndarray, 
                velocity: np.ndarray, 
                heading: float, 
                length: float = 4.0, 
                width: float = 2.0, 
                timestamp: float = 0.0):
        """
        Initialize an agent state.
        
        Args:
            agent_id: Unique identifier for the agent
            agent_type: Type of agent (vehicle, pedestrian, cyclist, etc.)
            position: (x, y) position
            velocity: (vx, vy) velocity
            heading: Heading angle in radians
            length: Length of the agent in meters
            width: Width of the agent in meters
            timestamp: Timestamp of the state
        """
        self.agent_id = agent_id
        self.agent_type = agent_type
        self.position = np.array(position)
        self.velocity = np.array(velocity)
        self.heading = heading
        self.length = length
        self.width = width
        self.timestamp = timestamp
        
        # Derived properties
        self.speed = np.linalg.norm(velocity)
        self.acceleration = np.zeros(2)  # Will be calculated from velocity changes
        self.yaw_rate = 0.0  # Will be calculated from heading changes
        
        # Future trajectory (ground truth)
        self.future_trajectory = None
        
        # Additional properties
        self.lane_id = None
        self.on_intersection = False
        self.distance_to_lane = 0.0
        self.distance_to_intersection = float('inf')
        self.nearest_agents = {}  # {agent_id: distance}
        self.traffic_rules = {}
        self.intention = None
        self.intention_prob = None
        
    def update_derived_properties(self, prev_state: Optional['AgentState'] = None):
        """
        Update derived properties based on current state and previous state.
        
        Args:
            prev_state: Previous state of the agent
        """
        self.speed = np.linalg.norm(self.velocity)
        
        if prev_state is not None:
            dt = self.timestamp - prev_state.timestamp
            if dt > 0:
                # Calculate acceleration
                self.acceleration = (self.velocity - prev_state.velocity) / dt
                
                # Calculate yaw rate
                heading_diff = self.heading - prev_state.heading
                # Normalize to [-pi, pi]
                heading_diff = (heading_diff + np.pi) % (2 * np.pi) - np.pi
                self.yaw_rate = heading_diff / dt
    
    def to_vector(self) -> np.ndarray:
        """
        Convert agent state to feature vector.
        
        Returns:
            Feature vector representing the agent state
        """
        # Basic state features: [x, y, vx, vy, heading, speed, acc_x, acc_y, yaw_rate]
        features = np.zeros(9)
        features[0:2] = self.position
        features[2:4] = self.velocity
        features[4] = self.heading
        features[5] = self.speed
        features[6:8] = self.acceleration
        features[8] = self.yaw_rate
        
        return features
        
    def predict_state(self, dt: float, motion_model: str = "CV") -> 'AgentState':
        """
        Predict future state after time dt using a motion model.
        
        Args:
            dt: Time increment in seconds
            motion_model: Motion model to use (CV: constant velocity, 
                                             CA: constant acceleration,
                                             CTRV: constant turn rate and velocity)
            
        Returns:
            Predicted agent state
        """
        # Create a copy of current state
        predicted_state = copy.deepcopy(self)
        predicted_state.timestamp += dt
        
        if motion_model == "CV":  # Constant velocity
            # Update position
            predicted_state.position = self.position + self.velocity * dt
            
            # Velocity remains constant
            predicted_state.velocity = self.velocity
            
            # Heading remains constant
            predicted_state.heading = self.heading
            
        elif motion_model == "CA":  # Constant acceleration
            # Update position
            predicted_state.position = self.position + self.velocity * dt + 0.5 * self.acceleration * dt**2
            
            # Update velocity
            predicted_state.velocity = self.velocity + self.acceleration * dt
            
            # Heading update based on velocity direction if speed is sufficient
            if np.linalg.norm(predicted_state.velocity) > 0.5:
                predicted_state.heading = np.arctan2(predicted_state.velocity[1], predicted_state.velocity[0])
            
        elif motion_model == "CTRV":  # Constant turn rate and velocity
            # Handle the case where yaw_rate is close to zero
            if abs(self.yaw_rate) < 1e-6:
                # Same as CV model if yaw_rate is zero
                predicted_state.position = self.position + self.velocity * dt
                predicted_state.heading = self.heading
            else:
                # Calculate heading
                predicted_state.heading = self.heading + self.yaw_rate * dt
                
                # Calculate position change
                v_over_omega = self.speed / self.yaw_rate
                predicted_state.position[0] = self.position[0] + v_over_omega * (
                    np.sin(predicted_state.heading) - np.sin(self.heading))
                predicted_state.position[1] = self.position[1] + v_over_omega * (
                    -np.cos(predicted_state.heading) + np.cos(self.heading))
            
            # Speed remains constant, but direction changes
            speed = self.speed
            predicted_state.velocity = np.array([
                speed * np.cos(predicted_state.heading),
                speed * np.sin(predicted_state.heading)
            ])
            
        else:
            raise ValueError(f"Unknown motion model: {motion_model}")
        
        return predicted_state
    
    def to_dict(self) -> Dict:
        """Convert agent state to dictionary."""
        return {
            'agent_id': self.agent_id,
            'agent_type': self.agent_type,
            'position': self.position.tolist(),
            'velocity': self.velocity.tolist(),
            'heading': self.heading,
            'length': self.length,
            'width': self.width,
            'timestamp': self.timestamp,
            'speed': self.speed,
            'acceleration': self.acceleration.tolist(),
            'yaw_rate': self.yaw_rate,
            'lane_id': self.lane_id,
            'on_intersection': self.on_intersection,
            'distance_to_lane': self.distance_to_lane,
            'distance_to_intersection': self.distance_to_intersection,
            'intention': self.intention,
            'intention_prob': self.intention_prob
        }
    
    @classmethod
    def from_dict(cls, data: Dict) -> 'AgentState':
        """Create agent state from dictionary."""
        state = cls(
            agent_id=data['agent_id'],
            agent_type=data['agent_type'],
            position=np.array(data['position']),
            velocity=np.array(data['velocity']),
            heading=data['heading'],
            length=data['length'],
            width=data['width'],
            timestamp=data['timestamp']
        )
        
        # Set additional properties if available
        if 'acceleration' in data:
            state.acceleration = np.array(data['acceleration'])
        if 'yaw_rate' in data:
            state.yaw_rate = data['yaw_rate']
        if 'lane_id' in data:
            state.lane_id = data['lane_id']
        if 'on_intersection' in data:
            state.on_intersection = data['on_intersection']
        if 'distance_to_lane' in data:
            state.distance_to_lane = data['distance_to_lane']
        if 'distance_to_intersection' in data:
            state.distance_to_intersection = data['distance_to_intersection']
        if 'intention' in data:
            state.intention = data['intention']
        if 'intention_prob' in data:
            state.intention_prob = data['intention_prob']
        
        return state

    def get_bounding_box(self) -> np.ndarray:
        """
        Get the bounding box of the agent.
        
        Returns:
            Array of (x, y) coordinates for the 4 corners of the bounding box
        """
        # Rotation matrix
        c, s = np.cos(self.heading), np.sin(self.heading)
        R = np.array([[c, -s], [s, c]])
        
        # Half-dimensions
        hl = self.length / 2
        hw = self.width / 2
        
        # Corners in vehicle frame
        corners_local = np.array([
            [-hl, -hw],  # rear-left
            [hl, -hw],   # front-left
            [hl, hw],    # front-right
            [-hl, hw]    # rear-right
        ])
        
        # Corners in world frame
        corners_world = np.array([self.position + R @ corner for corner in corners_local])
        
        return corners_world
    
    def check_collision(self, other: 'AgentState') -> bool:
        """
        Check if this agent collides with another agent.
        
        Args:
            other: The other agent state
            
        Returns:
            True if collision, False otherwise
        """
        # Simple distance-based collision check
        dist = np.linalg.norm(self.position - other.position)
        min_dist = (self.length + other.length) / 2
        
        return dist < min_dist


class SceneState:
    """
    Represents the state of the entire scene, including all agents and environmental context.
    """
    
    def __init__(self, 
                timestamp: float,
                map_id: str,
                ego_agent_id: str,
                weather_info: Optional[Dict] = None,
                time_info: Optional[Dict] = None):
        """
        Initialize a scene state.
        
        Args:
            timestamp: Timestamp of the scene
            map_id: ID of the map
            ego_agent_id: ID of the ego agent
            weather_info: Weather information
            time_info: Time information (time of day, etc.)
        """
        self.timestamp = timestamp
        self.map_id = map_id
        self.ego_agent_id = ego_agent_id
        self.weather_info = weather_info if weather_info is not None else {}
        self.time_info = time_info if time_info is not None else {}
        
        # Agents in the scene
        self.agents: Dict[str, AgentState] = {}
        
        # Scene graph representation
        self.scene_graph = nx.DiGraph()
        
        # Traffic rules applicable to the scene
        self.traffic_rules = {}
    
    def add_agent(self, agent: AgentState):
        """Add an agent to the scene."""
        self.agents[agent.agent_id] = agent
        self.update_scene_graph()
    
    def remove_agent(self, agent_id: str):
        """Remove an agent from the scene."""
        if agent_id in self.agents:
            del self.agents[agent_id]
            self.update_scene_graph()
    
    def get_ego_agent(self) -> Optional[AgentState]:
        """Get the ego agent."""
        return self.agents.get(self.ego_agent_id)
    
    def update_scene_graph(self):
        """Update the scene graph representation."""
        # Clear the graph
        self.scene_graph.clear()
        
        # Add nodes for all agents
        for agent_id, agent in self.agents.items():
            self.scene_graph.add_node(agent_id, 
                                     state=agent, 
                                     position=agent.position,
                                     agent_type=agent.agent_type)
        
        # Add edges based on proximity and interaction potential
        for agent1_id, agent1 in self.agents.items():
            for agent2_id, agent2 in self.agents.items():
                if agent1_id != agent2_id:
                    # Calculate distance between agents
                    distance = np.linalg.norm(agent1.position - agent2.position)
                    rel_velocity = np.linalg.norm(agent1.velocity - agent2.velocity)
                    
                    # Only connect agents that are close enough and potentially interacting
                    if distance < 50.0:  # 50 meters threshold
                        # Calculate time-to-collision (TTC) if relative motion exists
                        ttc = float('inf')
                        if rel_velocity > 0.1:
                            rel_pos = agent2.position - agent1.position
                            ttc_vec = rel_pos / (agent1.velocity - agent2.velocity + 1e-6)
                            if np.all(ttc_vec > 0):  # Only consider positive TTC
                                ttc = np.min(ttc_vec)
                        
                        # Add edge with attributes
                        self.scene_graph.add_edge(
                            agent1_id, agent2_id,
                            distance=distance,
                            rel_velocity=rel_velocity,
                            ttc=ttc
                        )
    
    def get_neighbor_agents(self, agent_id: str, distance_threshold: float = 50.0) -> Dict[str, AgentState]:
        """
        Get neighboring agents within a distance threshold.
        
        Args:
            agent_id: ID of the agent
            distance_threshold: Maximum distance to consider agents as neighbors
            
        Returns:
            Dictionary of neighboring agents {agent_id: agent}
        """
        if agent_id not in self.agents:
            return {}
        
        agent = self.agents[agent_id]
        neighbors = {}
        
        for other_id, other_agent in self.agents.items():
            if other_id != agent_id:
                distance = np.linalg.norm(agent.position - other_agent.position)
                if distance < distance_threshold:
                    neighbors[other_id] = other_agent
        
        return neighbors
    
    def to_dict(self) -> Dict:
        """Convert scene state to dictionary."""
        return {
            'timestamp': self.timestamp,
            'map_id': self.map_id,
            'ego_agent_id': self.ego_agent_id,
            'weather_info': self.weather_info,
            'time_info': self.time_info,
            'agents': {agent_id: agent.to_dict() for agent_id, agent in self.agents.items()},
            'traffic_rules': self.traffic_rules
        }
    
    @classmethod
    def from_dict(cls, data: Dict) -> 'SceneState':
        """Create scene state from dictionary."""
        scene = cls(
            timestamp=data['timestamp'],
            map_id=data['map_id'],
            ego_agent_id=data['ego_agent_id'],
            weather_info=data.get('weather_info'),
            time_info=data.get('time_info')
        )
        
        # Add agents
        for agent_id, agent_data in data['agents'].items():
            scene.agents[agent_id] = AgentState.from_dict(agent_data)
        
        # Set traffic rules if available
        if 'traffic_rules' in data:
            scene.traffic_rules = data['traffic_rules']
        
        # Update scene graph
        scene.update_scene_graph()
        
        return scene
    
    def check_collisions(self) -> List[Tuple[str, str]]:
        """
        Check for collisions between all agents in the scene.
        
        Returns:
            List of (agent1_id, agent2_id) pairs that are colliding
        """
        collisions = []
        
        # Check all pairs of agents
        agent_ids = list(self.agents.keys())
        for i in range(len(agent_ids)):
            for j in range(i+1, len(agent_ids)):
                agent1 = self.agents[agent_ids[i]]
                agent2 = self.agents[agent_ids[j]]
                
                if agent1.check_collision(agent2):
                    collisions.append((agent_ids[i], agent_ids[j]))
        
        return collisions


class TrajectoryDataset(Dataset):
    """Dataset for trajectory prediction."""
    
    def __init__(self, 
                config: Config, 
                map_data: MapData, 
                weather_data: WeatherData,
                traffic_rules: TrafficRules,
                split: str = 'train',
                transform: Optional[Callable] = None):
        """
        Initialize the trajectory dataset.
        
        Args:
            config: Configuration object
            map_data: Map data object
            weather_data: Weather data object
            traffic_rules: Traffic rules object
            split: Dataset split ('train', 'val', 'test')
            transform: Data transformation function
        """
        self.config = config
        self.map_data = map_data
        self.weather_data = weather_data
        self.traffic_rules = traffic_rules
        self.split = split
        self.transform = transform
        
        # Data storage
        self.scenes: List[SceneState] = []
        self.trajectories: Dict[str, Dict[str, List[AgentState]]] = {}  # {scene_id: {agent_id: [states]}}
        self.scene_ids: List[str] = []
        self.agent_ids: Dict[str, List[str]] = {}  # {scene_id: [agent_ids]}
        
        # Load the data
        self._load_data()
    
    def _load_data(self):
        """Load trajectory data from files."""
        print(f"Loading {self.split} data...")
        
        if self.config.dataset_type == "nuScenes":
            self._load_nuscenes_data()
        elif self.config.dataset_type == "Argoverse":
            self._load_argoverse_data()
        elif self.config.dataset_type == "Waymo":
            self._load_waymo_data()
        else:
            # Example implementation for a custom dataset
            self._load_custom_data()
        
        print(f"Loaded {len(self.scene_ids)} scenes with {sum(len(agents) for agents in self.agent_ids.values())} agents.")
    
    def _load_nuscenes_data(self):
        """Load data from nuScenes dataset."""
        # This is a placeholder - in a real implementation, use nuScenes API
        try:
            from nuscenes.nuscenes import NuScenes
            
            print("Loading nuScenes dataset...")
            nusc = NuScenes(version='v1.0-mini', dataroot=self.config.data_path, verbose=True)
            
            # Process scenes
            scene_idx = 0
            for scene in nusc.scene:
                # Filter scenes based on split
                if self.split == 'train' and scene_idx % 10 in range(0, 7):
                    pass
                elif self.split == 'val' and scene_idx % 10 in range(7, 9):
                    pass
                elif self.split == 'test' and scene_idx % 10 == 9:
                    pass
                else:
                    scene_idx += 1
                    continue
                
                scene_id = f"scene_{scene_idx}"
                self.scene_ids.append(scene_id)
                
                # Get the first sample in the scene
                sample_token = scene['first_sample_token']
                
                # Process samples in the scene
                while sample_token:
                    sample = nusc.get('sample', sample_token)
                    timestamp = sample['timestamp'] / 1e6  # Convert to seconds
                    
                    # Create scene state
                    scene_state = SceneState(
                        timestamp=timestamp,
                        map_id=nusc.get('log', scene['log_token'])['location'],
                        ego_agent_id='ego'
                    )
                    
                    # Add ego vehicle
                    ego_state = AgentState(
                        agent_id='ego',
                        agent_type='vehicle',
                        position=np.array([0.0, 0.0]),  # Placeholder
                        velocity=np.array([0.0, 0.0]),  # Placeholder
                        heading=0.0,
                        timestamp=timestamp
                    )
                    scene_state.add_agent(ego_state)
                    
                    # Process annotations
                    for ann_token in sample['anns']:
                        ann = nusc.get('sample_annotation', ann_token)
                        
                        # Skip annotations without tracking
                        if not ann['instance_token']:
                            continue
                        
                        # Get agent type
                        agent_type = ann['category_name'].split('.')[0]
                        if agent_type not in ['vehicle', 'pedestrian', 'bicycle']:
                            continue
                        
                        # Get agent state
                        translation = np.array(ann['translation'][:2])
                        rotation = ann['rotation']
                        # Convert quaternion to heading (simplified for 2D)
                        heading = np.arctan2(2 * (rotation[0] * rotation[1] + rotation[2] * rotation[3]),
                                           1 - 2 * (rotation[1]**2 + rotation[2]**2))
                        
                        # Get size
                        size = ann['size']
                        length, width = size[0], size[1]
                        
                        # Create agent state
                        agent_state = AgentState(
                            agent_id=ann['instance_token'],
                            agent_type=agent_type,
                            position=translation,
                            velocity=np.array([0.0, 0.0]),  # Will be calculated later
                            heading=heading,
                            length=length,
                            width=width,
                            timestamp=timestamp
                        )
                        
                        # Add agent to scene
                        scene_state.add_agent(agent_state)
                        
                        # Store agent ID
                        if scene_id not in self.agent_ids:
                            self.agent_ids[scene_id] = []
                        if ann['instance_token'] not in self.agent_ids[scene_id]:
                            self.agent_ids[scene_id].append(ann['instance_token'])
                    
                    # Add scene to list
                    self.scenes.append(scene_state)
                    
                    # Get next sample
                    sample_token = sample['next']
                
                scene_idx += 1
                
            # Process trajectories
            for scene_id, scene_agent_ids in self.agent_ids.items():
                self.trajectories[scene_id] = {}
                for agent_id in scene_agent_ids:
                    self.trajectories[scene_id][agent_id] = []
                    
                    # Find all states for this agent
                    for scene in self.scenes:
                        if agent_id in scene.agents:
                            self.trajectories[scene_id][agent_id].append(scene.agents[agent_id])
                    
                    # Sort by timestamp
                    self.trajectories[scene_id][agent_id].sort(key=lambda x: x.timestamp)
                    
                    # Calculate velocities and accelerations
                    for i in range(1, len(self.trajectories[scene_id][agent_id])):
                        curr_state = self.trajectories[scene_id][agent_id][i]
                        prev_state = self.trajectories[scene_id][agent_id][i-1]
                        dt = curr_state.timestamp - prev_state.timestamp
                        
                        if dt > 0:
                            # Calculate velocity
                            curr_state.velocity = (curr_state.position - prev_state.position) / dt
                            
                            # Update derived properties
                            curr_state.update_derived_properties(prev_state)
            
        except ImportError:
            print("Warning: nuScenes API not available. Using synthetic data.")
            # Generate synthetic data
            self._generate_synthetic_data()
    
    def _load_argoverse_data(self):
        """Load data from Argoverse dataset."""
        # This is a placeholder - in a real implementation, use Argoverse API
        print("Warning: Argoverse loader not implemented. Using synthetic data.")
        self._generate_synthetic_data()
    
    def _load_waymo_data(self):
        """Load data from Waymo dataset."""
        # This is a placeholder - in a real implementation, use Waymo API
        print("Warning: Waymo loader not implemented. Using synthetic data.")
        self._generate_synthetic_data()
    
    def _load_custom_data(self):
        """Load data from custom dataset."""
        # This is a placeholder - in a real implementation, load actual data
        print("No custom dataset available. Using synthetic data.")
        self._generate_synthetic_data()
    
    def _generate_synthetic_data(self):
        """Generate synthetic trajectory data for testing."""
        print("Generating synthetic data...")
        
        # Parameters
        num_scenes = 100
        agents_per_scene = (5, 20)  # (min, max)
        scene_duration = 10.0  # seconds
        sampling_rate = 10  # Hz
        dt = 1.0 / sampling_rate
        
        # Motion patterns
        motion_patterns = [
            "straight", "left_turn", "right_turn", "accelerate", "decelerate", 
            "lane_change_left", "lane_change_right", "stationary"
        ]
        
        # Agent types and sizes
        agent_types = {
            "vehicle": (4.5, 2.0),  # (length, width)
            "pedestrian": (0.8, 0.8),
            "cyclist": (1.8, 0.8)
        }
        
        # Weather types
        weather_types = ["clear", "rain", "snow", "fog", "cloudy"]
        
        # Time of day
        time_periods = ["morning", "afternoon", "evening", "night"]
        
        # Map IDs
        map_ids = ["map_01", "map_02", "map_03", "map_04", "map_05"]
        
        # Generate scenes
        for scene_idx in range(num_scenes):
            scene_id = f"scene_{scene_idx}"
            self.scene_ids.append(scene_id)
            self.agent_ids[scene_id] = []
            self.trajectories[scene_id] = {}
            
            # Scene parameters
            num_agents = random.randint(*agents_per_scene)
            scene_map_id = random.choice(map_ids)
            start_time = random.uniform(0, 72000)  # Seconds (20 hours)
            
            # Generate weather and time info
            weather_type = random.choice(weather_types)
            time_period = random.choice(time_periods)
            weather_info = {
                "weather_type": weather_type,
                "temperature": random.uniform(-10, 40),
                "precipitation": random.uniform(0, 50) if weather_type in ["rain", "snow"] else 0,
                "visibility": random.uniform(50, 1000) if weather_type == "fog" else random.uniform(1000, 10000),
                "wind_speed": random.uniform(0, 30)
            }
            time_info = {
                "time_period": time_period,
                "is_day": time_period in ["morning", "afternoon"],
                "time_of_day": random.uniform(0, 24)  # Hour of day
            }
            
            # Generate agents
            for agent_idx in range(num_agents):
                agent_id = f"agent_{scene_id}_{agent_idx}"
                agent_type = random.choices(list(agent_types.keys()), weights=[0.7, 0.2, 0.1])[0]
                length, width = agent_types[agent_type]
                
                # Initial position (random within a reasonable area)
                position = np.array([
                    random.uniform(-100, 100),
                    random.uniform(-100, 100)
                ])
                
                # Initial heading
                heading = random.uniform(-np.pi, np.pi)
                
                # Initial velocity
                if agent_type == "pedestrian":
                    speed = random.uniform(0, 2)  # m/s
                elif agent_type == "cyclist":
                    speed = random.uniform(0, 6)  # m/s
                else:  # vehicle
                    speed = random.uniform(0, 15)  # m/s
                
                velocity = np.array([
                    speed * np.cos(heading),
                    speed * np.sin(heading)
                ])
                
                # Select motion pattern
                motion_pattern = random.choice(motion_patterns)
                
                # Generate trajectory
                agent_states = []
                for t in range(int(scene_duration * sampling_rate) + 1):
                    timestamp = start_time + t * dt
                    
                    # Create agent state
                    state = AgentState(
                        agent_id=agent_id,
                        agent_type=agent_type,
                        position=position.copy(),
                        velocity=velocity.copy(),
                        heading=heading,
                        length=length,
                        width=width,
                        timestamp=timestamp
                    )
                    
                    agent_states.append(state)
                    
                    # Update position for next step
                    if motion_pattern == "straight":
                        position += velocity * dt
                    elif motion_pattern == "left_turn":
                        # Gradual left turn
                        turn_rate = 0.1  # rad/s
                        heading += turn_rate * dt
                        velocity = np.array([
                            speed * np.cos(heading),
                            speed * np.sin(heading)
                        ])
                        position += velocity * dt
                    elif motion_pattern == "right_turn":
                        # Gradual right turn
                        turn_rate = -0.1  # rad/s
                        heading += turn_rate * dt
                        velocity = np.array([
                            speed * np.cos(heading),
                            speed * np.sin(heading)
                        ])
                        position += velocity * dt
                    elif motion_pattern == "accelerate":
                        # Gradual acceleration
                        accel = 0.5  # m/s^2
                        speed += accel * dt
                        velocity = np.array([
                            speed * np.cos(heading),
                            speed * np.sin(heading)
                        ])
                        position += velocity * dt
                    elif motion_pattern == "decelerate":
                        # Gradual deceleration
                        accel = -0.5  # m/s^2
                        speed = max(0, speed + accel * dt)
                        velocity = np.array([
                            speed * np.cos(heading),
                            speed * np.sin(heading)
                        ])
                        position += velocity * dt
                    elif motion_pattern == "lane_change_left":
                        # Lane change (simplified)
                        lateral_speed = 0.5  # m/s
                        lateral_dir = np.array([-np.sin(heading), np.cos(heading)])
                        position += velocity * dt + lateral_dir * lateral_speed * dt
                    elif motion_pattern == "lane_change_right":
                        # Lane change (simplified)
                        lateral_speed = -0.5  # m/s
                        lateral_dir = np.array([-np.sin(heading), np.cos(heading)])
                        position += velocity * dt + lateral_dir * lateral_speed * dt
                    elif motion_pattern == "stationary":
                        # No movement
                        pass
                    
                    # Add some noise
                    if self.config.random_noise:
                        position += np.random.normal(0, 0.05, 2)  # Small position noise
                        heading += np.random.normal(0, 0.01)  # Small heading noise
                
                # Add agent to the dataset
                self.agent_ids[scene_id].append(agent_id)
                self.trajectories[scene_id][agent_id] = agent_states
                
                # Update derived properties
                for i in range(1, len(agent_states)):
                    agent_states[i].update_derived_properties(agent_states[i-1])
            
            # Create scene states (one per timestep)
            for t in range(int(scene_duration * sampling_rate) + 1):
                timestamp = start_time + t * dt
                
                # Create scene state
                scene_state = SceneState(
                    timestamp=timestamp,
                    map_id=scene_map_id,
                    ego_agent_id=self.agent_ids[scene_id][0],  # First agent as ego
                    weather_info=weather_info,
                    time_info=time_info
                )
                
                # Add agents to scene
                for agent_id in self.agent_ids[scene_id]:
                    agent_state = self.trajectories[scene_id][agent_id][t]
                    scene_state.add_agent(agent_state)
                
                # Add scene to list
                self.scenes.append(scene_state)
                
        print(f"Generated {len(self.scene_ids)} synthetic scenes with {len(self.scenes)} scene states.")
    
    def __len__(self) -> int:
        """Get the number of trajectory samples in the dataset."""
        return sum(len(agent_ids) for agent_ids in self.agent_ids.values())
    
    def __getitem__(self, idx: int) -> Dict:
        """
        Get a trajectory sample by index.
        
        Args:
            idx: Sample index
            
        Returns:
            Dictionary containing the trajectory sample data
        """
        # Map flat index to (scene_id, agent_id)
        scene_idx = 0
        agent_idx = idx
        for scene_id in self.scene_ids:
            num_agents = len(self.agent_ids[scene_id])
            if agent_idx < num_agents:
                agent_id = self.agent_ids[scene_id][agent_idx]
                break
            agent_idx -= num_agents
            scene_idx += 1
        else:
            raise IndexError(f"Index {idx} out of range")
        
        scene_id = self.scene_ids[scene_idx]
        agent_id = self.agent_ids[scene_id][agent_idx]
        
        # Get agent trajectory
        trajectory = self.trajectories[scene_id][agent_id]
        
        # Get relevant scenes for this trajectory
        relevant_scenes = []
        for scene in self.scenes:
            if scene.timestamp >= trajectory[0].timestamp and scene.timestamp <= trajectory[-1].timestamp:
                relevant_scenes.append(scene)
        
        # Sort scenes by timestamp
        relevant_scenes.sort(key=lambda x: x.timestamp)
        
        # Extract history and future trajectories
        history_len = self.config.history_len
        future_len = self.config.future_len
        
        # Find current time index (last history point)
        current_idx = min(history_len, len(trajectory))
        
        # Get history and future trajectories
        history = trajectory[:current_idx]
        future = trajectory[current_idx:current_idx + future_len]
        
        # Pad history if needed
        if len(history) < history_len:
            # Duplicate first state to pad history
            first_state = history[0]
            padding = [copy.deepcopy(first_state) for _ in range(history_len - len(history))]
            history = padding + history
        
        # Get current scene (at current_idx)
        current_time = trajectory[current_idx - 1].timestamp if current_idx > 0 else trajectory[0].timestamp
        current_scene = None
        for scene in relevant_scenes:
            if abs(scene.timestamp - current_time) < 1e-3:  # Small epsilon for float comparison
                current_scene = scene
                break
        
        if current_scene is None:
            # Use the closest scene by timestamp
            relevant_scenes.sort(key=lambda x: abs(x.timestamp - current_time))
            current_scene = relevant_scenes[0]
        
        # Get surrounding agents
        surrounding_agents = []
        for other_id, other_agent in current_scene.agents.items():
            if other_id != agent_id:
                surrounding_agents.append(other_agent)
                
        # Sort surrounding agents by distance to target agent
        target_pos = trajectory[current_idx - 1].position if current_idx > 0 else trajectory[0].position
        surrounding_agents.sort(key=lambda x: np.linalg.norm(x.position - target_pos))
        
        # Limit to max number of surrounding agents
        surrounding_agents = surrounding_agents[:self.config.max_surrounding_agents]
        
        # Get map features around the agent
        map_features = self._get_map_features(
            scene_id=scene_id,
            position=target_pos,
            heading=trajectory[current_idx - 1].heading if current_idx > 0 else trajectory[0].heading,
            radius=self.config.map_radius
        )
        
        # Get traffic rules
        traffic_features = self._get_traffic_features(scene_id, agent_id, current_scene)
        
        # Get weather features
        weather_features = self._get_weather_features(current_scene)
        
        # Create sample
        sample = {
            'scene_id': scene_id,
            'agent_id': agent_id,
            'agent_type': trajectory[0].agent_type,
            'history': self._process_trajectory_for_model(history),
            'future': self._process_trajectory_for_model(future),
            'surrounding_agents': [self._process_agent_for_model(agent) for agent in surrounding_agents],
            'map_features': map_features,
            'traffic_features': traffic_features,
            'weather_features': weather_features,
            'current_state': self._process_agent_for_model(trajectory[current_idx - 1] if current_idx > 0 else trajectory[0]),
            'current_timestamp': current_time,
            'future_timestamps': [state.timestamp for state in future]
        }
        
        # Apply transforms if available
        if self.transform:
            sample = self.transform(sample)
            
        return sample
    
    def _process_trajectory_for_model(self, trajectory: List[AgentState]) -> Dict:
        """
        Process trajectory data for model input.
        
        Args:
            trajectory: List of agent states
            
        Returns:
            Dictionary with processed trajectory features
        """
        if not trajectory:
            return {
                'positions': np.zeros((0, 2)),
                'velocities': np.zeros((0, 2)),
                'accelerations': np.zeros((0, 2)),
                'headings': np.zeros(0),
                'angular_velocities': np.zeros(0),
                'timestamps': np.zeros(0)
            }
        
        # Extract features
        positions = np.array([state.position for state in trajectory])
        velocities = np.array([state.velocity for state in trajectory])
        headings = np.array([state.heading for state in trajectory])
        timestamps = np.array([state.timestamp for state in trajectory])
        
        # Calculate accelerations
        if len(velocities) >= 2:
            time_diffs = timestamps[1:] - timestamps[:-1]
            time_diffs = np.where(time_diffs == 0, 1e-6, time_diffs)  # Avoid division by zero
            accelerations = np.diff(velocities, axis=0) / time_diffs[:, np.newaxis]
            # Pad with zeros for the first timestep
            accelerations = np.vstack([np.zeros_like(velocities[0:1]), accelerations])
        else:
            accelerations = np.zeros_like(velocities)
        
        # Calculate angular velocities
        if len(headings) >= 2:
            # Ensure heading differences are between -pi and pi
            heading_diffs = np.diff(headings)
            heading_diffs = np.where(heading_diffs > np.pi, heading_diffs - 2*np.pi, heading_diffs)
            heading_diffs = np.where(heading_diffs < -np.pi, heading_diffs + 2*np.pi, heading_diffs)
            
            time_diffs = timestamps[1:] - timestamps[:-1]
            time_diffs = np.where(time_diffs == 0, 1e-6, time_diffs)  # Avoid division by zero
            angular_velocities = heading_diffs / time_diffs
            # Pad with zeros for the first timestep
            angular_velocities = np.concatenate([np.zeros(1), angular_velocities])
        else:
            angular_velocities = np.zeros(len(headings))
        
        return {
            'positions': positions,
            'velocities': velocities,
            'accelerations': accelerations,
            'headings': headings,
            'angular_velocities': angular_velocities,
            'timestamps': timestamps
        }
    
    def _process_agent_for_model(self, agent: AgentState) -> Dict:
        """
        Process agent state for model input.
        
        Args:
            agent: Agent state
            
        Returns:
            Dictionary with processed agent features
        """
        return {
            'id': agent.agent_id,
            'type': agent.agent_type,
            'position': agent.position,
            'velocity': agent.velocity,
            'acceleration': agent.acceleration if hasattr(agent, 'acceleration') else np.zeros(2),
            'heading': agent.heading,
            'angular_velocity': agent.angular_velocity if hasattr(agent, 'angular_velocity') else 0.0,
            'length': agent.length,
            'width': agent.width,
            'timestamp': agent.timestamp
        }
    
    def _get_map_features(self, scene_id: str, position: np.ndarray, heading: float, radius: float) -> Dict:
        """
        Get map features around a position.
        
        Args:
            scene_id: Scene ID
            position: Center position
            heading: Agent heading
            radius: Radius to consider
            
        Returns:
            Dictionary with map features
        """
        # Get scene map ID
        scene_map_id = None
        for scene in self.scenes:
            if hasattr(scene, 'map_id'):
                scene_map_id = scene.map_id
                break
        
        if scene_map_id is None:
            return {
                'lanes': [],
                'crosswalks': [],
                'traffic_signals': [],
                'stop_signs': [],
                'speed_bumps': []
            }
        
        # Get map features from map data
        return self.map_data.get_features_in_radius(
            map_id=scene_map_id,
            position=position,
            radius=radius,
            heading=heading
        )
    
    def _get_traffic_features(self, scene_id: str, agent_id: str, scene: SceneState) -> Dict:
        """
        Get traffic rule features for an agent.
        
        Args:
            scene_id: Scene ID
            agent_id: Agent ID
            scene: Current scene state
            
        Returns:
            Dictionary with traffic features
        """
        if not hasattr(self, 'traffic_rules') or self.traffic_rules is None:
            return {
                'speed_limit': 13.89,  # 50 km/h default
                'is_intersection': False,
                'has_right_of_way': True,
                'is_crosswalk_ahead': False,
                'priority_agents': []
            }
        
        # Get agent state
        agent_state = scene.agents.get(agent_id)
        if agent_state is None:
            return {
                'speed_limit': 13.89,
                'is_intersection': False,
                'has_right_of_way': True,
                'is_crosswalk_ahead': False,
                'priority_agents': []
            }
        
        # Get map ID
        map_id = scene.map_id if hasattr(scene, 'map_id') else None
        
        # Get traffic rules
        return self.traffic_rules.get_rules_for_agent(
            map_id=map_id,
            agent_state=agent_state,
            surrounding_agents=[a for a_id, a in scene.agents.items() if a_id != agent_id]
        )
    
    def _get_weather_features(self, scene: SceneState) -> Dict:
        """
        Get weather features for a scene.
        
        Args:
            scene: Scene state
            
        Returns:
            Dictionary with weather features
        """
        if not hasattr(scene, 'weather_info') or scene.weather_info is None:
            return {
                'weather_type': 'clear',
                'temperature': 20.0,
                'precipitation': 0.0,
                'visibility': 10000.0,
                'wind_speed': 0.0,
                'is_day': True
            }
        
        # Get weather info from scene
        weather_info = scene.weather_info
        
        # Get time info from scene
        time_info = scene.time_info if hasattr(scene, 'time_info') else {
            'time_period': 'day',
            'is_day': True,
            'time_of_day': 12.0
        }
        
        # Combine weather and time info
        return {
            'weather_type': weather_info.get('weather_type', 'clear'),
            'temperature': weather_info.get('temperature', 20.0),
            'precipitation': weather_info.get('precipitation', 0.0), 
            'visibility': weather_info.get('visibility', 10000.0),
            'wind_speed': weather_info.get('wind_speed', 0.0),
            'is_day': time_info.get('is_day', True),
            'time_of_day': time_info.get('time_of_day', 12.0)
        }
    
    def get_scene_by_id(self, scene_id: str) -> List[SceneState]:
        """
        Get all scene states for a specific scene ID.
        
        Args:
            scene_id: Scene ID
            
        Returns:
            List of scene states for the scene
        """
        scene_states = []
        for scene in self.scenes:
            if hasattr(scene, 'scene_id') and scene.scene_id == scene_id:
                scene_states.append(scene)
        
        # Sort by timestamp
        scene_states.sort(key=lambda x: x.timestamp)
        
        return scene_states
    
    def get_agent_trajectory(self, scene_id: str, agent_id: str) -> List[AgentState]:
        """
        Get the trajectory of a specific agent in a scene.
        
        Args:
            scene_id: Scene ID
            agent_id: Agent ID
            
        Returns:
            List of agent states for the agent
        """
        if scene_id not in self.trajectories or agent_id not in self.trajectories[scene_id]:
            return []
        
        return self.trajectories[scene_id][agent_id]
    
    def get_dataset_statistics(self) -> Dict:
        """
        Get statistics about the dataset.
        
        Returns:
            Dictionary with dataset statistics
        """
        # Count agents by type
        agent_types_count = {}
        for scene_id in self.scene_ids:
            for agent_id in self.agent_ids[scene_id]:
                if scene_id in self.trajectories and agent_id in self.trajectories[scene_id]:
                    agent_type = self.trajectories[scene_id][agent_id][0].agent_type
                    agent_types_count[agent_type] = agent_types_count.get(agent_type, 0) + 1
        
        # Calculate trajectory statistics
        traj_lengths = []
        durations = []
        avg_speeds = []
        
        for scene_id in self.scene_ids:
            for agent_id in self.agent_ids[scene_id]:
                if scene_id in self.trajectories and agent_id in self.trajectories[scene_id]:
                    traj = self.trajectories[scene_id][agent_id]
                    traj_lengths.append(len(traj))
                    
                    if len(traj) > 1:
                        # Calculate duration
                        duration = traj[-1].timestamp - traj[0].timestamp
                        durations.append(duration)
                        
                        # Calculate average speed
                        total_dist = 0
                        for i in range(1, len(traj)):
                            total_dist += np.linalg.norm(traj[i].position - traj[i-1].position)
                        
                        avg_speed = total_dist / duration if duration > 0 else 0
                        avg_speeds.append(avg_speed)
        
        return {
            'num_scenes': len(self.scene_ids),
            'num_agents': sum(len(agents) for agents in self.agent_ids.values()),
            'agent_types': agent_types_count,
            'avg_trajectory_length': np.mean(traj_lengths) if traj_lengths else 0,
            'min_trajectory_length': min(traj_lengths) if traj_lengths else 0,
            'max_trajectory_length': max(traj_lengths) if traj_lengths else 0,
            'avg_trajectory_duration': np.mean(durations) if durations else 0,
            'avg_speed': np.mean(avg_speeds) if avg_speeds else 0
        }
    
    def get_multimodal_samples(self, difficulty: str = 'medium', num_samples: int = 10) -> List[Dict]:
        """
        Get samples that require multimodal prediction (e.g., at intersections).
        
        Args:
            difficulty: Difficulty level ('easy', 'medium', 'hard')
            num_samples: Number of samples to return
            
        Returns:
            List of samples that require multimodal prediction
        """
        # Identify potentially multimodal scenarios
        multimodal_samples = []
        
        for scene_id in self.scene_ids:
            for agent_id in self.agent_ids[scene_id]:
                if scene_id not in self.trajectories or agent_id not in self.trajectories[scene_id]:
                    continue
                
                trajectory = self.trajectories[scene_id][agent_id]
                
                # Skip trajectories that are too short
                if len(trajectory) < self.config.history_len + self.config.future_len:
                    continue
                
                # Check if trajectory has complex behavior
                is_multimodal = self._is_potentially_multimodal(trajectory)
                
                if is_multimodal:
                    # Find the corresponding index in the dataset
                    idx = 0
                    for s_idx, s_id in enumerate(self.scene_ids):
                        if s_id == scene_id:
                            for a_idx, a_id in enumerate(self.agent_ids[s_id]):
                                if a_id == agent_id:
                                    idx = sum(len(self.agent_ids[self.scene_ids[i]]) for i in range(s_idx)) + a_idx
                                    break
                            break
                    
                    # Get sample
                    sample = self.__getitem__(idx)
                    
                    # Add metadata
                    sample['is_multimodal'] = True
                    sample['scenario_difficulty'] = self._assess_difficulty(trajectory, scene_id)
                    
                    multimodal_samples.append(sample)
                    
                    if len(multimodal_samples) >= num_samples:
                        break
            
            if len(multimodal_samples) >= num_samples:
                break
        
        # Filter by difficulty
        filtered_samples = [s for s in multimodal_samples if s['scenario_difficulty'] == difficulty]
        
        # If not enough samples of the specified difficulty, return what we have
        if len(filtered_samples) < num_samples:
            return multimodal_samples[:num_samples]
        
        return filtered_samples[:num_samples]
    
    def _is_potentially_multimodal(self, trajectory: List[AgentState]) -> bool:
        """
        Check if a trajectory potentially requires multimodal prediction.
        
        Args:
            trajectory: Agent trajectory
            
        Returns:
            Boolean indicating if the trajectory is potentially multimodal
        """
        if len(trajectory) < 10:
            return False
        
        # Check for significant changes in heading
        headings = np.array([state.heading for state in trajectory])
        heading_diffs = np.diff(headings)
        
        # Normalize heading differences to [-π, π]
        heading_diffs = np.where(heading_diffs > np.pi, heading_diffs - 2*np.pi, heading_diffs)
        heading_diffs = np.where(heading_diffs < -np.pi, heading_diffs + 2*np.pi, heading_diffs)
        
        # Calculate absolute heading changes
        abs_heading_changes = np.abs(heading_diffs)
        
        # Check if there are significant heading changes
        if np.max(abs_heading_changes) > 0.5:  # ~30 degrees
            return True
        
        # Check for significant speed changes
        speeds = np.array([np.linalg.norm(state.velocity) for state in trajectory])
        speed_diffs = np.diff(speeds)
        
        # Check if there are significant acceleration or deceleration events
        if np.max(np.abs(speed_diffs)) > 2.0:  # 2 m/s^2
            return True
        
        return False
    
    def _assess_difficulty(self, trajectory: List[AgentState], scene_id: str) -> str:
        """
        Assess the difficulty of a scenario.
        
        Args:
            trajectory: Agent trajectory
            scene_id: Scene ID
            
        Returns:
            Difficulty level ('easy', 'medium', 'hard')
        """
        # Count surrounding agents
        surrounding_count = 0
        for scene in self.scenes:
            if hasattr(scene, 'scene_id') and scene.scene_id == scene_id:
                surrounding_count = max(surrounding_count, len(scene.agents) - 1)  # -1 for the agent itself
        
        # Check trajectory complexity
        headings = np.array([state.heading for state in trajectory])
        heading_diffs = np.diff(headings)
        
        # Normalize heading differences to [-π, π]
        heading_diffs = np.where(heading_diffs > np.pi, heading_diffs - 2*np.pi, heading_diffs)
        heading_diffs = np.where(heading_diffs < -np.pi, heading_diffs + 2*np.pi, heading_diffs)
        
        max_heading_change = np.max(np.abs(heading_diffs))
        
        # Check speed changes
        speeds = np.array([np.linalg.norm(state.velocity) for state in trajectory])
        speed_diffs = np.diff(speeds)
        max_speed_change = np.max(np.abs(speed_diffs))
        
        # Assess difficulty
        if surrounding_count > 5 and (max_heading_change > 1.0 or max_speed_change > 3.0):
            return 'hard'
        elif surrounding_count > 2 and (max_heading_change > 0.5 or max_speed_change > 2.0):
            return 'medium'
        else:
            return 'easy'
    
    def create_data_loader(self, batch_size: int = 32, shuffle: bool = True, num_workers: int = 4) -> DataLoader:
        """
        Create a DataLoader for this dataset.
        
        Args:
            batch_size: Batch size
            shuffle: Whether to shuffle the data
            num_workers: Number of worker processes
            
        Returns:
            DataLoader for this dataset
        """
        return DataLoader(
            dataset=self,
            batch_size=batch_size,
            shuffle=shuffle,
            num_workers=num_workers,
            collate_fn=self._collate_fn
        )
    
    def _collate_fn(self, batch: List[Dict]) -> Dict[str, Any]:
        """
        Custom collate function for batching samples.
        
        Args:
            batch: List of samples
            
        Returns:
            Batched samples
        """
        # Get batch size
        batch_size = len(batch)
        
        # Get maximum sequence lengths
        max_history_len = max(len(sample['history']['positions']) for sample in batch)
        max_future_len = max(len(sample['future']['positions']) for sample in batch)
        max_surrounding_agents = max(len(sample['surrounding_agents']) for sample in batch)
        
        # Initialize batch tensors with padding
        history_positions = np.zeros((batch_size, max_history_len, 2))
        history_velocities = np.zeros((batch_size, max_history_len, 2))
        history_accelerations = np.zeros((batch_size, max_history_len, 2))
        history_headings = np.zeros((batch_size, max_history_len))
        history_masks = np.zeros((batch_size, max_history_len), dtype=bool)
        
        future_positions = np.zeros((batch_size, max_future_len, 2))
        future_velocities = np.zeros((batch_size, max_future_len, 2))
        future_headings = np.zeros((batch_size, max_future_len))
        future_masks = np.zeros((batch_size, max_future_len), dtype=bool)
        
        surrounding_positions = np.zeros((batch_size, max_surrounding_agents, 2))
        surrounding_velocities = np.zeros((batch_size, max_surrounding_agents, 2))
        surrounding_headings = np.zeros((batch_size, max_surrounding_agents))
        surrounding_types = np.zeros((batch_size, max_surrounding_agents), dtype=int)
        surrounding_masks = np.zeros((batch_size, max_surrounding_agents), dtype=bool)
        
        # Other metadata
        scene_ids = []
        agent_ids = []
        agent_types = []
        map_features_list = []
        traffic_features_list = []
        weather_features_list = []
        
        # Fill batch tensors
        for i, sample in enumerate(batch):
            # History
            h_len = len(sample['history']['positions'])
            if h_len > 0:
                history_positions[i, :h_len] = sample['history']['positions']
                history_velocities[i, :h_len] = sample['history']['velocities']
                history_accelerations[i, :h_len] = sample['history']['accelerations']
                history_headings[i, :h_len] = sample['history']['headings']
                history_masks[i, :h_len] = True
            
            # Future
            f_len = len(sample['future']['positions'])
            if f_len > 0:
                future_positions[i, :f_len] = sample['future']['positions']
                future_velocities[i, :f_len] = sample['future']['velocities']
                future_headings[i, :f_len] = sample['future']['headings']
                future_masks[i, :f_len] = True
            
            # Surrounding agents
            s_len = len(sample['surrounding_agents'])
            for j in range(s_len):
                surrounding_positions[i, j] = sample['surrounding_agents'][j]['position']
                surrounding_velocities[i, j] = sample['surrounding_agents'][j]['velocity']
                surrounding_headings[i, j] = sample['surrounding_agents'][j]['heading']
                # Convert agent type to integer
                if sample['surrounding_agents'][j]['type'] == 'vehicle':
                    surrounding_types[i, j] = 0
                elif sample['surrounding_agents'][j]['type'] == 'pedestrian':
                    surrounding_types[i, j] = 1
                elif sample['surrounding_agents'][j]['type'] == 'cyclist':
                    surrounding_types[i, j] = 2
                else:
                    surrounding_types[i, j] = 3  # Other
                surrounding_masks[i, j] = True
            
            # Metadata
            scene_ids.append(sample['scene_id'])
            agent_ids.append(sample['agent_id'])
            
            # Convert agent type to integer
            if sample['agent_type'] == 'vehicle':
                agent_types.append(0)
            elif sample['agent_type'] == 'pedestrian':
                agent_types.append(1)
            elif sample['agent_type'] == 'cyclist':
                agent_types.append(2)
            else:
                agent_types.append(3)  # Other
            
            # Other features
            map_features_list.append(sample['map_features'])
            traffic_features_list.append(sample['traffic_features'])
            weather_features_list.append(sample['weather_features'])
        
        # Create batch dictionary
        batch_dict = {
            'scene_ids': scene_ids,
            'agent_ids': agent_ids,
            'agent_types': np.array(agent_types),
            
            'history_positions': history_positions,
            'history_velocities': history_velocities,
            'history_accelerations': history_accelerations,
            'history_headings': history_headings,
            'history_masks': history_masks,
            
            'future_positions': future_positions,
            'future_velocities': future_velocities,
            'future_headings': future_headings,
            'future_masks': future_masks,
            
            'surrounding_positions': surrounding_positions,
            'surrounding_velocities': surrounding_velocities,
            'surrounding_headings': surrounding_headings,
            'surrounding_types': surrounding_types,
            'surrounding_masks': surrounding_masks,
            
            'map_features': map_features_list,
            'traffic_features': traffic_features_list,
            'weather_features': weather_features_list
        }
        
        return batch_dict