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

# -------- Database --------
conn = sqlite3.connect("vehicles.db", check_same_thread=False)
cursor = conn.cursor()
cursor.execute("""
CREATE TABLE IF NOT EXISTS vehicles(
    plate TEXT,
    type TEXT,
    time TEXT
)
""")
conn.commit()

# -------- Models --------
processor = YolosImageProcessor.from_pretrained("nickmuchi/yolos-small-finetuned-license-plate-detection")
model = YolosForObjectDetection.from_pretrained("nickmuchi/yolos-small-finetuned-license-plate-detection")
model.eval()
reader = easyocr.Reader(['en'])

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

    green = np.sum(cv2.inRange(hsv, (35,40,40),(85,255,255)))
    yellow = np.sum(cv2.inRange(hsv,(15,50,50),(35,255,255)))
    white = np.sum(cv2.inRange(hsv,(0,0,200),(180,30,255)))

    if green > yellow and green > white:
        return "EV"
    elif yellow > green and yellow > white:
        return "Commercial"
    else:
        return "Personal"

# -------- OCR --------
def read_plate(img):
    results = reader.readtext(np.array(img))
    if results:
        return results[0][1]
    return "UNKNOWN"

# -------- Dashboard --------
def get_stats():
    df = pd.read_sql("SELECT * FROM vehicles", conn)
    if len(df)==0:
        return "No vehicles yet"
    return df['type'].value_counts().to_string()

# -------- Pipeline --------
def process_image(img):
    image = Image.fromarray(img)
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    results = processor.post_process_object_detection(
        outputs, threshold=0.3,
        target_sizes=torch.tensor([[image.size[1], image.size[0]]])
    )[0]

    draw = ImageDraw.Draw(image)

    if len(results["boxes"])==0:
        return image,"No Plate","",get_stats()

    box = results["boxes"][0].tolist()
    x1,y1,x2,y2 = map(int,box)

    plate = image.crop((x1,y1,x2,y2))
    plate_text = read_plate(plate)
    vtype = classify_plate_color(plate)

    cursor.execute("INSERT INTO vehicles VALUES (?,?,?)",
                   (plate_text,vtype,datetime.now().strftime("%H:%M:%S")))
    conn.commit()

    draw.rectangle([x1,y1,x2,y2],outline="red",width=3)
    draw.text((x1,y1-10),f"{plate_text} | {vtype}",fill="red")

    return image, plate_text, vtype, get_stats()

# -------- UI --------
with gr.Blocks() as demo:
    gr.Markdown("# 🚗 Smart Traffic & EV Detection System")

    img = gr.Image(type="numpy", sources=["upload","webcam"])
    out_img = gr.Image()
    plate = gr.Textbox(label="Number Plate")
    vtype = gr.Textbox(label="Vehicle Type")
    stats = gr.Textbox(label="Dashboard")

    btn = gr.Button("Scan Vehicle")
    btn.click(process_image, img, [out_img,plate,vtype,stats])

demo.launch()