import joblib
import pandas as pd
import numpy as np
import gradio as gr

MODEL_PATH = "models/predictix_pm_calibrated_model.joblib"
COLS_PATH  = "models/predictix_pm_feature_columns.joblib"

model = joblib.load(MODEL_PATH)
feature_cols = joblib.load(COLS_PATH)

# Alert thresholds (use your tuned policy)
THRESH_WARNING = 0.60
THRESH_CRITICAL = 0.85

def risk_level(p: float) -> str:
    if p >= THRESH_CRITICAL:
        return "CRITICAL"
    if p >= THRESH_WARNING:
        return "WARNING"
    return "NORMAL"

def model_confidence(p: float) -> float:
    # Confidence = distance from 0.5 scaled to 0..1 (industrial-friendly)
    return float(abs(p - 0.5) * 2)

def predict(
    Year_of_Manufacture,
    Usage_Hours,
    Load_Capacity,
    Actual_Load,
    Engine_Temperature,
    Tire_Pressure,
    Fuel_Consumption,
    Battery_Status_Score,
    Vibration_Levels,
    Oil_Quality,
    Brake_Condition_Score,
    Load_Ratio,
    Overload_Flag,
    Days_Since_Last_Maintenance,
    Vehicle_Age_Years,
    Total_Operating_Hours,
    Total_Mileage_km,
    Lifetime_Maintenance_Count,
    Lifetime_Failure_Count,
    Lifetime_Downtime_Hours,
    Maintenance_Overdue_Flag,
    Vehicle_Type,         # categorical
    Route_Info,           # categorical
    Weather_Conditions,   # categorical
    Road_Conditions,      # categorical
    Make_and_Model        # categorical
):
    record = {
        "Year_of_Manufacture": Year_of_Manufacture,
        "Usage_Hours": Usage_Hours,
        "Load_Capacity": Load_Capacity,
        "Actual_Load": Actual_Load,
        "Engine_Temperature": Engine_Temperature,
        "Tire_Pressure": Tire_Pressure,
        "Fuel_Consumption": Fuel_Consumption,
        "Battery_Status_Score": Battery_Status_Score,
        "Vibration_Levels": Vibration_Levels,
        "Oil_Quality": Oil_Quality,
        "Brake_Condition_Score": Brake_Condition_Score,
        "Load_Ratio": Load_Ratio,
        "Overload_Flag": int(Overload_Flag),
        "Days_Since_Last_Maintenance": int(Days_Since_Last_Maintenance),
        "Vehicle_Age_Years": Vehicle_Age_Years,
        "Total_Operating_Hours": int(Total_Operating_Hours),
        "Total_Mileage_km": int(Total_Mileage_km),
        "Lifetime_Maintenance_Count": int(Lifetime_Maintenance_Count),
        "Lifetime_Failure_Count": int(Lifetime_Failure_Count),
        "Lifetime_Downtime_Hours": int(Lifetime_Downtime_Hours),
        "Maintenance_Overdue_Flag": int(Maintenance_Overdue_Flag),
        "Vehicle_Type": Vehicle_Type,
        "Route_Info": Route_Info,
        "Weather_Conditions": Weather_Conditions,
        "Road_Conditions": Road_Conditions,
        "Make_and_Model": Make_and_Model,
    }

    x = pd.DataFrame([record])
    x = pd.get_dummies(x)

    # add missing columns
    for c in feature_cols:
        if c not in x.columns:
            x[c] = 0
    # drop any extra columns not seen in training
    x = x.reindex(columns=feature_cols, fill_value=0)

    p = float(model.predict_proba(x)[:, 1][0])
    conf = model_confidence(p)

    return {
        "maintenance_probability_percent": round(p * 100, 2),
        "model_confidence_percent": round(conf * 100, 2),
        "risk_level": risk_level(p),
    }

demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Number(label="Year_of_Manufacture", value=2018),
        gr.Number(label="Usage_Hours", value=3500),
        gr.Number(label="Load_Capacity", value=5000),
        gr.Number(label="Actual_Load", value=4800),
        gr.Number(label="Engine_Temperature (°C)", value=96.5),
        gr.Number(label="Tire_Pressure", value=34),
        gr.Number(label="Fuel_Consumption", value=12.5),
        gr.Number(label="Battery_Status_Score", value=45.5),
        gr.Number(label="Vibration_Levels", value=5.1),
        gr.Number(label="Oil_Quality", value=78),
        gr.Number(label="Brake_Condition_Score", value=0.6),
        gr.Number(label="Load_Ratio", value=0.96),
        gr.Checkbox(label="Overload_Flag", value=False),
        gr.Number(label="Days_Since_Last_Maintenance", value=120),
        gr.Number(label="Vehicle_Age_Years", value=6),
        gr.Number(label="Total_Operating_Hours", value=58000),
        gr.Number(label="Total_Mileage_km", value=210000),
        gr.Number(label="Lifetime_Maintenance_Count", value=85),
        gr.Number(label="Lifetime_Failure_Count", value=2),
        gr.Number(label="Lifetime_Downtime_Hours", value=310),
        gr.Checkbox(label="Maintenance_Overdue_Flag", value=True),
        gr.Dropdown(["Van", "Truck"], label="Vehicle_Type", value="Van"),
        gr.Dropdown(["Highway", "Urban", "Rural"], label="Route_Info", value="Highway"),
        gr.Dropdown(["Clear", "Rainy", "Snowy", "Windy"], label="Weather_Conditions", value="Clear"),
        gr.Dropdown(["Highway", "Urban", "Rural"], label="Road_Conditions", value="Urban"),
        gr.Dropdown(["Ford F-150", "Chevy Silverado", "Volvo FH", "Tesla Semi"], label="Make_and_Model", value="Volvo FH"),
    ],
    outputs=gr.JSON(label="PredictiX Output"),
    title="PredictiX Predictive Maintenance (PDM) Demo",
    description="Returns maintenance probability (%), model confidence (%), and alert level."
)

if __name__ == "__main__":
    demo.launch()