import gradio as gr
from PIL import Image
import os

# Import modules
from detect_anomaly import detect_hazard
from generate_hazards import generate_unsafe_image
from explain_alert import explain_hazard
from detect_helmet import detect_helmet_in_image

last_result = {"anomalous": False, "error": 0.0}

def analyze_image(img):
    global last_result
    if img is None:
        return "Please upload an image", None, None

    pil_img = Image.fromarray(img).convert("RGB")
    is_anomalous, error, recon = detect_hazard(pil_img)
    last_result = {"anomalous": is_anomalous, "error": error}

    alert = "HAZARD DETECTED: No helmet or danger zone!" if is_anomalous else "No hazard detected."

    explanation = explain_hazard(
        helmet="No" if is_anomalous else "Yes",
        zone="Danger" if is_anomalous else "Safe",
        lighting="Normal"
    )

    return alert, recon, explanation

def create_synthetic():
    img = generate_unsafe_image()
    return img

#Gradio Interface
with gr.Blocks(title="AI Safety Inspector") as demo:
    gr.Markdown("# AI Safety Inspector\nDetects missing helmets, danger zones using **Unsupervised + Gen AI**")

    with gr.Tabs():
        with gr.Tab("Analyze Image"):
            with gr.Row():
                input_img = gr.Image(label="Upload Site Photo")
                output_recon = gr.Image(label="Reconstructed (Autoencoder)")
            output_alert = gr.Label(label="Status")
            output_explain = gr.Textbox(label="AI Safety Officer Says")
            btn = gr.Button("Analyze")
            btn.click(analyze_image, inputs=input_img, outputs=[output_alert, output_recon, output_explain])

        with gr.Tab("Helmet Detection (YOLOv8)"):
            gr.Markdown("Detects workers and checks if they are wearing helmets using YOLOv8.")

            with gr.Row():
                yolo_input = gr.Image(label="Upload Image")
                yolo_output_img = gr.Image(label="Detected Helmets")

            yolo_output_count = gr.Number(label="Workers Without Helmet")
            yolo_output_labels = gr.Textbox(label="Detections")

            yolo_btn = gr.Button("Run Helmet Detection")
            yolo_btn.click(
                detect_helmet_in_image,
                inputs=yolo_input,
                outputs=[yolo_output_img, yolo_output_count, yolo_output_labels]
            )

        with gr.Tab("Generate Synthetic Hazard (Gen AI)"):
            gen_output = gr.Image(label="Generated Unsafe Scenario")
            gen_btn = gr.Button("Generate No-Helmet Danger Scene")
            gen_btn.click(create_synthetic, outputs=gen_output)

        with gr.Tab("ℹ About"):
            gr.Markdown("""
            ### How It Works
            - Uses **autoencoder** to detect anomalies (no labels needed!)
            - **YOLOv8** detects helmets with bounding boxes
            - **Stable Diffusion** generates synthetic unsafe images
            - **TinyLlama** explains alerts in natural language
            - Runs on **Google Colab**
            """)

#Launch
demo.launch(share=True)