import gradio as gr
import joblib
import numpy as np
import matplotlib.pyplot as plt

# Load trained model and scaler
model = joblib.load("aqi_model.pkl")
scaler = joblib.load("scaler.pkl")

# Function to create AQI gauge chart
def create_aqi_gauge(aqi):
    fig, ax = plt.subplots(figsize=(5, 2))

    # Define AQI categories and their color ranges
    categories = [
        (0, 50, "green"),
        (50, 100, "yellow"),
        (100, 150, "orange"),
        (150, 200, "red"),
        (200, 300, "purple"),
        (300, 500, "maroon")
    ]
    
    # Draw color bands
    for start, end, color in categories:
        ax.barh(0, end-start, left=start, color=color, height=5, align="center")

    # Plot AQI marker
    ax.scatter(aqi, 0, color="black", s=200, marker="o", label=f"AQI: {aqi}")
    
    # Formatting
    ax.set_xlim(0, 500)
    ax.set_yticks([])
    ax.set_xticks([0, 50, 100, 150, 200, 300, 500])
    ax.set_xticklabels(["0", "50", "100", "150", "200", "300", "500"])
    ax.set_title("AQI Level Indicator")

    return fig

# Prediction function
def predict_aqi(pm25, pm10, co, no2, so2, o3):
    input_data = np.array([[pm25, pm10, co, no2, so2, o3]])
    input_scaled = scaler.transform(input_data)  # Normalize input
    prediction = model.predict(input_scaled)
    aqi = round(prediction[0], 2)
    
    # Determine AQI category and suggestions
    if aqi <= 50:
        category = "Good"
        suggestion = "Air quality is satisfactory. Enjoy your day!"
    elif aqi <= 100:
        category = "Moderate"
        suggestion = "Air quality is acceptable. Sensitive groups should take care."
    elif aqi <= 150:
        category = "Unhealthy for Sensitive Groups"
        suggestion = "Sensitive groups should limit outdoor activities."
    elif aqi <= 200:
        category = "Unhealthy"
        suggestion = "Everyone should limit prolonged outdoor exertion."
    elif aqi <= 300:
        category = "Very Unhealthy"
        suggestion = "Health alert: Avoid outdoor activities."
    else:
        category = "Hazardous"
        suggestion = "Health warning: Stay indoors and avoid physical activities."
    
    # Generate AQI gauge chart
    gauge_plot = create_aqi_gauge(aqi)
    
    return f"AQI: {aqi} ({category})", suggestion, gauge_plot

# Function to clear input fields
def clear_inputs():
    return None, None, None, None, None, None, "", "", None  # Reset all fields

# Create Gradio UI
with gr.Blocks() as iface:
    gr.Markdown("<h1 style='text-align: center;'>🌎 AQI Predictor</h1>")
    gr.Markdown("<h3 style='text-align: center;'>Enter pollutant levels to predict the Air Quality Index (AQI).</h3>")

    # with gr.Row():
    #     pm25 = gr.Number(label="PM2.5 (µg/m³)", info="Typical: 0 - 500", value=None)
    #     pm10 = gr.Number(label="PM10 (µg/m³)", info="Typical: 0 - 600", value=None)
    #     co = gr.Number(label="CO (mg/m³)", info="Typical: 0 - 50", value=None)
    #     no2 = gr.Number(label="NO2 (µg/m³)", info="Typical: 0 - 200", value=None)
    #     so2 = gr.Number(label="SO2 (µg/m³)", info="Typical: 0 - 300", value=None)
    #     o3 = gr.Number(label="O3 (µg/m³)", info="Typical: 0 - 400", value=None)
    


    with gr.Row():
        pm25 = gr.Number(label="PM2.5 (µg/m³)", info="Examples: 12.5, 85.9", value=23.7)
        pm10 = gr.Number(label="PM10 (µg/m³)", info="Examples: 35.8, 280.5", value=78.3)
        co = gr.Number(label="CO (mg/m³)", info="Examples: 2.1, 11.2", value=4.9)
        no2 = gr.Number(label="NO2 (µg/m³)", info="Examples: 28.7, 95.0", value=42.1)
        so2 = gr.Number(label="SO2 (µg/m³)", info="Examples: 15.3, 70.6", value=25.8)
        o3 = gr.Number(label="O3 (µg/m³)", info="Examples: 65.0, 180.2", value=90.5)


    with gr.Row():
        predict_button = gr.Button("Predict AQI")
        clear_button = gr.Button("Clear")  # Added Clear button

    # Two-column layout for AQI Result & Suggestion vs. AQI Indicator
    with gr.Row():
        with gr.Column():
            aqi_output = gr.Textbox(label="AQI Result", interactive=False)
            suggestion_output = gr.Textbox(label="Suggestion", interactive=False)

        aqi_gauge = gr.Plot(label="AQI Indicator")  # AQI Indicator (Gauge Chart)

    # Button Click Actions
    predict_button.click(
        predict_aqi,
        inputs=[pm25, pm10, co, no2, so2, o3],
        outputs=[aqi_output, suggestion_output, aqi_gauge]
    )

    clear_button.click(
        clear_inputs,
        inputs=[],
        outputs=[pm25, pm10, co, no2, so2, o3, aqi_output, suggestion_output, aqi_gauge]
    )

    # # Example cases
    # gr.Markdown("<h3 style='text-align: center;'>🔥 Try These Examples:</h3>")
    # examples = gr.Examples(
    #     examples=[
    #         [10, 20, 0.3, 15, 5, 30],  # Low AQI example
    #         [180, 200, 2.5, 100, 40, 80]  # High AQI example
    #     ],
    #     inputs=[pm25, pm10, co, no2, so2, o3],
    #     outputs=[aqi_output, suggestion_output, aqi_gauge],
    #     fn=predict_aqi,
    #     cache_examples=True  # Speeds up example execution
    # )


    # Assuming pm25, pm10, co, no2, so2, o3, aqi_output, suggestion_output, aqi_gauge, predict_aqi are defined elsewhere

    # Example cases
    gr.Markdown("<h3 style='text-align: center;'>🔥 Try These Examples:</h3>")
    examples = gr.Examples(
        examples=[
            [12.5, 35.8, 2.1, 28.7, 15.3, 65.0],  # Moderate AQI example
            [85.9, 280.5, 11.2, 95.0, 70.6, 180.2], # High AQI example
        ],
        inputs=[pm25, pm10, co, no2, so2, o3],
        outputs=[aqi_output, suggestion_output, aqi_gauge],
        fn=predict_aqi,
        cache_examples=True  # Speeds up example execution
    )


# Run Gradio app
if __name__ == "__main__":
    iface.launch()