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

# Load model, feature names, and scaler
model = joblib.load("churn_model.pkl")
feature_cols = joblib.load("feature_cols.pkl")
scaler = joblib.load("scaler.pkl")

# Create properly encoded sample data
def make_sample(values_dict):
    """Create a DataFrame row matching feature_cols from a simple dict of values."""
    row = {col: 0 for col in feature_cols}  # default zeros
    for key, val in values_dict.items():
        if key in row:
            row[key] = val
    return row

# Example pre-encoded samples
high_risk_sample = make_sample({
    "gender_Male": 1,
    "SeniorCitizen": 1,
    "tenure": 2,
    "MonthlyCharges": 95,
    "Contract_Two year": 0,
    "Contract_One year": 0,
    "Contract_Month-to-month": 1
})

low_risk_sample = make_sample({
    "gender_Female": 1,
    "SeniorCitizen": 0,
    "tenure": 60,
    "MonthlyCharges": 30,
    "Contract_Two year": 1,
    "Contract_One year": 0,
    "Contract_Month-to-month": 0
})

# Prediction function
def predict_churn(*inputs):
    data = np.array(inputs).reshape(1, -1)
    X_scaled = scaler.transform(data)
    pred = model.predict(X_scaled)[0]
    prob = model.predict_proba(X_scaled)[0][1]
    risk_label = "🔴 High Risk" if pred == 1 else "🟢 Low Risk"
    return f"{risk_label} ({prob:.2%} chance of churn)"

# UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 📊 Customer Churn Prediction")

    with gr.Row():
        inputs = []
        for col in feature_cols:
            inputs.append(gr.Number(label=col, value=0))

    with gr.Row():
        predict_btn = gr.Button("🚀 Predict")
        high_btn = gr.Button("📌 Load High Risk Sample")
        low_btn = gr.Button("📌 Load Low Risk Sample")

    output = gr.Textbox(label="Prediction")

    predict_btn.click(fn=predict_churn, inputs=inputs, outputs=output)

    def fill_high():
        return list(high_risk_sample.values())

    def fill_low():
        return list(low_risk_sample.values())

    high_btn.click(fn=fill_high, outputs=inputs)
    low_btn.click(fn=fill_low, outputs=inputs)

demo.launch()