import pandas as pd
import numpy as np
import requests
import json
import time
from datetime import datetime, timedelta
import plotly.graph_objs as go
import plotly.express as px
from plotly.subplots import make_subplots
import gradio as gr
from dataclasses import dataclass
from typing import List, Dict, Tuple
import threading
import queue
import random

# Configuration
class FleetConfig:
    def __init__(self):
        self.num_vehicles = 50
        self.vehicle_capacity = 4
        self.max_distance = 100  # km
        self.base_cost_per_km = 0.5
        self.weather_impact = {
            'clear': 1.0,
            'rain': 1.2,
            'snow': 1.5,
            'storm': 2.0
        }
        self.traffic_impact = {
            'low': 1.0,
            'medium': 1.3,
            'high': 1.8,
            'severe': 2.5
        }

@dataclass
class Vehicle:
    id: int
    location: Tuple[float, float]  # lat, lng
    status: str  # 'available', 'busy', 'maintenance'
    capacity: int
    current_load: int
    total_distance: float
    earnings: float
    last_update: datetime

@dataclass
class Demand:
    id: int
    pickup_location: Tuple[float, float]
    dropoff_location: Tuple[float, float]
    passengers: int
    priority: int  # 1-5, 5 being highest
    timestamp: datetime
    status: str  # 'pending', 'assigned', 'completed'

@dataclass
class WeatherData:
    location: Tuple[float, float]
    condition: str
    temperature: float
    wind_speed: float
    visibility: float
    timestamp: datetime

@dataclass
class TrafficData:
    location: Tuple[float, float]
    congestion_level: str  # 'low', 'medium', 'high', 'severe'
    average_speed: float
    delay_minutes: float
    timestamp: datetime

class FleetOptimizer:
    def __init__(self):
        self.config = FleetConfig()
        self.vehicles = []
        self.demands = []
        self.weather_data = {}
        self.traffic_data = {}
        self.simulation_running = False
        self.data_queue = queue.Queue()
        
        # Initialize vehicles
        self._initialize_vehicles()
        
        # Simulation parameters
        self.simulation_time = datetime.now()
        self.time_step = 60  # seconds
        
    def _initialize_vehicles(self):
        """Initialize fleet vehicles with random locations"""
        for i in range(self.config.num_vehicles):
            vehicle = Vehicle(
                id=i,
                location=(random.uniform(40.7, 40.8), random.uniform(-74.0, -73.9)),  # NYC area
                status='available',
                capacity=self.config.vehicle_capacity,
                current_load=0,
                total_distance=0.0,
                earnings=0.0,
                last_update=datetime.now()
            )
            self.vehicles.append(vehicle)
    
    def generate_demand(self):
        """Generate realistic demand patterns"""
        # Simulate demand hotspots
        hotspots = [
            (40.7589, -73.9851),  # Times Square
            (40.7505, -73.9934),  # Penn Station
            (40.7527, -73.9772),  # Grand Central
            (40.7484, -73.9857),  # Empire State Building
            (40.7587, -73.9787),  # Rockefeller Center
        ]
        
        # Generate demand based on time patterns
        hour = self.simulation_time.hour
        base_demand_rate = 0.3  # Increased base rate
        
        # Peak hours (7-9 AM, 5-7 PM)
        if 7 <= hour <= 9 or 17 <= hour <= 19:
            base_demand_rate = 0.6  # Higher peak rate
        elif 22 <= hour or hour <= 6:
            base_demand_rate = 0.1  # Higher night rate
            
        # Generate multiple demands per step
        num_demands = random.choices([0, 1, 2], weights=[0.4, 0.4, 0.2])[0]
        
        for _ in range(num_demands):
            if random.random() < base_demand_rate:
                pickup = random.choice(hotspots)
                dropoff = (
                    pickup[0] + random.uniform(-0.01, 0.01),
                    pickup[1] + random.uniform(-0.01, 0.01)
                )
                
                demand = Demand(
                    id=len(self.demands),
                    pickup_location=pickup,
                    dropoff_location=dropoff,
                    passengers=random.randint(1, 4),
                    priority=random.randint(1, 5),
                    timestamp=self.simulation_time,
                    status='pending'
                )
                self.demands.append(demand)
    
    def update_weather_data(self):
        """Simulate weather data updates"""
        # Simulate weather conditions
        conditions = ['clear', 'rain', 'snow', 'storm']
        weights = [0.7, 0.2, 0.08, 0.02]  # Mostly clear weather
        
        for vehicle in self.vehicles:
            condition = random.choices(conditions, weights=weights)[0]
            weather = WeatherData(
                location=vehicle.location,
                condition=condition,
                temperature=random.uniform(-5, 35),
                wind_speed=random.uniform(0, 30),
                visibility=random.uniform(0.1, 10),
                timestamp=self.simulation_time
            )
            self.weather_data[vehicle.id] = weather
    
    def update_traffic_data(self):
        """Simulate traffic data updates"""
        # Traffic patterns based on time
        hour = self.simulation_time.hour
        if 7 <= hour <= 9 or 17 <= hour <= 19:
            congestion_levels = ['medium', 'high', 'severe']
            weights = [0.3, 0.5, 0.2]
        else:
            congestion_levels = ['low', 'medium', 'high']
            weights = [0.6, 0.3, 0.1]
            
        for vehicle in self.vehicles:
            congestion = random.choices(congestion_levels, weights=weights)[0]
            traffic = TrafficData(
                location=vehicle.location,
                congestion_level=congestion,
                average_speed=random.uniform(10, 60),
                delay_minutes=random.uniform(0, 15),
                timestamp=self.simulation_time
            )
            self.traffic_data[vehicle.id] = traffic
    
    def calculate_distance(self, loc1, loc2):
        """Calculate distance between two locations (simplified)"""
        return np.sqrt((loc1[0] - loc2[0])**2 + (loc1[1] - loc2[1])**2) * 111  # km
    
    def calculate_cost(self, vehicle_id, pickup_loc, dropoff_loc):
        """Calculate cost considering weather and traffic"""
        distance = self.calculate_distance(pickup_loc, dropoff_loc)
        
        # Get weather and traffic impacts
        weather = self.weather_data.get(vehicle_id)
        traffic = self.traffic_data.get(vehicle_id)
        
        weather_multiplier = self.config.weather_impact.get(weather.condition, 1.0) if weather else 1.0
        traffic_multiplier = self.config.traffic_impact.get(traffic.congestion_level, 1.0) if traffic else 1.0
        
        total_cost = distance * self.config.base_cost_per_km * weather_multiplier * traffic_multiplier
        return total_cost, distance
    
    def optimize_vehicle_allocation(self):
        """AI-powered vehicle allocation optimization"""
        pending_demands = [d for d in self.demands if d.status == 'pending']
        available_vehicles = [v for v in self.vehicles if v.status == 'available']
        
        if not pending_demands or not available_vehicles:
            return
        
        # Create cost matrix for assignment problem
        cost_matrix = []
        for vehicle in available_vehicles:
            vehicle_costs = []
            for demand in pending_demands:
                cost, distance = self.calculate_cost(vehicle.id, vehicle.location, demand.pickup_location)
                # Add penalty for distance and priority
                penalty = distance * 0.1 + (6 - demand.priority) * 2
                total_cost = cost + penalty
                vehicle_costs.append(total_cost)
            cost_matrix.append(vehicle_costs)
        
        # Simple greedy assignment (can be improved with Hungarian algorithm)
        assignments = []
        used_vehicles = set()
        used_demands = set()
        
        # Sort demands by priority (highest first)
        sorted_demands = sorted(pending_demands, key=lambda x: x.priority, reverse=True)
        
        for demand in sorted_demands:
            best_vehicle = None
            best_cost = float('inf')
            
            for i, vehicle in enumerate(available_vehicles):
                if i in used_vehicles:
                    continue
                    
                if vehicle.current_load + demand.passengers <= vehicle.capacity:
                    cost = cost_matrix[i][pending_demands.index(demand)]
                    if cost < best_cost:
                        best_cost = cost
                        best_vehicle = i
            
            if best_vehicle is not None:
                assignments.append((available_vehicles[best_vehicle], demand))
                used_vehicles.add(best_vehicle)
                used_demands.add(demand.id)
        
        # Execute assignments (limit to prevent all vehicles being assigned at once)
        max_assignments = min(len(assignments), 5)  # Max 5 assignments per step
        for vehicle, demand in assignments[:max_assignments]:
            self._assign_vehicle_to_demand(vehicle, demand)
    
    def _assign_vehicle_to_demand(self, vehicle, demand):
        """Assign vehicle to demand and update status"""
        vehicle.status = 'busy'
        vehicle.current_load = demand.passengers
        demand.status = 'assigned'
        
        # Calculate trip details
        pickup_distance = self.calculate_distance(vehicle.location, demand.pickup_location)
        trip_distance = self.calculate_distance(demand.pickup_location, demand.dropoff_location)
        
        # Update vehicle metrics
        vehicle.total_distance += pickup_distance + trip_distance
        vehicle.earnings += self.calculate_cost(vehicle.id, demand.pickup_location, demand.dropoff_location)[0]
        vehicle.location = demand.dropoff_location
        vehicle.last_update = self.simulation_time
        
        # Simulate trip completion after some time
        completion_time = self.simulation_time + timedelta(minutes=random.randint(5, 20))
        self.data_queue.put(('complete_trip', vehicle.id, completion_time))
    
    def complete_trips(self):
        """Complete trips that have finished"""
        current_time = self.simulation_time
        
        # Check for completed trips
        while not self.data_queue.empty():
            try:
                action, vehicle_id, completion_time = self.data_queue.get_nowait()
                if action == 'complete_trip' and completion_time <= current_time:
                    vehicle = next(v for v in self.vehicles if v.id == vehicle_id)
                    vehicle.status = 'available'
                    vehicle.current_load = 0
            except queue.Empty:
                break
    
    def run_simulation_step(self):
        """Run one simulation step"""
        if not self.simulation_running:
            return
        
        # Update simulation time (advance by 1 hour for more realistic demand patterns)
        self.simulation_time = self.simulation_time + timedelta(hours=1)
        
        # Generate new demand
        self.generate_demand()
        
        # Update weather and traffic data
        self.update_weather_data()
        self.update_traffic_data()
        
        # Complete finished trips
        self.complete_trips()
        
        # Optimize vehicle allocation
        self.optimize_vehicle_allocation()
    
    def start_simulation(self):
        """Start the simulation"""
        self.simulation_running = True
        print("🚗 Fleet optimization simulation started!")
        
        while self.simulation_running:
            self.run_simulation_step()
            time.sleep(1)  # Real-time simulation
    
    def stop_simulation(self):
        """Stop the simulation"""
        self.simulation_running = False
        print("🛑 Simulation stopped")
    
    def get_simulation_stats(self):
        """Get current simulation statistics"""
        total_earnings = sum(v.earnings for v in self.vehicles)
        total_distance = sum(v.total_distance for v in self.vehicles)
        available_vehicles = len([v for v in self.vehicles if v.status == 'available'])
        busy_vehicles = len([v for v in self.vehicles if v.status == 'busy'])
        pending_demands = len([d for d in self.demands if d.status == 'pending'])
        
        return {
            'total_earnings': total_earnings,
            'total_distance': total_distance,
            'available_vehicles': available_vehicles,
            'busy_vehicles': busy_vehicles,
            'pending_demands': pending_demands,
            'simulation_time': self.simulation_time.strftime('%H:%M:%S'),
            'total_demands': len(self.demands)
        }
    
    def create_dashboard(self):
        """Create interactive dashboard"""
        # Vehicle locations
        vehicle_locations = pd.DataFrame([
            {
                'id': v.id,
                'lat': v.location[0],
                'lng': v.location[1],
                'status': v.status,
                'earnings': v.earnings,
                'distance': v.total_distance
            }
            for v in self.vehicles
        ])
        
        # Demand locations
        demand_locations = pd.DataFrame([
            {
                'id': d.id,
                'lat': d.pickup_location[0],
                'lng': d.pickup_location[1],
                'status': d.status,
                'priority': d.priority
            }
            for d in self.demands if d.status in ['pending', 'assigned']
        ])
        
        # Create map
        fig = go.Figure()
        
        # Add vehicle markers
        for status in ['available', 'busy']:
            vehicles = vehicle_locations[vehicle_locations['status'] == status]
            if not vehicles.empty:
                fig.add_trace(go.Scattermapbox(
                    lat=vehicles['lat'],
                    lon=vehicles['lng'],
                    mode='markers',
                    marker=go.scattermapbox.Marker(
                        size=10,
                        color='green' if status == 'available' else 'red'
                    ),
                    name=f'Vehicles ({status})',
                    text=vehicles['id'],
                    hovertemplate='Vehicle %{text}<br>Earnings: $%{customdata[0]:.2f}<br>Distance: %{customdata[1]:.1f}km<extra></extra>',
                    customdata=vehicles[['earnings', 'distance']].values
                ))
        
        # Add demand markers
        if not demand_locations.empty:
            fig.add_trace(go.Scattermapbox(
                lat=demand_locations['lat'],
                lon=demand_locations['lng'],
                mode='markers',
                marker=go.scattermapbox.Marker(
                    size=8,
                    color='blue',
                    symbol='diamond'
                ),
                name='Demands',
                text=demand_locations['id'],
                hovertemplate='Demand %{text}<br>Priority: %{customdata}<extra></extra>',
                customdata=demand_locations['priority']
            ))
        
        fig.update_layout(
            mapbox=dict(
                style='open-street-map',
                center=dict(lat=40.7589, lon=-73.9851),
                zoom=12
            ),
            title='Fleet Optimization Dashboard',
            height=600
        )
        
        return fig

# Global optimizer instance
optimizer = FleetOptimizer()

def start_fleet_simulation():
    """Start the fleet optimization simulation"""
    if not optimizer.simulation_running:
        thread = threading.Thread(target=optimizer.start_simulation, daemon=True)
        thread.start()
        return "🚗 Fleet optimization simulation started! Check the dashboard for real-time updates."
    return "Simulation is already running!"

def stop_fleet_simulation():
    """Stop the fleet optimization simulation"""
    optimizer.stop_simulation()
    return "🛑 Simulation stopped"

def get_fleet_stats():
    """Get current fleet statistics"""
    stats = optimizer.get_simulation_stats()
    return json.dumps(stats, indent=2)

def update_fleet_dashboard():
    """Update the fleet dashboard"""
    return optimizer.create_dashboard()

# Gradio interface
def create_fleet_interface():
    with gr.Blocks(title="Fleet Resource Optimization Simulator", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🚗 Fleet Resource Optimization with AI Agents")
        gr.Markdown("### Dynamic vehicle allocation based on traffic, weather, and demand")
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 🎮 Simulation Controls")
                start_btn = gr.Button("🚀 Start Simulation", variant="primary")
                stop_btn = gr.Button("🛑 Stop Simulation", variant="secondary")
                
                gr.Markdown("### 📊 Real-time Statistics")
                stats_btn = gr.Button("📈 Update Stats")
                stats_output = gr.Textbox(label="Fleet Statistics", lines=10, interactive=False)
                
                gr.Markdown("### ⚙️ Configuration")
                gr.Markdown(f"""
                - **Total Vehicles**: {optimizer.config.num_vehicles}
                - **Vehicle Capacity**: {optimizer.config.vehicle_capacity} passengers
                - **Max Distance**: {optimizer.config.max_distance} km
                - **Base Cost**: ${optimizer.config.base_cost_per_km}/km
                """)
                
            with gr.Column(scale=2):
                gr.Markdown("### 🗺️ Live Fleet Dashboard")
                dashboard_output = gr.Plot(label="Vehicle Locations & Demand")
        
        # Event handlers
        start_btn.click(
            fn=start_fleet_simulation,
            outputs=gr.Textbox(label="Status", lines=2)
        )
        
        stop_btn.click(
            fn=stop_fleet_simulation,
            outputs=gr.Textbox(label="Status", lines=2)
        )
        
        stats_btn.click(
            fn=get_fleet_stats,
            outputs=stats_output
        )
        
        # Auto-refresh dashboard
        demo.load(
            fn=update_fleet_dashboard,
            outputs=dashboard_output
        )
        
        # Periodic updates
        demo.load(
            fn=lambda: None,
            every=5  # Update every 5 seconds
        )
    
    return demo

if __name__ == "__main__":
    demo = create_fleet_interface()
    demo.launch(share=True)