import cv2
import gradio as gr
from ultralytics import YOLO
import tempfile
import os
import numpy as np
from collections import defaultdict, deque
import time
import json
from datetime import datetime
import matplotlib.pyplot as plt
import pandas as pd

# Fix Ultralytics config warning
os.environ['YOLO_CONFIG_DIR'] = os.path.expanduser('~/.config/Ultralytics')
os.makedirs(os.environ['YOLO_CONFIG_DIR'], exist_ok=True)

# Initialize YOLO models
models = {
    "YOLOv8n (Fast)": YOLO("yolov8n.pt"),
    "YOLOv8s (Balanced)": YOLO("yolov8s.pt"),
    "YOLOv8m (Accurate)": YOLO("yolov8m.pt")
}

# Enhanced vehicle classes with more details
VEHICLE_CLASSES = {
    2: {"name": "car", "color": (0, 255, 0), "icon": "🚗"},
    3: {"name": "motorcycle", "color": (0, 255, 255), "icon": "🏍️"},
    5: {"name": "bus", "color": (0, 0, 255), "icon": "🚌"},
    7: {"name": "truck", "color": (255, 165, 0), "icon": "🚛"},
    1: {"name": "bicycle", "color": (255, 0, 255), "icon": "🚲"}
}

class VehicleTracker:
    def __init__(self):
        self.vehicle_counts = defaultdict(int)
        self.tracked_ids = set()
        self.speed_data = defaultdict(list)
        self.position_history = defaultdict(deque)
        self.frame_count = 0
        
    def calculate_speed(self, track_id, current_pos, fps, pixels_per_meter=10):
        """Calculate vehicle speed based on position history"""
        self.position_history[track_id].append((current_pos, self.frame_count))
        
        # Keep only recent positions (last 30 frames)
        if len(self.position_history[track_id]) > 30:
            self.position_history[track_id].popleft()
        
        if len(self.position_history[track_id]) >= 2:
            # Calculate speed using last two positions
            (x1, y1), frame1 = self.position_history[track_id][-2]
            (x2, y2), frame2 = self.position_history[track_id][-1]
            
            # Distance in pixels
            pixel_distance = np.sqrt((x2-x1)**2 + (y2-y1)**2)
            
            # Time difference
            time_diff = (frame2 - frame1) / fps
            
            if time_diff > 0:
                # Speed in pixels per second, convert to km/h (rough estimation)
                speed_pixels_per_sec = pixel_distance / time_diff
                speed_kmh = (speed_pixels_per_sec / pixels_per_meter) * 3.6
                return min(speed_kmh, 200)  # Cap at reasonable speed
        
        return 0

def create_analytics_chart(vehicle_counts):
    """Create a pie chart of vehicle distribution"""
    if not vehicle_counts:
        return None
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
    
    # Pie chart
    labels = [f"{VEHICLE_CLASSES.get(k, {}).get('icon', '')} {k}" for k in vehicle_counts.keys()]
    colors = [np.array(VEHICLE_CLASSES.get(k, {}).get('color', (128, 128, 128)))/255 for k in vehicle_counts.keys()]
    
    ax1.pie(vehicle_counts.values(), labels=labels, colors=colors, autopct='%1.1f%%', startangle=90)
    ax1.set_title('Vehicle Distribution')
    
    # Bar chart
    ax2.bar(vehicle_counts.keys(), vehicle_counts.values(), color=colors)
    ax2.set_title('Vehicle Count by Type')
    ax2.set_xlabel('Vehicle Type')
    ax2.set_ylabel('Count')
    
    plt.tight_layout()
    
    # Save to temp file
    temp_path = tempfile.NamedTemporaryFile(suffix='.png', delete=False).name
    plt.savefig(temp_path, dpi=150, bbox_inches='tight')
    plt.close()
    
    return temp_path

def draw_detection_zone(frame, zone_coords):
    """Draw detection zone on frame"""
    if zone_coords:
        pts = np.array(zone_coords, np.int32)
        cv2.polylines(frame, [pts], True, (255, 255, 0), 2)
        cv2.fillPoly(frame, [pts], (255, 255, 0, 30))

def is_in_zone(center_point, zone_coords):
    """Check if point is inside detection zone"""
    if not zone_coords:
        return True
    return cv2.pointPolygonTest(np.array(zone_coords, np.int32), center_point, False) >= 0

def process_video(input_video, model_choice, confidence_threshold, show_speed, show_trails, 
                 detection_zone, max_frames, output_format):
    """Enhanced video processing with advanced features"""
    
    if input_video is None:
        raise gr.Error("Please upload a video file")
    
    # Get selected model
    model = models[model_choice]
    
    # Create temp files
    input_path = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False).name
    output_path = tempfile.NamedTemporaryFile(suffix=f'.{output_format}', delete=False).name
    
    # Save uploaded file
    with open(input_path, 'wb') as f:
        if isinstance(input_video, str):
            with open(input_video, 'rb') as src:
                f.write(src.read())
        else:
            f.write(input_video)
    
    # Open video
    cap = cv2.VideoCapture(input_path)
    if not cap.isOpened():
        raise gr.Error("Could not open video file")
    
    # Get video properties
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    # Limit frames if specified
    if max_frames > 0:
        total_frames = min(total_frames, max_frames)
    
    # Setup video writer
    fourcc = cv2.VideoWriter_fourcc(*'mp4v') if output_format == 'mp4' else cv2.VideoWriter_fourcc(*'XVID')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    
    # Initialize tracker
    tracker = VehicleTracker()
    trail_points = defaultdict(deque)
    
    # Parse detection zone
    zone_coords = None
    if detection_zone:
        try:
            # Expected format: "x1,y1;x2,y2;x3,y3;x4,y4"
            points = detection_zone.split(';')
            zone_coords = [(int(p.split(',')[0]), int(p.split(',')[1])) for p in points]
        except:
            pass
    
    # Processing loop with progress
    frame_count = 0
    processing_times = []
    
    while cap.isOpened() and frame_count < total_frames:
        start_time = time.time()
        ret, frame = cap.read()
        if not ret:
            break
        
        tracker.frame_count = frame_count
        
        # Draw detection zone
        if zone_coords:
            draw_detection_zone(frame, zone_coords)
        
        # Run detection/tracking
        results = model.track(frame, persist=True, verbose=False, conf=confidence_threshold)
        
        # Process detections
        current_detections = []
        if results[0].boxes is not None and results[0].boxes.id is not None:
            for box, cls, conf, track_id in zip(
                results[0].boxes.xyxy.cpu().numpy(),
                results[0].boxes.cls.cpu().numpy().astype(int),
                results[0].boxes.conf.cpu().numpy(),
                results[0].boxes.id.cpu().numpy().astype(int)
            ):
                if cls in VEHICLE_CLASSES:
                    x1, y1, x2, y2 = map(int, box)
                    center_x, center_y = (x1 + x2) // 2, (y1 + y2) // 2
                    
                    # Check if in detection zone
                    if not is_in_zone((center_x, center_y), zone_coords):
                        continue
                    
                    vehicle_info = VEHICLE_CLASSES[cls]
                    label = vehicle_info["name"]
                    color = vehicle_info["color"]
                    icon = vehicle_info["icon"]
                    
                    # Calculate speed if enabled
                    speed = 0
                    if show_speed:
                        speed = tracker.calculate_speed(track_id, (center_x, center_y), fps)
                    
                    # Draw bounding box with enhanced styling
                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3)
                    
                    # Enhanced label with confidence and speed
                    label_text = f"{icon} {label} #{track_id}"
                    if show_speed and speed > 0:
                        label_text += f" {speed:.1f}km/h"
                    label_text += f" ({conf:.2f})"
                    
                    # Draw label background
                    (text_width, text_height), _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
                    cv2.rectangle(frame, (x1, y1-text_height-10), (x1+text_width, y1), color, -1)
                    cv2.putText(frame, label_text, (x1, y1-5),
                              cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
                    
                    # Draw trails if enabled
                    if show_trails:
                        trail_points[track_id].append((center_x, center_y))
                        if len(trail_points[track_id]) > 50:  # Keep last 50 points
                            trail_points[track_id].popleft()
                        
                        # Draw trail
                        points = list(trail_points[track_id])
                        for i in range(1, len(points)):
                            alpha = i / len(points)
                            cv2.line(frame, points[i-1], points[i], color, max(1, int(3*alpha)))
                    
                    # Count unique vehicles
                    if track_id not in tracker.tracked_ids:
                        tracker.vehicle_counts[label] += 1
                        tracker.tracked_ids.add(track_id)
                    
                    current_detections.append({
                        'id': track_id,
                        'type': label,
                        'confidence': conf,
                        'speed': speed,
                        'position': (center_x, center_y)
                    })
        
        # Draw enhanced statistics panel
        panel_height = 120
        panel = np.zeros((panel_height, width, 3), dtype=np.uint8)
        panel[:] = (40, 40, 40)  # Dark background
        
        # Title
        cv2.putText(panel, "VEHICLE DETECTION ANALYTICS", (10, 25),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
        
        # Vehicle counts
        y_offset = 50
        x_offset = 10
        for vehicle_type, count in tracker.vehicle_counts.items():
            icon = next((v["icon"] for v in VEHICLE_CLASSES.values() if v["name"] == vehicle_type), "")
            text = f"{icon} {vehicle_type.upper()}: {count}"
            cv2.putText(panel, text, (x_offset, y_offset),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
            x_offset += 150
            if x_offset > width - 150:
                x_offset = 10
                y_offset += 25
        
        # Processing info
        processing_time = time.time() - start_time
        processing_times.append(processing_time)
        avg_processing_time = np.mean(processing_times[-30:])  # Last 30 frames
        
        info_text = f"Frame: {frame_count+1}/{total_frames} | FPS: {1/avg_processing_time:.1f} | Active: {len(current_detections)}"
        cv2.putText(panel, info_text, (10, panel_height-10),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 1)
        
        # Combine frame with panel
        frame_with_panel = np.vstack([frame, panel])
        out.write(frame_with_panel)
        
        frame_count += 1
    
    # Cleanup
    cap.release()
    out.release()
    os.unlink(input_path)
    
    # Generate analytics
    analytics_chart = create_analytics_chart(tracker.vehicle_counts)
    
    # Create detailed statistics
    stats = {
        "summary": dict(tracker.vehicle_counts),
        "total_vehicles": sum(tracker.vehicle_counts.values()),
        "processing_info": {
            "total_frames": frame_count,
            "avg_fps": 1/np.mean(processing_times) if processing_times else 0,
            "model_used": model_choice,
            "confidence_threshold": confidence_threshold
        },
        "timestamp": datetime.now().isoformat()
    }
    
    return output_path, stats, analytics_chart

# Enhanced Gradio Interface
with gr.Blocks(theme=gr.themes.Soft(), title="🚗 Advanced Vehicle Detection System") as demo:
    gr.HTML("""
    <div style="text-align: center; padding: 20px;">
        <h1 style="color: #2E86AB; margin-bottom: 10px;">🚗 Advanced Vehicle Detection & Analytics</h1>
        <p style="color: #666; font-size: 16px;">AI-powered vehicle detection with real-time tracking, speed estimation, and comprehensive analytics</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📁 Input Configuration")
            video_input = gr.File(
                label="Upload Video File",
                file_types=[".mp4", ".avi", ".mov", ".mkv"],
                type="binary"
            )
            
            with gr.Accordion("⚙️ Detection Settings", open=True):
                model_choice = gr.Dropdown(
                    choices=list(models.keys()),
                    value="YOLOv8s (Balanced)",
                    label="AI Model",
                    info="Choose speed vs accuracy tradeoff"
                )
                
                confidence_threshold = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.5,
                    step=0.05,
                    label="Confidence Threshold",
                    info="Higher values = fewer false positives"
                )
            
            with gr.Accordion("🎯 Advanced Features", open=False):
                show_speed = gr.Checkbox(
                    label="Speed Estimation",
                    value=True,
                    info="Estimate vehicle speeds (approximate)"
                )
                
                show_trails = gr.Checkbox(
                    label="Motion Trails",
                    value=False,
                    info="Show vehicle movement paths"
                )
                
                detection_zone = gr.Textbox(
                    label="Detection Zone (Optional)",
                    placeholder="x1,y1;x2,y2;x3,y3;x4,y4",
                    info="Define polygon detection area"
                )
                
                max_frames = gr.Number(
                    label="Max Frames to Process",
                    value=0,
                    info="0 = process entire video"
                )
                
                output_format = gr.Dropdown(
                    choices=["mp4", "avi"],
                    value="mp4",
                    label="Output Format"
                )
        
        with gr.Column(scale=2):
            gr.Markdown("### 📊 Results & Analytics")
            
            with gr.Tabs():
                with gr.TabItem("🎥 Processed Video"):
                    video_output = gr.Video(
                        label="Detection Results",
                        show_label=False
                    )
                
                with gr.TabItem("📈 Analytics Dashboard"):
                    with gr.Row():
                        stats_output = gr.JSON(
                            label="Detailed Statistics",
                            show_label=True
                        )
                        analytics_chart = gr.Image(
                            label="Vehicle Distribution",
                            show_label=True
                        )
                
                with gr.TabItem("ℹ️ Help & Examples"):
                    gr.Markdown("""
                    ### 🚀 How to Use:
                    1. **Upload Video**: Support for MP4, AVI, MOV formats
                    2. **Choose Model**: 
                       - YOLOv8n: Fastest processing
                       - YOLOv8s: Balanced speed/accuracy  
                       - YOLOv8m: Highest accuracy
                    3. **Adjust Settings**: Fine-tune confidence and features
                    4. **Process**: Click analyze and wait for results
                    
                    ### 🎯 Advanced Features:
                    - **Speed Estimation**: Approximate vehicle speeds
                    - **Motion Trails**: Visualize movement patterns
                    - **Detection Zones**: Focus on specific areas
                    - **Real-time Analytics**: Live statistics overlay
                    
                    ### 📋 Supported Vehicles:
                    🚗 Cars | 🏍️ Motorcycles | 🚌 Buses | 🚛 Trucks | 🚲 Bicycles
                    """)
    
    # Examples section
    gr.Examples(
        examples=[
            ["examples/traffic.mp4", "YOLOv8s (Balanced)", 0.5, True, False, "", 0, "mp4"],
            ["examples/highway.mp4", "YOLOv8m (Accurate)", 0.6, True, True, "", 0, "mp4"]
        ],
        inputs=[video_input, model_choice, confidence_threshold, show_speed, 
                show_trails, detection_zone, max_frames, output_format],
        label="📎 Example Videos"
    )
    
    # Process button
    analyze_btn = gr.Button(
        "🔍 Analyze Video",
        variant="primary",
        size="lg"
    )
    
    analyze_btn.click(
        fn=process_video,
        inputs=[video_input, model_choice, confidence_threshold, show_speed,
                show_trails, detection_zone, max_frames, output_format],
        outputs=[video_output, stats_output, analytics_chart],
        show_progress=True
    )
    
    # Footer
    gr.HTML("""
    <div style="text-align: center; padding: 20px; color: #666; border-top: 1px solid #eee; margin-top: 20px;">
        <p>🤖 Powered by YOLOv8 & Ultralytics | Built with ❤️ using Gradio</p>
    </div>
    """)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )