import gradio as gr
from risk_model import predict_risk, retrain_model, get_history_df
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")
    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", "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)

        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}  
- **Timestamp**: {timestamp}  
- **Alert**: {alert_msg}
"""

        df = get_history_df()

        # Convert Risk to numeric for plotting
        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")
        a