import gradio as gr
from risk_model import predict_risk, get_history_df
import pandas as pd
import matplotlib.pyplot as plt

with gr.Blocks() as demo:
    gr.Markdown("## 🔥 Heating Mantle Safety Risk Predictor")

    # 🔄 Reset button at the top
    with gr.Row():
        reset_btn = gr.Button("🔄 Reset")

    # Input fields
    with gr.Row():
        temp = gr.Number(label="Max Temperature (°C)", value=100)
        duration = gr.Number(label="Duration (min)", value=30)

    # Predict button
    with gr.Row():
        predict_btn = gr.Button("🔍 Predict")

    # Output fields
    result = gr.Textbox(label="Risk Prediction")
    alert = gr.Textbox(label="🚨 Alert Message")
    ist_time = gr.Textbox(label="Timestamp (IST)")
    risk_score = gr.Textbox(label="📊 Risk Score")  # NEW FIELD
    summary = gr.Markdown()
    history_table = gr.Dataframe(headers=["Temperature", "Duration", "Risk", "Timestamp"], label="📈 Prediction History")
    plot = gr.Plot(label="📊 Risk Trend Chart")

    def classify(temp, duration):
        if temp <= 0 or duration <= 0:
            return "❌ Invalid", "Invalid", "Invalid", "", "", pd.DataFrame(), plt.figure()

        risk, timestamp = predict_risk(temp, duration)
        score = round((temp + duration) / 2, 2)

        alert_msg = {
            "Low": "✅ SAFE - No action needed",
            "Moderate": "⚠️ SAFE - Monitor closely",
            "High": "🔥 SHUTDOWN - Immediate attention needed"
        }.get(risk, "Unknown")

        summary_md = f"""
### 🔎 Summary
- **Max Temp**: {temp} °C  
- **Duration**: {duration} min  
- **Risk Score**: {score}  
- **Risk**: {risk}  
- **Timestamp**: {timestamp}  
- **Alert**: {alert_msg}
"""

        df = get_history_df()
        risk_map = {'Low': 1, 'Moderate': 2, 'High': 3}
        df["Risk_Num"] = df["Risk"].map(risk_map)

        fig, ax = plt.subplots(figsize=(6, 3))
        ax.plot(df["Timestamp"], df["Risk_Num"], marker="o", linestyle="-", color="red")
        ax.set_ylim(0.5, 3.5)
        ax.set_yticks([1, 2, 3])
        ax.set_yticklabels(['Low', 'Moderate', 'High'])
        ax.set_title("Risk Level Over Time")
        ax.set_xlabel("Timestamp")
        ax.set_ylabel("Risk Level")
        ax.tick_params(axis='x', rotation=45)
        plt.tight_layout()

        df_display = df[["Temperature", "Duration", "Risk", "Timestamp"]]

        return risk, alert_msg, timestamp, score, summary_md, df_display, fig

    # Function to reset all fields
    def reset_all():
        return (
            100, 30, "", "", "", "", "", pd.DataFrame(columns=["Temperature", "Duration", "Risk", "Timestamp"]), plt.figure()
        )

    predict_btn.click(
        classify,
        inputs=[temp, duration],
        outputs=[result, alert, ist_time, risk_score, summary, history_table, plot]
    )

    reset_btn.click(
        reset_all,
        inputs=[],
        outputs=[temp, duration, result, alert, ist_time, risk_score, summary, history_table, plot]
    )

demo.launch()