app.py

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

# ======================
# CUSTOM STACKING CLASS
# ======================
class SimpleStacking:
    def __init__(self, base_models, final_model):
        self.base_models = base_models
        self.final_model = final_model

    def predict(self, X):
        meta_features = []

        for model in self.base_models.values():
            if hasattr(model, "predict_proba"):
                probs = model.predict_proba(X)
                meta_features.append(probs)
            else:
                preds = model.predict(X)
                meta_features.append(preds)

        meta_X = np.column_stack(meta_features)
        return self.final_model.predict(meta_X)

    def predict_proba(self, X):
        meta_features = []

        for model in self.base_models.values():
            if hasattr(model, "predict_proba"):
                probs = model.predict_proba(X)
                meta_features.append(probs)
            else:
                preds = model.predict(X)
                meta_features.append(preds)

        meta_X = np.column_stack(meta_features)

        if hasattr(self.final_model, "predict_proba"):
            return self.final_model.predict_proba(meta_X)
        else:
            preds = self.final_model.predict(meta_X)
            # For simplicity, assume predict gives class, but need probs
            return preds


# ======================
# LOAD MODELS
# ======================
xgboost_model = joblib.load("models/xgboost_final_complete.pkl")
svm_model = joblib.load("models/svm_final_complete.pkl")
knn_model = joblib.load("models/knn_final_complete.pkl")

stacking_model = joblib.load("models/stacking_final_complete.pkl")

scaler = joblib.load("models/scaler_final_complete.pkl")
imputer = joblib.load("models/imputer_final_complete.pkl")


# ======================
# PREDICT FUNCTION
# ======================
def predict_potability(ph, tds, turbidity, temperature):

    if ph is None or not (0 <= ph <= 14):
        return "❌ pH must be 0–14", "", "", ""

    if tds is None or tds < 0:
        return "❌ Invalid TDS", "", "", ""

    if turbidity is None or turbidity < 0:
        return "❌ Invalid Turbidity", "", "", ""

    if temperature is None or temperature < 0 or temperature > 100:
        return "❌ Temperature must be 0–100°C", "", "", ""

    # Correct pH for temperature
    corrected_ph = ph + 0.0002 * (temperature - 25)

    input_df = pd.DataFrame(
        [[corrected_ph, tds, turbidity]],
        columns=["ph", "TDS", "Turbidity"]
    )

    imputed = imputer.transform(input_df)
    scaled = scaler.transform(imputed)

    models = {
        "XGBoost": xgboost_model,
        "SVM": svm_model,
        "KNN": knn_model,
        "Stacking": stacking_model
    }

    class_names = ["Non-Potable", "Borderline", "Potable"]

    outputs = {}

    for name, model in models.items():
        pred = model.predict(scaled)[0]

        if hasattr(model, "predict_proba"):
            probs = model.predict_proba(scaled)[0]
            prob_str = ", ".join([f"{class_names[i]}: {probs[i]:.4f}" for i in range(3)])
        else:
            prob_str = "N/A"

        label = class_names[pred]

        outputs[name] = f"{label} (Probabilities: {prob_str})"

    return (
        outputs["XGBoost"],
        outputs["SVM"],
        outputs["KNN"],
        outputs["Stacking"]
    )


# ======================
# UI
# ======================
with gr.Blocks(title="Water Potability Predictor") as demo:

    gr.Markdown("# 💧 Water Potability Prediction")
    gr.Markdown("Enter water parameters to check if water is drinkable.")

    with gr.Row():
        ph = gr.Number(label="pH", value=7.0)
        tds = gr.Number(label="TDS (mg/L)", value=200)
        turbidity = gr.Number(label="Turbidity (NTU)", value=3.0)
        temperature = gr.Number(label="Temperature (°C)", value=20.0)

    predict_btn = gr.Button("Predict")

    with gr.Row():
        xgboost_out = gr.Textbox(label="XGBoost", interactive=False)
        svm_out = gr.Textbox(label="SVM", interactive=False)

    with gr.Row():
        knn_out = gr.Textbox(label="KNN", interactive=False)
        stacking_out = gr.Textbox(label="Stacking (Final)", interactive=False)

    predict_btn.click(
        fn=predict_potability,
        inputs=[ph, tds, turbidity, temperature],
        outputs=[xgboost_out, svm_out, knn_out, stacking_out]
    )

demo.launch(server_name="0.0.0.0", server_port=7860,share=True)