import gradio as gr
import joblib
import numpy as np
import pandas as pd
from huggingface_hub import hf_hub_download

# Download model & scaler from Hugging Face Hub
model_path = hf_hub_download(repo_id="thalaivanan/pm-decision-tree", filename="model.joblib")
scaler_path = hf_hub_download(repo_id="thalaivanan/pm-decision-tree", filename="scaler.joblib")

model = joblib.load(model_path)
scaler = joblib.load(scaler_path)

feature_names = ["engine_rpm", "lub_oil_pressure", "fuel_pressure", "coolant_pressure", "lub_oil_temp", "coolant_temp"]

# Single prediction function
def predict_single(engine_rpm, lub_oil_pressure, fuel_pressure, coolant_pressure, lub_oil_temp, coolant_temp):
    features = np.array([[engine_rpm, lub_oil_pressure, fuel_pressure, coolant_pressure, lub_oil_temp, coolant_temp]])
    scaled_features = scaler.transform(features)
    prediction = model.predict(scaled_features)[0]
    return "Faulty" if prediction == 1 else "Active"

# Batch prediction function
def predict_batch(file):
    df = pd.read_csv(file)
    if not all(col in df.columns for col in feature_names):
        return "Error: CSV must contain columns: " + ", ".join(feature_names)
    scaled_features = scaler.transform(df[feature_names])
    predictions = model.predict(scaled_features)
    df["prediction"] = ["Faulty" if p == 1 else "Active" for p in predictions]
    output_file = "predictions.csv"
    df.to_csv(output_file, index=False)
    return output_file

# Gradio UI
single_inputs = [
    gr.Number(label="Engine RPM"),
    gr.Number(label="Lub Oil Pressure"),
    gr.Number(label="Fuel Pressure"),
    gr.Number(label="Coolant Pressure"),
    gr.Number(label="Lub Oil Temp"),
    gr.Number(label="Coolant Temp")
]

single_output = gr.Textbox(label="Engine Condition")

batch_input = gr.File(label="Upload CSV", file_types=[".csv"])
batch_output = gr.File(label="Download Predictions")

tab1 = gr.Interface(fn=predict_single, inputs=single_inputs, outputs=single_output, title="Single Prediction")
tab2 = gr.Interface(fn=predict_batch, inputs=batch_input, outputs=batch_output, title="Batch Prediction")

demo = gr.TabbedInterface([tab1, tab2], ["Single Prediction", "Batch Prediction"])
demo.launch(share=True)