import gradio as gr
import torch
import numpy as np
import cv2
import sqlite3
import pandas as pd
import matplotlib.pyplot as plt
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'], gpu=False)

# -------------------- Plate Color Classifier --------------------
def classify_plate_color(plate_img):
    img = np.array(plate_img)
    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(plate_img):
    results = reader.readtext(np.array(plate_img))
    if results:
        return results[0][1]
    return "UNKNOWN"

# -------------------- Dashboard --------------------
def get_dashboard():
    df = pd.read_sql("SELECT * FROM vehicles", conn)

    fig, ax = plt.subplots(figsize=(8, 5))

    if len(df) == 0:
        ax.text(0.5, 0.5, "No vehicles scanned yet",
                ha="center", va="center", fontsize=10)
        ax.axis("off")
        return fig

    counts = df["type"].value_counts()
    counts.plot(kind="bar", ax=ax)

    ax.set_title("Vehicle Classification Dashboard")
    ax.set_xlabel("Vehicle Type")
    ax.set_ylabel("Count")
    ax.grid(axis="y")

    return fig

# -------------------- Main 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 detected", "", get_dashboard()

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

    plate_text = read_plate(plate_img)
    vehicle_type = classify_plate_color(plate_img)

    cursor.execute(
        "INSERT INTO vehicles VALUES (?, ?, ?)",
        (plate_text, vehicle_type, datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
    )
    conn.commit()

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

    return image, plate_text, vehicle_type, get_dashboard()

# -------------------- Gradio UI --------------------
css = """
textarea {
    white-space: pre-wrap !important;
    word-break: break-word !important;
    overflow-wrap: break-word !important;
}
"""

with gr.Blocks() as demo:
    gr.Markdown("## Smart Traffic & EV Analytics System")
    gr.Markdown(
        "Detects license plates, reads number plate text, "
        "classifies EV / Commercial / Personal vehicles, "
        "and shows live analytics."
    )

    with gr.Row():
        input_img = gr.Image(type="numpy", sources=["upload", "webcam"])
        output_img = gr.Image()

    with gr.Row():
        plate_box = gr.Textbox(label="Number Plate")
        type_box = gr.Textbox(label="Vehicle Type")

    dashboard = gr.Plot(label="Live Vehicle Dashboard")

    btn = gr.Button("Scan Vehicle")
    btn.click(
        process_image,
        input_img,
        [output_img, plate_box, type_box, dashboard]
    )

demo.launch(css=css)