import os
import joblib
import pandas as pd
import gradio as gr
from huggingface_hub import hf_hub_download

# 1) SET YOUR MODEL REPO HERE (Model Hub repo, not dataset repo)
MODEL_REPO = os.getenv("MODEL_REPO", "SabarnaDeb/Capstone_PredictiveMaintenance_Model")
MODEL_FILE = os.getenv("MODEL_FILE", "model.joblib") 

# 2) Download model file from Hugging Face Model Hub
model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE, repo_type="model")
model = joblib.load(model_path)

# 3) Feature list must match your training columns
FEATURES = [
    "engine_rpm",
    "lub_oil_pressure",
    "fuel_pressure",
    "coolant_pressure",
    "lub_oil_temperature",
    "coolant_temperature"
]

def predict(engine_rpm, lub_oil_pressure, fuel_pressure, coolant_pressure,
            lub_oil_temperature, coolant_temperature):

    # 4) Save inputs into a DataFrame
    data = {
        "engine_rpm": [engine_rpm],
        "lub_oil_pressure": [lub_oil_pressure],
        "fuel_pressure": [fuel_pressure],
        "coolant_pressure": [coolant_pressure],
        "lub_oil_temperature": [lub_oil_temperature],
        "coolant_temperature": [coolant_temperature],
    }
    input_df = pd.DataFrame(data)

    # 5) Predict
    pred = model.predict(input_df[FEATURES])[0]

    prob = None
    if hasattr(model, "predict_proba"):
        prob = float(model.predict_proba(input_df[FEATURES])[:, 1][0])

    # 6) Business-friendly output
    if int(pred) == 1:
        msg = "⚠️ Maintenance Needed"
    else:
        msg = "✅ Normal Operation"

    if prob is not None:
        msg += f"\nConfidence (maintenance probability): {prob:.2f}"

    return msg, input_df

demo = gr.Interface(
    fn=predict,
    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 Temperature"),
        gr.Number(label="Coolant Temperature"),
    ],
    outputs=[
        gr.Textbox(label="Prediction Result"),
        gr.Dataframe(label="Input Data (saved as DataFrame)")
    ],
    title="Predictive Maintenance – Engine Health",
    description="Enter engine sensor readings to predict whether maintenance is needed."
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)