import gradio as gr
import pandas as pd
import joblib
import matplotlib.pyplot as plt
import seaborn as sns

# Load pre-trained models and scalers
scaler_initial = joblib.load("scaler_initial.pkl")
scaler_with_cluster = joblib.load("scaler_with_cluster.pkl")
kmeans = joblib.load("kmeans.pkl")
linear_model = joblib.load("linear_model.pkl")
poly_model = joblib.load("poly_model.pkl")
ridge_model = joblib.load("ridge_model.pkl")
rf_regressor = joblib.load("rf_regressor.pkl")
logistic_model = joblib.load("logistic_model.pkl")
rf_classifier = joblib.load("rf_classifier.pkl")


# Prediction function
def predict_aqi(pm25, pm10, no2, co, temp, humidity):
    # Create input dataframe with initial features
    input_data = pd.DataFrame([[pm25, pm10, no2, co, temp, humidity]],
                              columns=["PM2.5", "PM10", "NO2", "CO", "Temperature", "Humidity"])

    # Scale initial features for clustering
    input_scaled_initial = scaler_initial.transform(input_data)

    # Apply K-means clustering
    cluster = kmeans.predict(input_scaled_initial)[0]
    input_data['Cluster'] = cluster

    # Scale data with Cluster feature
    input_scaled_with_cluster = scaler_with_cluster.transform(input_data)

    # Regression predictions
    linear_pred = linear_model.predict(input_scaled_with_cluster)[0]
    poly_pred = poly_model.predict(input_scaled_with_cluster)[0]
    ridge_pred = ridge_model.predict(input_scaled_with_cluster)[0]
    rf_pred = rf_regressor.predict(input_scaled_with_cluster)[0]

    # Classification predictions
    logistic_class = logistic_model.predict(input_scaled_with_cluster)[0]
    rf_class = rf_classifier.predict(input_scaled_with_cluster)[0]

    # Create performance plot
    models = ["Linear", "Polynomial", "Ridge", "Random Forest"]
    predictions = [linear_pred, poly_pred, ridge_pred, rf_pred]
    plt.figure(figsize=(8, 4))
    sns.barplot(x=models, y=predictions)
    plt.title("AQI Predictions by Model")
    plt.ylabel("Predicted AQI")
    plt.savefig("aqi_plot.png")
    plt.close()

    output_text = (
        f"Linear Regression AQI: {linear_pred:.2f}\n"
        f"Polynomial Regression AQI: {poly_pred:.2f}\n"
        f"Ridge Regression AQI: {ridge_pred:.2f}\n"
        f"Random Forest AQI: {rf_pred:.2f}\n"
        f"Logistic Classification: {'Safe' if logistic_class == 0 else 'Unsafe'}\n"
        f"Random Forest Classification: {'Safe' if rf_class == 0 else 'Unsafe'}"
    )
    return output_text, "aqi_plot.png"


# Gradio interface
iface = gr.Interface(
    fn=predict_aqi,
    inputs=[
        gr.Slider(0, 200, label="PM2.5 (µg/m³)", value=50),
        gr.Slider(0, 300, label="PM10 (µg/m³)", value=80),
        gr.Slider(0, 100, label="NO2 (µg/m³)", value=20),
        gr.Slider(0, 10, label="CO (mg/m³)", value=1),
        gr.Slider(-10, 40, label="Temperature (°C)", value=20),
        gr.Slider(0, 100, label="Humidity (%)", value=50)
    ],
    outputs=[
        gr.Textbox(label="Predictions"),
        gr.Image(label="Model Comparison Plot")
    ],
    title="Air Quality Prediction and Classification",
    description="Enter pollutant levels and weather conditions to predict AQI and classify air quality. Built with multiple machine learning models to address urban air pollution."
)

if __name__ == "__main__":
    iface.launch()