# Author Sarvamangala Kokatanur
# Import libraries
import gradio as gr
import torch
import numpy as np
import cv2
import sqlite3
import pandas as pd
import matplotlib.pyplot as plt

from datetime import datetime, timedelta

from PIL import Image, ImageDraw

from transformers import (
    YolosImageProcessor,
    YolosForObjectDetection
)

# 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,

timestamp TEXT,

license_plate TEXT,

vehicle_status TEXT,

discount INTEGER

)
""")

conn.commit()

# -------- Plate Color Classifier -------- #
def classify_plate_color(plate_img):

    img = np.array(plate_img)

    # Convert RGB → HSV
    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)

    h, w = hsv.shape[:2]

    # Only inspect the left 20% of the plate
    left = hsv[:, :int(w*0.2)]

    # Green mask
    green = cv2.inRange(
        left,
        (35, 40, 40),
        (90, 255, 255)
    )

    green_ratio = np.count_nonzero(green) / green.size

    # If more than 15% of the left strip is green,
    # classify as EV
    if green_ratio > 0.15:
        return "EV"

    return "Non-EV"

# Avoid duplicate vehicle details to the dashboard

def is_duplicate_vehicle(plate_number):

    cursor.execute("""
    SELECT timestamp
    FROM vehicles
    WHERE license_plate=?
    ORDER BY id DESC
    LIMIT 1
    """,(plate_number,))

    row = cursor.fetchone()

    if row is None:
        return False

    last_time = datetime.strptime(
        row[0],
        "%Y-%m-%d %H:%M:%S"
    )

    if datetime.now() - last_time < timedelta(minutes=5):

        return True

    return False

# to save vehicle details in the database

def save_vehicle(plate,status):

    if status=="EV":
        discount=50
    else:
        discount=0

    if is_duplicate_vehicle(plate):
        return "Duplicate"

    current_time=datetime.now().strftime(
        "%Y-%m-%d %H:%M:%S"
    )

    cursor.execute("""

    INSERT INTO vehicles(

    timestamp,

    license_plate,

    vehicle_status,

    discount

    )

    VALUES(?,?,?,?)

    """,

    (

    current_time,

    plate,

    status,

    discount

    ))

    conn.commit()
    return "Saved"

# --------get_dashboard function ------#
def get_dashboard():

    df = pd.read_sql(
        "SELECT * FROM vehicles",
        conn
    )

    fig, axs = plt.subplots(2, 2, figsize=(8, 6))

    if df.empty:

        for ax in axs.flatten():
            ax.text(
                0.5,
                0.5,
                "No Data Available",
                ha="center",
                va="center",
                fontsize=10
            )
            ax.axis("off")

        plt.tight_layout()
        return fig

    status_counts = df["vehicle_status"].value_counts()

    axs[0,0].bar(
        status_counts.index,
        status_counts.values
    )

    axs[0,0].set_title("EV vs Non-EV")
    if status_counts.empty:
        plt.tight_layout()
        return fig
    
    axs[0,1].pie(
        status_counts.values,
        labels=status_counts.index,
        autopct="%1.1f%%"
    )

    axs[0,1].set_title("Vehicle Distribution")

    total_discount = df["discount"].sum()

    axs[1,0].bar(
        ["Discount"],
        [total_discount]
    )

    axs[1,0].set_title("Total Discount")

    axs[1,1].axis("off")

    report = (
        f"Today's Report\n\n"
        f"Total Vehicles : {len(df)}\n\n"
        f"EV : {len(df[df.vehicle_status=='EV'])}\n\n"
        f"Non-EV : {len(df[df.vehicle_status=='Non-EV'])}\n\n"
        f"Discount Given : ₹{total_discount}"
    )

    axs[1,1].text(
        0,
        1,
        report,
        fontsize=10,
        va="top"
    )

    plt.tight_layout()

    plt.close(fig)
    return fig
    
# -------- Main Pipeline -------- #
def process_image(img):

    if img is None:
        return (
            None,
            "Please upload an image.",
            "0",
            "0",
            "₹0",
            get_dashboard()
        )

    image = Image.fromarray(img)

    inputs = processor(images=image, return_tensors="pt")

    with torch.no_grad():
        outputs = model(**inputs)

    target_sizes = torch.tensor([[image.size[1], image.size[0]]])

    results = processor.post_process_object_detection(
        outputs,
        threshold=0.3,
        target_sizes=target_sizes
    )[0]

    draw = ImageDraw.Draw(image)

    ev_count = 0
    non_ev_count = 0
    discount_total = 0

    output_text = ""

    # No detection
    if len(results["boxes"]) == 0:
        return (
            image,
            "No license plate detected.",
            "0",
            "0",
            "₹0",
            get_dashboard()
        )

    # Process each detected plate
    for i, box in enumerate(results["boxes"]):

        x1, y1, x2, y2 = map(int, box.tolist())

        plate = image.crop((x1, y1, x2, y2))

        status = classify_plate_color(plate)

        # Temporary plate number
        # Replace with OCR later
        plate_number = f"Vehicle_{datetime.now().strftime('%H%M%S%f')}_{i}"

        saved = save_vehicle(
            plate_number,
            status
        )

        # Skip duplicate entries
        if saved == "Duplicate":
            continue

        if status == "EV":

            ev_count += 1
            discount = 50
            discount_total += discount

            color = "green"
            label = f"{plate_number}\nEV | ₹{discount}"

        else:

            non_ev_count += 1
            discount = 0

            color = "red"
            label = f"{plate_number}\nNon-EV"

        # Draw bounding box
        draw.rectangle(
            [x1, y1, x2, y2],
            outline=color,
            width=3
        )

        # Draw label
        draw.text(
            (x1, max(0, y1 - 30)),
            label,
            fill=color
        )

        output_text += (
            f"Vehicle {i+1}\n"
            f"Plate : {plate_number}\n"
            f"Status : {status}\n"
            f"Discount : ₹{discount}\n\n"
        )

    conn.commit()

    dashboard = get_dashboard()

    return (
        image,
        output_text,
        str(ev_count),
        str(non_ev_count),
        f"₹{discount_total}",
        dashboard
    )
# -------- Gradio UI -------- #
css = """
.gradio-container{
    max-width:98% !important;
    margin:auto;
}

textarea{
    font-size:15px !important;
}

h1{
    text-align:center;
}

.block{
    border-radius:12px;
}

.gr-image{
    min-height:350px;
}
"""
with gr.Blocks() as demo:

    gr.Markdown("# Smart Traffic & EV Analytics System")
    gr.Markdown(
        "Automatic License Plate Detection • EV Classification • Discount Calculation"
    )

    ##################################################
    # TOP SECTION
    ##################################################

    with gr.Row(equal_height=True):

        # LEFT
        with gr.Column(scale=1):

            input_img = gr.Image(
                type="numpy",
                label="Input Image",
                sources=["upload","webcam"],
                height=350
            )

            btn = gr.Button(
                "Scan Vehicle",
                variant="primary",
                size="lg"
            )

        # CENTER
        with gr.Column(scale=1):

            output_img = gr.Image(
                label="Detection Result",
                height=350
            )

        # RIGHT
        with gr.Column(scale=1):

            summary_box = gr.Textbox(
                label="Detection Summary",
                lines=16
            )

    ##################################################
    # STATISTICS
    ##################################################

    with gr.Row():

        ev_box = gr.Textbox(
            label="EV Vehicles"
        )

        non_ev_box = gr.Textbox(
            label="Non-EV Vehicles"
        )

        discount_box = gr.Textbox(
            label="Total Discount"
        )

    ##################################################
    # DASHBOARD
    ##################################################

    gr.Markdown("## 📊 Today's Traffic Dashboard")

    dashboard = gr.Plot(label="Analytics")
    

    btn.click(
        fn=process_image,
        inputs=input_img,
        outputs=[
            output_img,
            summary_box,
            ev_box,
            non_ev_box,
            discount_box,
            dashboard
        ]
    )
if __name__ == "__main__":
    demo.queue()
    demo.launch(
        ssr_mode=False
    )