import cv2
import numpy as np
import gradio as gr
from ultralytics import YOLO
import json
import random
import requests

def request_fire_data():
    url = "https://kauil-fire-estimator-208352106463.us-central1.run.app/predict_new_fire"

    # Define the data payload (JSON format)
    bodies = [
        {"longitude": "-115.077305804044", "latitude": "56.7542346185555", "fire_start_date": "2024-01-09"},
        {"longitude": "-110.951883007294", "latitude": "56.0765930131132", "fire_start_date": "2024-06-07"},
        {"longitude": "-133.899560229577", "latitude": "63.4377263695738", "fire_start_date": "2024-09-22"}
    ]

    # Randomly select a body
    selected_body = random.choice(bodies)

    
    headers = {
        "Content-Type": "application/json",
    }
    
    # Send the POST request
    response = requests.post(url, json=selected_body, headers=headers)

    return response.json()

def process_image_and_risk(image_path: str):
    # YOLO prediction
    model = YOLO("yolo11m.yaml")
    model = YOLO("kauil_smoke_detection.pt")
    results = model.predict(source=image_path)
    
    # Simulating API response (replace this with your actual API call)
    api_response = request_fire_data()
    
    # Process the image
    annotated_image = None
    for r in results:
        annotated_image = r.plot()
    
    # Create formatted risk level display
    risk_level = api_response["predicted_risk_level"]
    confidence = api_response["confidence"]
    
    # Define color schemes for different risk levels
    risk_colors = {
        "Low": "#2ECC71",    # Green
        "Medium": "#F1C40F", # Yellow
        "High": "#E74C3C"    # Red
    }
    
    # Create HTML for styled output
    color = risk_colors.get(risk_level, "#7F8C8D")  # Default gray if unknown level
    html_output = f"""
    <div style="padding: 20px; border-radius: 10px; background-color: {color}; color: white;">
        <h2 style="margin: 0; font-size: 24px;">Risk Level: {risk_level}</h2>
        <p style="margin: 10px 0 0 0; font-size: 18px;">Confidence: {confidence:.2%}</p>
    </div>
    """
    
    return annotated_image, html_output

# Create Gradio interface with custom layout
with gr.Blocks() as app:
    gr.Markdown("# Smoke Detection with Risk Assessment")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="filepath", label="Input Image")
        with gr.Column():
            output_image = gr.Image(type="numpy", label="Detection Result")
    
    with gr.Row():
        risk_display = gr.HTML(label="Risk Assessment")
    
    input_image.change(
        fn=process_image_and_risk,
        inputs=[input_image],
        outputs=[output_image, risk_display]
    )

app.launch(debug=False)