import gradio as gr
from ultralytics import YOLO

MODEL_PATHS = {
    "Baseline - YOLOv8n original dataset": "baseline.pt",
    "Experiment 1 - YOLOv8m original dataset": "experiment1.pt",
    "Experiment 2 - YOLOv8n expanded dataset": "experiment2.pt",
    "Final Model - YOLOv8m expanded dataset": "final_version.pt",
}

MODEL_INFO = {
    "Baseline - YOLOv8n original dataset": "Original baseline model trained with YOLOv8n on the original dataset.",
    "Experiment 1 - YOLOv8m original dataset": "Larger YOLOv8m model trained on the original dataset to test architecture improvement.",
    "Experiment 2 - YOLOv8n expanded dataset": "YOLOv8n model trained using the expanded dataset to test dataset enhancement.",
    "Final Model - YOLOv8m expanded dataset": "Final best model using YOLOv8m with the expanded dataset.",
}

loaded_models = {}

def get_model(model_name):
    if model_name not in loaded_models:
        loaded_models[model_name] = YOLO(MODEL_PATHS[model_name])
    return loaded_models[model_name]

def update_model_info(model_name):
    return MODEL_INFO.get(model_name, "")

def detect_objects(model_name, image, confidence):
    if image is None:
        raise gr.Error("Please upload an image first.")

    model = get_model(model_name)
    results = model.predict(image, conf=confidence)
    annotated_image = results[0].plot()

    return annotated_image

custom_css = """
#title {
    text-align: center;
    margin-bottom: 8px;
}

#subtitle {
    text-align: center;
    color: #555;
    font-size: 16px;
    margin-bottom: 24px;
}

.model-box {
    background: linear-gradient(135deg, #f8fafc, #eef2ff);
    border: 1px solid #dbe4ff;
    border-radius: 16px;
    padding: 16px;
    margin-bottom: 12px;
}

.footer {
    text-align: center;
    color: #666;
    font-size: 13px;
    margin-top: 24px;
}

.gradio-container {
    max-width: 1100px !important;
    margin: auto !important;
}

button.primary {
    border-radius: 12px !important;
    font-weight: 700 !important;
}
"""

with gr.Blocks() as demo:
    gr.Markdown(
        """
        <div id="title">
            <h1>AI-Based Vehicle Detection System</h1>
        </div>
        <div id="subtitle">
            Upload a image, choose a trained YOLO model version, and view the vehicle detection result.
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Model Settings")

            model_dropdown = gr.Dropdown(
                choices=list(MODEL_PATHS.keys()),
                value="Final Model - YOLOv8m expanded dataset",
                label="Choose model version"
            )

            model_description = gr.Textbox(
                value=MODEL_INFO["Final Model - YOLOv8m expanded dataset"],
                label="Selected model description",
                interactive=False,
                lines=3
            )

            confidence_slider = gr.Slider(
                minimum=0.05,
                maximum=0.95,
                value=0.25,
                step=0.05,
                label="Confidence threshold",
                info="Lower values show more detections. Higher values show only more confident detections."
            )

            gr.Markdown(
                """
                <div class="model-box">
                <b>Model Versions</b><br>
                1. Baseline: YOLOv8n + original dataset<br>
                2. Experiment 1: YOLOv8m + original dataset<br>
                3. Experiment 2: YOLOv8n + expanded dataset<br>
                4. Final Model: YOLOv8m + expanded dataset
                </div>
                """
            )

        with gr.Column(scale=2):
            gr.Markdown("### Detection Demo")

            with gr.Row():
                input_image = gr.Image(
                    type="pil",
                    label="Upload image",
                    height=420
                )

                output_image = gr.Image(
                    type="numpy",
                    label="Detection result",
                    height=420
                )

            with gr.Row():
                clear_btn = gr.ClearButton(
                    components=[input_image, output_image],
                    value="Clear"
                )

                submit_btn = gr.Button(
                    value="Run Detection",
                    variant="primary"
                )

    gr.Markdown(
        """
        <div class="footer">
        BCS407 Artificial Intelligence Project · Smart Vehicle Detection · YOLOv8
        </div>
        """
    )

    model_dropdown.change(
        fn=update_model_info,
        inputs=model_dropdown,
        outputs=model_description
    )

    submit_btn.click(
        fn=detect_objects,
        inputs=[model_dropdown, input_image, confidence_slider],
        outputs=output_image
    )

demo.launch(
    theme=gr.themes.Soft(),
    css=custom_css,
    ssr_mode=False
)