import gradio as gr
import torch
import numpy as np
import cv2
import sqlite3
from PIL import Image, ImageDraw
from transformers import YolosImageProcessor, YolosForObjectDetection
from datetime import datetime

# -----------------------------
# Load Model
# -----------------------------
processor = YolosImageProcessor.from_pretrained(
    "nickmuchi/yolos-small-finetuned-license-plate-detection"
)
model = YolosForObjectDetection.from_pretrained(
    "nickmuchi/yolos-small-finetuned-license-plate-detection"
)
model.eval()

# -----------------------------
# Database
# -----------------------------
conn = sqlite3.connect("vehicles.db", check_same_thread=False)
cursor = conn.cursor()
cursor.execute("""
CREATE TABLE IF NOT EXISTS vehicles (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    time TEXT,
    plate_id TEXT,
    vehicle_type TEXT
)
""")
conn.commit()

# -----------------------------
# Plate Color Classifier
# -----------------------------
def classify_plate_color(plate_img):
    img = np.array(plate_img)
    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)

    green = cv2.inRange(hsv, (35,40,40), (85,255,255))
    yellow = cv2.inRange(hsv, (15,50,50), (35,255,255))

    if np.sum(green) > np.sum(yellow):
        return "EV"
    elif np.sum(yellow) > 0:
        return "Commercial"
    return "Personal"

# -----------------------------
# Bar Chart (OpenCV)
# -----------------------------
def generate_bar_chart():
    cursor.execute("""
        SELECT vehicle_type, COUNT(*) FROM vehicles
        GROUP BY vehicle_type
    """)
    data = cursor.fetchall()

    chart = np.ones((300, 400, 3), dtype=np.uint8) * 255
    cv2.putText(chart, "Vehicle Count Report", (60,30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0,0,0), 2)

    max_count = max([c for _, c in data], default=1)
    x = 70

    for vtype, count in data:
        bar_height = int((count / max_count) * 180)
        cv2.rectangle(chart, (x,2