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

# Dummy models for demonstration (replace with actual ones)
class DummyLSTM:
    def predict(self, x):
        # Simulate next-day pollution by averaging inputs and adding noise
        return np.mean(x, axis=1) + np.random.rand(1, 5) * 5

class DummyRF:
    def predict(self, x):
        # Health score based on average pollution
        avg = np.mean(x)
        return [100 - min(avg, 100)]

# Load or simulate models
lstm_model = DummyLSTM()
rf_model = DummyRF()
pollutants = ['Ozone', 'SO2', 'NO2', 'PM10', 'PM2.5']

def predict_pollution(*inputs):
    try:
        input_data = np.array([inputs[i:i+5] for i in range(0, 15, 5)], dtype=np.float32)  # shape (3, 5)
        reshaped = input_data.reshape(1, 3, 5)  # batch size 1

        pred = lstm_model.predict(reshaped)
        pred = np.clip(pred, 0, 500)  # limit extreme values

        health_score = rf_model.predict(pred)[0]

        # Visualization
        plt.figure(figsize=(8, 4))
        sns.barplot(x=pollutants, y=pred[0])
        plt.title("Predicted Pollutant Levels")
        plt.ylabel("µg/m³")
        plot_path = "plot.png"
        plt.savefig(plot_path)
        plt.close()

        # Health risk evaluation
        if health_score > 80:
            risk = "🔴 High Risk"
            color = "#ff6b6b"
            advice = "Avoid outdoor activities."
        elif health_score > 60:
            risk = "🟠 Moderate Risk"
            color = "#ffd166"
            advice = "Limit outdoor exposure."
        else:
            risk = "🟢 Low Risk"
            color = "#06d6a0"
            advice = "Air quality is acceptable."

        health_html = f"""
        <div style="background:{color}20; padding:15px; border-radius:8px; border-left:6px solid {color};">
            <h3>{risk}</h3>
            <p><strong>Health Score:</strong> {health_score:.1f}/100</p>
            <p>{advice}</p>
        </div>
        """

        results = [[p, f"{val:.1f} µg/m³"] for p, val in zip(pollutants, pred[0])]
        return plot_path, health_html, results

    except Exception as e:
        return None, f"<div style='color:red;'>Error: {str(e)}</div>", []

# Gradio interface
with gr.Blocks(title="Air Quality Predictor", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🌍 Air Quality Prediction Tool")
    gr.Markdown("Enter pollutant data for the past 3 days. The app will predict pollution for the next day and estimate health impact.")

    with gr.Row():
        with gr.Column():
            inputs = []
            for i in range(3):
                gr.Markdown(f"### Day {i+1}")
                inputs.extend([
                    gr.Number(label="Ozone", value=30 + i * 5),
                    gr.Number(label="SO2", value=20 + i * 2),
                    gr.Number(label="NO2", value=35 + i),
                    gr.Number(label="PM10", value=120 + i * 5),
                    gr.Number(label="PM2.5", value=90 - i * 3),
                ])
            btn = gr.Button("Predict", variant="primary")

        with gr.Column():
            with gr.Tabs():
                with gr.Tab("Chart"):
                    plot = gr.Image()
                with gr.Tab("Health Impact"):
                    health = gr.HTML()
                with gr.Tab("Details"):
                    table = gr.DataFrame(headers=["Pollutant", "Predicted µg/m³"], datatype=["str", "str"], interactive=False)

    btn.click(predict_pollution, inputs=inputs, outputs=[plot, health, table])

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