import gradio as gr
import shap
import xgboost as xgb
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# -----------------------------
# 1. Train a simple clean model
# -----------------------------
np.random.seed(0)
X = pd.DataFrame({
    "age": np.random.randint(20, 80, 200),
    "bmi": np.random.uniform(18, 40, 200),
    "bp": np.random.uniform(80, 160, 200)
})
y = 0.3 * X["age"] + 0.5 * X["bmi"] + 0.2 * X["bp"] + np.random.normal(0, 3, 200)

model = xgb.XGBRegressor(
    n_estimators=100,
    max_depth=3,
    learning_rate=0.1,
    objective="reg:squarederror"
)
model.fit(X, y)

# Use TreeExplainer ONLY with XGBoost <= 1.7.6
explainer = shap.TreeExplainer(model)

# -----------------------------
# 2. Prediction + SHAP plot
# -----------------------------
def predict_and_explain(age, bmi, bp):

    df = pd.DataFrame([{"age": age, "bmi": bmi, "bp": bp}])
    pred = float(model.predict(df)[0])

    shap_values = explainer(df)

    # Make SHAP waterfall plot
    plt.figure(figsize=(8, 5))
    shap.plots.waterfall(shap_values[0], show=False)
    plt.tight_layout()
    plt.savefig("shap_plot.png")
    plt.close()

    return pred, "shap_plot.png"

# -----------------------------
# 3. Gradio UI
# -----------------------------
inputs = [
    gr.Number(label="Age"),
    gr.Number(label="BMI"),
    gr.Number(label="Blood Pressure"),
]

outputs = [
    gr.Number(label="Prediction"),
    gr.Image(label="SHAP Explanation"),
]

app = gr.Interface(
    fn=predict_and_explain,
    inputs=inputs,
    outputs=outputs,
    title="XGBoost + SHAP Deployment Fix",
    description="Guaranteed working version for Hugging Face Spaces."
)

if __name__ == "__main__":
    app.launch(share=True)