import gradio as gr
import cv2
import numpy as np
import easyocr

reader = easyocr.Reader(['en'], gpu=False)

feedback_data = []

#########################################################
# 1️⃣ Plate Colour Based Indian Vehicle Classification
#########################################################

def classify_vehicle_by_plate_color(plate_img):

    hsv = cv2.cvtColor(plate_img, cv2.COLOR_BGR2HSV)

    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", True
    elif yellow > green and yellow > white:
        return "Commercial", False
    else:
        return "Personal", False


#########################################################
# 2️⃣ Detection + OCR + EV Benefits
#########################################################

def detect_vehicles(image):

    if image is None:
        return None, "Upload image first."

    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    edges = cv2.Canny(gray, 100, 200)
    contours, _ = cv2.findContours(edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    detected_summary = []
    count = 0

    for cnt in contours:

        x, y, w, h = cv2.boundingRect(cnt)

        # plate shape filtering
        if h == 0:
            continue

        ratio = w / h

        if 2 < ratio < 6 and w > 120 and h > 30:

            plate_img = img[y:y+h, x:x+w]

            # OCR
            results = reader.readtext(plate_img)

            plate_number = "Unknown"

            if len(results) > 0:
                plate_number = results[0][1]

            vehicle_type, is_ev = classify_vehicle_by_plate_color(plate_img)

            # EV Benefits
            if is_ev:
                benefit = "EV Benefits: Toll Discount + Parking Discount"
            else:
                benefit = "No EV Benefits"

            # draw detection
            cv2.rectangle(img, (x,y), (x+w,y+h), (0,255,0), 2)

            label = f"{plate_number} | {vehicle_type}"

            cv2.putText(
                img,
                label,
                (x, y-10),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.6,
                (0,255,0),
                2
            )

            detected_summary.append(
                f"Plate: {plate_number} | Type: {vehicle_type} | {benefit}"
            )

            count += 1

    if count == 0:
        summary = "No number plate detected."
    else:
        summary = "\n".join(detected_summary)

    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    return img, summary


#########################################################
# 3️⃣ Feedback + Evaluation
#########################################################

def submit_feedback(is_correct):

    feedback_data.append(is_correct)

    total = len(feedback_data)

    correct = sum(feedback_data)

    accuracy = (correct / total) * 100

    return f"""
Evaluation Summary
-------------------
Total Samples : {total}
Correct       : {correct}
Accuracy      : {accuracy:.2f} %
"""


#########################################################
# 4️⃣ Gradio UI
#########################################################

with gr.Blocks() as demo:

    gr.Markdown("## 🚦 Smart Traffic & EV Classification System")

    with gr.Row():

        with gr.Column(scale=2):

            image_input = gr.Image(type="pil", label="Upload Image")

            detect_btn = gr.Button("Detect", size="sm")

            output_image = gr.Image(label="Output")

            output_text = gr.Textbox(label="Detection Summary")

        with gr.Column(scale=1):

            gr.Markdown("### Feedback")

            correct_btn = gr.Button("Correct", size="sm")

            incorrect_btn = gr.Button("Incorrect", size="sm")

            summary_box = gr.Textbox(label="Evaluation Summary")

    detect_btn.click(
        fn=detect_vehicles,
        inputs=image_input,
        outputs=[output_image, output_text]
    )

    correct_btn.click(
        fn=lambda: submit_feedback(True),
        outputs=summary_box
    )

    incorrect_btn.click(
        fn=lambda: submit_feedback(False),
        outputs=summary_box
    )

demo.launch()