import gradio as gr
import pandas as pd
import joblib
import matplotlib.pyplot as plt
import numpy as np

# ----------------------------
# Load Model
# ----------------------------
model = joblib.load("engine_condition_rf_production.joblib")
saved_threshold = joblib.load("decision_threshold.joblib")

feature_names = model.feature_names_in_

# ----------------------------
# Single Prediction Function
# ----------------------------
def predict_engine(*inputs):
    input_df = pd.DataFrame([inputs], columns=feature_names)
    probability = model.predict_proba(input_df)[0][1]
    prediction = 1 if probability >= saved_threshold else 0
    
    if prediction == 1:
        result = "⚠ Engine Likely Faulty"
    else:
        result = "✅ Engine Operating Normally"
        
    return result, round(probability, 4)

# ----------------------------
# Batch Prediction Function
# ----------------------------
def batch_predict(file):
    df = pd.read_csv(file.name)
    
    missing_cols = [col for col in feature_names if col not in df.columns]
    
    if missing_cols:
        return f"Missing required columns: {missing_cols}"
    
    df = df[feature_names]
    probabilities = model.predict_proba(df)[:, 1]
    
    df["Probability_of_Failure"] = probabilities
    df["Prediction"] = (probabilities >= saved_threshold).astype(int)
    
    output_file = "engine_predictions.csv"
    df.to_csv(output_file, index=False)
    
    return output_file

# ----------------------------
# Build UI
# ----------------------------
with gr.Blocks() as demo:
    
    gr.Markdown("# 🚗 Engine Condition Classification System")
    
    gr.Markdown("## 🔧 Manual Prediction")
    
    inputs = []
    for feature in feature_names:
        inputs.append(gr.Number(label=feature))
    
    output_text = gr.Textbox(label="Prediction Result")
    output_prob = gr.Number(label="Failure Probability")
    
    btn = gr.Button("Predict Engine Condition")
    btn.click(predict_engine, inputs, [output_text, output_prob])
    
    gr.Markdown("## 📂 Batch Prediction (CSV Upload)")
    
    file_input = gr.File(label="Upload CSV File")
    file_output = gr.File(label="Download Predictions")
    
    file_input.change(batch_predict, file_input, file_output)

demo.launch()