generate_training_data.py

"""
Training data generation using current system architecture
Integrates with existing WorldConfig, DrivingAgent, and FeatureExtractor
"""

import sys
import time
import carla
import json
import os
from datetime import datetime

sys.path.append(r"C:\Users\eliav\Desktop\Uni\Workshop\CARLA_0.9.11\WindowsNoEditor\PythonAPI\carla")
from agents.navigation.local_planner import RoadOption

# Import our existing modules
from config import WorldConfig
from graph_builder import build_graph
from route_gen import plan_diverse_routes
from feature_extractor import RouteFeatureExtractor
from driving_agent import DrivingAgent


class TrainingDataGenerator:
    """Generate training data using our current taxi system architecture"""
    
    def __init__(self, training_data_dir="training_data"):
        self.training_data_dir = training_data_dir
        self.world_config = None
        self.world = None
        self.graph = None
        self.extractor = None
        self.taxi_vehicle = None
        self.driving_agent = None
        self.spawn_points = None
        self.route_counter = 0
        
        # Ensure training data directory exists
        if not os.path.exists(training_data_dir):
            os.makedirs(training_data_dir)
            print(f"📁 Created training data directory: {training_data_dir}")
    
    def setup_world(self):
        """Setup world using our existing WorldConfig"""
        print("=== SETTING UP WORLD FOR TRAINING DATA ===")
        
        # Use our existing WorldConfig
        self.world_config = WorldConfig()
        
        if not self.world_config.connect_to_carla():
            print("❌ Failed to connect to CARLA!")
            return False
        
        # Set weather conditions for training variety
        weather_conditions = ["clear_day", "overcast", "rain", "night", "foggy"]
        selected_weather = "clear_day"  # You can randomize this
        self.world_config.setup_weather(selected_weather)
        print(f"🌤️ Weather set to: {selected_weather}")
        
        # Clean up existing vehicles
        self.world_config.cleanup_existing_vehicles()
        
        # Spawn traffic vehicles for realistic conditions
        print("🚗 Spawning traffic vehicles...")
        traffic_vehicles = self.world_config.spawn_traffic_vehicles(num_vehicles=15)
        print(f"✅ Spawned {len(traffic_vehicles)} traffic vehicles")
        
        self.world = self.world_config.get_world()
        
        # Build graph using our existing builder
        print("\n=== BUILDING ROAD NETWORK GRAPH ===")
        self.graph = build_graph(self.world, resolution=2.0)
        print(f"📊 Graph built: {len(self.graph.nodes)} nodes, {len(self.graph.edges)} edges")
        
        # Initialize feature extractor
        self.extractor = RouteFeatureExtractor(self.world)
        
        # Get spawn points
        self.spawn_points = self.world.get_map().get_spawn_points()
        print(f"📍 Available spawn points: {len(self.spawn_points)}")
        
        return True
    
    def spawn_training_taxi(self):
        """Spawn a taxi using our current system approach"""
        print("\n=== SPAWNING TRAINING TAXI ===")
        
        blueprint_library = self.world.get_blueprint_library()
        
        # Use same taxi model as in our system
        taxi_blueprint_name = "vehicle.toyota.prius"
        
        try:
            vehicle_bp = blueprint_library.filter(taxi_blueprint_name)[0]
            print(f"Using {taxi_blueprint_name} as training taxi")
        except:
            # Fallback to any vehicle
            vehicle_bp = blueprint_library.filter("vehicle.*")[0]
            print(f"Using fallback vehicle: {vehicle_bp.id}")
        
        # Spawn at first spawn point
        spawn_point = self.spawn_points[0]
        self.taxi_vehicle = self.world.try_spawn_actor(vehicle_bp, spawn_point)
        
        if not self.taxi_vehicle:
            print("❌ Failed to spawn training taxi!")
            return False
        
        print(f"✅ Training taxi spawned at spawn point 0: ({spawn_point.location.x:.0f}, {spawn_point.location.y:.0f})")
        
        # Create DrivingAgent using our existing class
        self.driving_agent = DrivingAgent(self.world, self.taxi_vehicle)
        self.driving_agent.setup_agent(behavior="normal", ignore_traffic_lights=False)
        print("🤖 DrivingAgent initialized")
        
        # Position camera for overview
        spectator = self.world.get_spectator()
        camera_transform = carla.Transform(
            carla.Location(x=0, y=0, z=100),
            carla.Rotation(pitch=-90, yaw=0)
        )
        spectator.set_transform(camera_transform)
        
        return True
    
    def generate_training_routes(self, num_routes=50):
        """Generate multiple training routes with diverse characteristics"""
        print(f"\n🎯 GENERATING {num_routes} TRAINING ROUTES")
        print("=" * 60)
        
        successful_routes = 0
        failed_routes = 0
        
        # Predefined interesting route combinations for variety
        route_configs = [
            # Short routes
            {"start": 0, "end": 5, "type": "short"},
            {"start": 10, "end": 15, "type": "short"},
            {"start": 20, "end": 25, "type": "short"},
            
            # Medium routes
            {"start": 0, "end": 50, "type": "medium"},
            {"start": 25, "end": 100, "type": "medium"},
            {"start": 75, "end": 150, "type": "medium"},
            
            # Long routes
            {"start": 0, "end": 200, "type": "long"},
            {"start": 50, "end": 230, "type": "long"},
            {"start": 100, "end": 10, "type": "long"},
            
            # Cross-city routes (if available)
            {"start": 1, "end": 240, "type": "cross_city"},
            {"start": 30, "end": 180, "type": "cross_city"},
        ]
        
        # Generate additional random routes to reach target number
        import random
        while len(route_configs) < num_routes:
            start_idx = random.randint(0, min(250, len(self.spawn_points) - 1))
            end_idx = random.randint(0, min(250, len(self.spawn_points) - 1))
            
            # Ensure different start and end
            if start_idx != end_idx:
                # Classify route by distance
                start_loc = self.spawn_points[start_idx].location
                end_loc = self.spawn_points[end_idx].location
                distance = start_loc.distance(end_loc)
                
                if distance < 500:
                    route_type = "short"
                elif distance < 1500:
                    route_type = "medium"
                else:
                    route_type = "long"
                
                route_configs.append({
                    "start": start_idx, 
                    "end": end_idx, 
                    "type": route_type
                })
        
        # Execute routes
        for i, route_config in enumerate(route_configs[:num_routes]):
            route_num = i + 1
            start_idx = route_config["start"]
            end_idx = route_config["end"]
            route_type = route_config["type"]
            
            print(f"\n🚗 TRAINING ROUTE {route_num}/{num_routes} ({route_type.upper()})")
            print(f"📍 From spawn {start_idx} → spawn {end_idx}")
            
            success = self.execute_training_route(
                route_num, start_idx, end_idx, route_type
            )
            
            if success:
                successful_routes += 1
                print(f"✅ Route {route_num} completed successfully")
            else:
                failed_routes += 1
                print(f"❌ Route {route_num} failed")
            
            # Brief pause between routes
            if route_num < num_routes:
                print("⏸️ Pausing 3 seconds before next route...")
                time.sleep(3)
        
        # Summary
        print(f"\n📊 TRAINING DATA GENERATION SUMMARY")
        print("=" * 50)
        print(f"Total routes attempted: {num_routes}")
        print(f"Successful routes: {successful_routes}")
        print(f"Failed routes: {failed_routes}")
        print(f"Success rate: {(successful_routes/num_routes)*100:.1f}%")
        print(f"Training data saved in: {self.training_data_dir}")
        
        return successful_routes > 0
    
    def execute_training_route(self, route_num, start_spawn_idx, end_spawn_idx, route_type):
        """Execute a single training route and collect data"""
        
        # Get locations
        start_location = self.spawn_points[start_spawn_idx].location
        end_location = self.spawn_points[end_spawn_idx].location
        current_location = self.taxi_vehicle.get_location()
        
        print(f"Current taxi location: ({current_location.x:.0f}, {current_location.y:.0f})")
        print(f"Target: ({end_location.x:.0f}, {end_location.y:.0f})")
        
        # 1. Plan route using our existing route planner
        print("🗺️ Planning route...")
        candidate_routes = plan_diverse_routes(
            self.graph,
            current_location,
            end_location,
            threshold=0.2,
            max_routes=3
        )
        
        if not candidate_routes:
            print("❌ No route found!")
            return False
        
        # Use the shortest route for training consistency
        route_name, route_nodes = candidate_routes[0]
        print(f"📊 Using route: {route_name} with {len(route_nodes)} nodes")
        
        # 2. Extract features using our existing feature extractor
        print("🔍 Extracting route features...")
        route_waypoints = []
        for node_id in route_nodes:
            if node_id in self.graph.nodes:
                waypoint = self.graph.nodes[node_id]["waypoint"]
                road_option = RoadOption.LANEFOLLOW
                route_waypoints.append((waypoint, road_option))
        
        if not route_waypoints:
            print("❌ Could not convert route to waypoints!")
            return False
        
        features = self.extractor.extract_all_features(route_waypoints)
        
        # 3. Drive the route and measure actual time
        print("🚗 Starting drive...")
        start_time = time.time()
        
        success = self.driving_agent.drive_planned_route(
            route_nodes, 
            self.graph,
            end_spawn_idx,
            end_location,
            max_steps=8000  # Longer timeout for training
        )
        
        actual_drive_time = time.time() - start_time
        
        if not success:
            print(f"❌ Drive failed after {actual_drive_time:.1f}s")
            return False
        
        print(f"✅ Drive completed in {actual_drive_time:.1f}s")
        
        # 4. Save training data
        training_data = {
            "route_id": f"training_route_{route_num:03d}",
            "route_number": route_num,
            "route_type": route_type,
            "timestamp": datetime.now().isoformat(),
            
            # Route info
            "start_spawn_idx": start_spawn_idx,
            "end_spawn_idx": end_spawn_idx,
            "start_location": {
                "x": float(current_location.x),
                "y": float(current_location.y),
                "z": float(current_location.z)
            },
            "end_location": {
                "x": float(end_location.x),
                "y": float(end_location.y),
                "z": float(end_location.z)
            },
            
            # Route characteristics
            "route_name": route_name,
            "route_nodes_count": len(route_nodes),
            "waypoints_count": len(route_waypoints),
            
            # Timing data
            "actual_drive_time_seconds": float(actual_drive_time),
            "drive_successful": success,
            
            # All extracted features
            **{k: self._convert_to_serializable(v) for k, v in features.items()}
        }
        
        # Save to JSON file
        filename = f"training_route_{route_num:03d}_{route_type}.json"
        filepath = os.path.join(self.training_data_dir, filename)
        
        try:
            with open(filepath, 'w') as f:
                json.dump(training_data, f, indent=2)
            print(f"💾 Training data saved: {filename}")
        except Exception as e:
            print(f"❌ Error saving training data: {e}")
            return False
        
        # Print summary
        print(f"📊 Route Summary:")
        print(f"   Distance: {features.get('total_distance', 0):.0f}m")
        print(f"   Traffic lights: {features.get('traffic_lights', 0)}")
        print(f"   Junctions: {features.get('junctions', 0)}")
        print(f"   Turns: {features.get('total_turns', 0)}")
        print(f"   Actual time: {actual_drive_time:.1f}s")
        
        return True
    
    def _convert_to_serializable(self, obj):
        """Convert numpy/CARLA types to JSON serializable types"""
        import numpy as np
        
        if isinstance(obj, (np.float32, np.float64)):
            return float(obj)
        elif isinstance(obj, (np.int32, np.int64)):
            return int(obj)
        elif hasattr(obj, 'item'):  # numpy scalar
            return obj.item()
        elif isinstance(obj, dict):
            return {key: self._convert_to_serializable(value) for key, value in obj.items()}
        elif isinstance(obj, list):
            return [self._convert_to_serializable(item) for item in obj]
        else:
            return obj
    
    def cleanup(self):
        """Clean up resources"""
        print("\n🧹 Cleaning up training environment...")
        
        if self.taxi_vehicle:
            try:
                self.taxi_vehicle.destroy()
                print("🚗 Training taxi destroyed")
            except:
                pass
        
        if self.world_config:
            self.world_config.cleanup()
            print("🌍 World cleaned up")
        
        print("✅ Cleanup complete")


def main():
    """Main function for training data generation"""
    print("🚖 TAXI TRAINING DATA GENERATOR")
    print("Using Current System Architecture")
    print("=" * 60)
    
    # Get user input for number of routes
    try:
        num_routes_input = input("Enter number of training routes to generate (default: 25): ").strip()
        if num_routes_input:
            num_routes = int(num_routes_input)
        else:
            num_routes = 25
            
        if num_routes <= 0 or num_routes > 200:
            print("❌ Please enter a number between 1 and 200")
            return
            
    except ValueError:
        print("❌ Invalid input, using default of 25 routes")
        num_routes = 25
    except KeyboardInterrupt:
        print("\n❌ Cancelled by user")
        return
    
    print(f"🎯 Will generate {num_routes} training routes")
    
    # Initialize generator
    generator = TrainingDataGenerator()
    
    try:
        # Setup world
        if not generator.setup_world():
            print("❌ World setup failed!")
            return
        
        # Spawn taxi
        if not generator.spawn_training_taxi():
            print("❌ Taxi spawn failed!")
            return
        
        print("\n✅ Setup complete! Starting training data generation...")
        
        # Generate training routes
        success = generator.generate_training_routes(num_routes)
        
        if success:
            print("\n🎉 Training data generation completed!")
            print("📊 Check the 'training_data/' directory for JSON files")
            print("💡 Each file contains route features and actual drive time")
        else:
            print("\n❌ Training data generation failed")
        
        print("\nPress Ctrl+C to exit...")
        
        # Keep world alive for inspection
        while True:
            time.sleep(1)
            
    except KeyboardInterrupt:
        print("\n🛑 Training interrupted by user")
    
    finally:
        generator.cleanup()
        print("🏁 Training data generation session ended")


if __name__ == "__main__":
    main()