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

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

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

    with gr.Row():
        predict_btn = gr.Button("🔍 Predict")
        retrain_btn = gr.Button("🔁 Retrain Model")

    result = gr.Textbox(label="Risk Prediction")
    score = gr.Textbox(label="Confidence (%)")
    alert = gr.Textbox(label="🚨 Alert Message")
    ist_time = gr.Textbox(label="Timestamp (IST)")
    retrain_output = gr.Textbox(label="Retrain Status")

    summary = gr.Markdown()
    history_table = gr.Dataframe(headers=["Temperature", "Duration", "Risk", "Confidence", "Timestamp"], label="📈 Prediction History")

    plot = gr.Plot(label="📊 Risk Distribution Chart")

    def classify(temp, duration):
        if temp <= 0 or duration <= 0:
            return "❌ Invalid Input", "Use values > 0", "N/A", "N/A", "", pd.DataFrame(), plt.figure()

        risk, confidence, timestamp = predict_risk(temp, duration)

        if risk == "Low":
            alert_msg = "✅ SAFE - No action needed"
        elif risk == "Moderate":
            alert_msg = "⚠️ SAFE - Monitor closely"
        else:
            alert_msg = "🔥 SHUTDOWN - Immediate attention needed"

        summary_md = f"""
### 🔎 Summary
- **Max Temp**: {temp} °C  
- **Duration**: {duration} min  
- **Risk**: {risk}  
- **Confidence**: {confidence:.2f}%  
- **Timestamp**: {timestamp}  
- **Alert**: {alert_msg}
"""

        df = get_history_df()

        # Plot chart
        fig, ax = plt.subplots()
        risk_counts = df["Risk"].value_counts()
        risk_counts.plot(kind="bar", ax=ax, color=["green", "orange", "red"])
        ax.set_title("Risk Level Distribution")
        ax.set_ylabel("Count")

        return f"{risk}", f"{confidence:.2f}", alert_msg, timestamp, summary_md, df, fig

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

    retrain_btn.click(retrain_model, outputs=[retrain_output])

demo.launch()