import os
import joblib
import pandas as pd
import gradio as gr

# --- 1) Загружаем модель ---
# Положи свой файл рядом (например, model.joblib)
# или укажи путь через переменную окружения MODEL_PATH
MODEL_PATH = os.getenv("MODEL_PATH", "model.joblib")
model = joblib.load(MODEL_PATH)

# Имена колонок должны совпадать с теми, на которых обучалась модель
FEATURE_COLUMNS = [
    "Type",
    "Air temperature [K]",
    "Process temperature [K]",
    "Rotational speed [rpm]",
    "Torque [Nm]",
    "Tool wear [min]",
]

TYPE_CHOICES = ["L", "M", "H"]  # под твой датасет

# --- 2) Функция предсказания для одной записи ---
def predict_single(_type, air_k, proc_k, rpm, torque, wear):
    row = {
        "Type": _type,
        "Air temperature [K]": float(air_k),
        "Process temperature [K]": float(proc_k),
        "Rotational speed [rpm]": float(rpm),
        "Torque [Nm]": float(torque),
        "Tool wear [min]": float(wear),
    }
    X = pd.DataFrame([row], columns=FEATURE_COLUMNS)

    if hasattr(model, "predict_proba"):
        proba = float(model.predict_proba(X)[:, 1][0])
        label = int(proba >= 0.5)
    else:
        # на случай, если модели нет predict_proba (вернём 0/1)
        label = int(model.predict(X)[0])
        proba = float(label)

    return round(proba, 4), int(label)

# --- 3) Пакетный режим (CSV/DataFrame) ---
def predict_batch(df: pd.DataFrame):
    # Проверяем, что нужные колонки есть
    missing = [c for c in FEATURE_COLUMNS if c not in df.columns]
    if missing:
        raise gr.Error(f"В CSV не хватает колонок: {missing}")

    X = df[FEATURE_COLUMNS].copy()
    if hasattr(model, "predict_proba"):
        probs = model.predict_proba(X)[:, 1]
        labels = (probs >= 0.5).astype(int)
    else:
        labels = model.predict(X)
        probs = labels.astype(float)

    out = df.copy()
    out["failure_probability"] = probs.round(4)
    out["predicted_label"] = labels
    return out

# --- 4) UI ---
with gr.Blocks(title="Predictive Maintenance (Gradio)") as demo:
    gr.Markdown("## Predictive Maintenance – Inference\nВвод параметров → предсказание вероятности отказа.")

    with gr.Tab("Single"):
        with gr.Row():
            t  = gr.Dropdown(TYPE_CHOICES, value="L", label="Type")
            aK = gr.Number(value=300.0, label="Air temperature [K]")
            pK = gr.Number(value=310.0, label="Process temperature [K]")
        with gr.Row():
            rpm = gr.Number(value=1500, label="Rotational speed [rpm]")
            tq  = gr.Number(value=40.0, label="Torque [Nm]")
            tw  = gr.Number(value=120.0, label="Tool wear [min]")

        btn = gr.Button("Predict")
        prob = gr.Number(label="failure_probability")
        labl = gr.Number(label="predicted_label (0/1)")

        # api_name="predict" — даёт REST-эндпоинт /run/predict
        btn.click(
            predict_single,
            inputs=[t, aK, pK, rpm, tq, tw],
            outputs=[prob, labl],
            api_name="predict"
        )

    with gr.Tab("Batch (CSV)"):
        gr.Markdown(
            "Загрузи CSV с колонками: "
            + ", ".join([f"`{c}`" for c in FEATURE_COLUMNS])
        )
        df_in  = gr.Dataframe(headers=FEATURE_COLUMNS, label="Input Data", row_count=2)
        btn_b  = gr.Button("Predict for all rows")
        df_out = gr.Dataframe(label="Output with predictions")
        btn_b.click(predict_batch, inputs=df_in, outputs=df_out, api_name="predict_batch")

# Запуск
if __name__ == "__main__":
    demo.launch()