import streamlit as st
import pandas as pd
import cv2
import numpy as np
from ultralytics import YOLO
from datetime import datetime
import plotly.express as px
import os
import torch
import torch.nn as nn
import sys
import subprocess
from collections import OrderedDict
import yaml
from pathlib import Path

# Whitelist safe modules for unpickling
SAFE_MODULES = {
    'torch.Size',
    'torch.LongStorage',
    'torch.HalfStorage',
    'torch.FloatStorage',
    'torch.nn.modules.container.Sequential',
    'torch.nn.modules.container.ModuleList',
    'torch.nn.modules.activation.SiLU',
    'torch.nn.modules.conv.Conv2d',
    'torch.nn.modules.batchnorm.BatchNorm2d',
    'torch._utils._rebuild_tensor_v2',
    'torch._utils._rebuild_parameter',
    'collections.OrderedDict',
    'numpy.core.multiarray.scalar',
    'numpy.dtype',
    'ultralytics.nn.modules.Detect',
    'ultralytics.nn.modules.SPPF',
    'ultralytics.nn.modules.DFL',
    'ultralytics.nn.modules.Conv',
    'ultralytics.nn.modules.Bottleneck',
    'ultralytics.nn.modules.C2f',
    'ultralytics.nn.modules.Concat',
    'ultralytics.nn.tasks.DetectionModel',
    'ultralytics.yolo.utils.IterableSimpleNamespace'
}

# Security settings for Streamlit
st.set_option('server.enableXsrfProtection', False)
st.set_option('server.enableCORS', False)

# Set page config
st.set_page_config(
    page_title="License Plate Detection System",
    page_icon="🚗",
    layout="wide"
)

# Initialize YOLO model
@st.cache_resource
def load_model():
    return YOLO('best.pt')

# Load model
model = load_model()

# Function to process video frame
def process_frame(frame, model):
    results = model.predict(frame)
    detections = pd.DataFrame(results[0].boxes.data).astype("float")
    return detections

# Initialize session state for data storage
if 'detected_plates' not in st.session_state:
    st.session_state.detected_plates = []

# Main interface
st.title("License Plate Detection System")

# Sidebar
st.sidebar.header("Controls")
video_source = st.sidebar.selectbox(
    "Select Video Source",
    ["Uploaded Video", "Sample Video"]
)

if video_source == "Uploaded Video":
    uploaded_file = st.sidebar.file_uploader("Upload Video", type=['mp4', 'avi', 'mov'])
    if uploaded_file:
        # Save uploaded file temporarily
        temp_file = "temp_video.mp4"
        with open(temp_file, "wb") as f:
            f.write(uploaded_file.read())
        video_path = temp_file
    else:
        video_path = None
else:
    video_path = "mycarplate.mp4"  # Default sample video

# Main content area
col1, col2 = st.columns([2, 1])

with col1:
    if video_path:
        # Video processing
        cap = cv2.VideoCapture(video_path)
        
        # Progress bar
        progress_bar = st.progress(0)
        frame_placeholder = st.empty()
        
        # Process video
        frame_count = 0
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
                
            frame_count += 1
            progress = int((frame_count / total_frames) * 100)
            progress_bar.progress(progress)
            
            # Process every 3rd frame for efficiency
            if frame_count % 3 == 0:
                detections = process_frame(frame, model)
                
                if len(detections) > 0:
                    for _, row in detections.iterrows():
                        # Extract detection information
                        confidence = row[4]
                        x1, y1, x2, y2 = map(int, row[:4])
                        
                        # Draw detection box
                        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                        
                        # Save detection info
                        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                        st.session_state.detected_plates.append({
                            'timestamp': current_time,
                            'confidence': float(confidence),
                            'box': [x1, y1, x2, y2]
                        })
            
            # Display the frame
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame_placeholder.image(frame_rgb, channels="RGB", use_column_width=True)
            
        cap.release()

with col2:
    # Detection statistics
    st.subheader("Detection Statistics")
    if st.session_state.detected_plates:
        df = pd.DataFrame(st.session_state.detected_plates)
        
        # Display metrics
        st.metric("Total Detections", len(df))
        st.metric("Average Confidence", f"{df['confidence'].mean():.2%}")
        
        # Plot confidence distribution
        fig = px.histogram(df, x='confidence', title='Confidence Distribution')
        st.plotly_chart(fig, use_container_width=True)
        
        # Recent detections table
        st.subheader("Recent Detections")
        st.dataframe(df.tail(10)[['timestamp', 'confidence']])
    else:
        st.info("No detections yet. Upload a video to begin processing.")