import streamlit as st
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.ticker as mtick
import seaborn as sns
import joblib
import os
from sklearn.metrics import precision_recall_curve

st.set_page_config(page_title="Customer Churn Predictor", page_icon="📉", layout="wide")

COLORS = {"Logistic Regression": "#888780", "Random Forest": "#5DCAA5", "XGBoost": "#D85A30"}
BASE = os.path.dirname(__file__)

@st.cache_resource
def load_models():
    return joblib.load(os.path.join(BASE, "churn_models.pkl"))

@st.cache_data
def load_test_data():
    return joblib.load(os.path.join(BASE, "churn_test_data.pkl"))

# ── Load instantly ──
results  = load_models()
data     = load_test_data()
X_test   = data["X_test"]
y_test   = data["y_test"]
best_name = data["best_name"]
df       = data["df"]
y        = data["y"]
X        = data["X"]

# ── Sidebar ──
st.sidebar.title("📉 Navigation")
page = st.sidebar.radio("Go to", [
    "🏠 Overview", "📊 Model Comparison", "💰 Business Impact", "🔬 Predict"
])

# ── PAGE 1 – Overview ──
if page == "🏠 Overview":
    st.title("📉 Customer Churn Prediction")
    st.markdown(
        "End-to-end churn prediction on the **IBM Telco dataset** using "
        "Logistic Regression, Random Forest, and XGBoost — with business impact analysis."
    )
    c1, c2, c3, c4 = st.columns(4)
    c1.metric("Customers", len(df))
    c2.metric("Churn Rate", f"{y.mean():.1%}")
    c3.metric("Features", len(X.columns))
    c4.metric("Best AUC", f"{results[best_name]['roc_auc']:.3f} ({best_name})")

    col1, col2 = st.columns(2)
    with col1:
        st.subheader("Churn Distribution")
        fig, ax = plt.subplots(figsize=(5, 3))
        counts = df["Churn"].value_counts()
        ax.bar(counts.index, counts.values, color=["#5DCAA5", "#D85A30"])
        ax.set_ylabel("Customers")
        st.pyplot(fig); plt.close()

    with col2:
        st.subheader("Monthly Charges by Churn")
        fig, ax = plt.subplots(figsize=(5, 3))
        sns.boxplot(data=df, x="Churn", y="MonthlyCharges",
                    palette={"No": "#5DCAA5", "Yes": "#D85A30"}, ax=ax)
        st.pyplot(fig); plt.close()

    st.subheader("Churn Rate by Contract Type")
    contract_churn = df.groupby("Contract")["Churn"].apply(lambda x: (x == "Yes").mean())
    fig, ax = plt.subplots(figsize=(6, 3))
    contract_churn.sort_values().plot(kind="barh", ax=ax, color="#D85A30")
    ax.xaxis.set_major_formatter(mtick.PercentFormatter(xmax=1))
    st.pyplot(fig); plt.close()

# ── PAGE 2 – Model Comparison ──
elif page == "📊 Model Comparison":
    st.title("📊 Model Comparison")

    summary = pd.DataFrame([
        {"Model": n, "ROC-AUC": r["roc_auc"]}
        for n, r in results.items()
    ]).sort_values("ROC-AUC", ascending=False).reset_index(drop=True)
    st.dataframe(summary.style.format({"ROC-AUC": "{:.4f}"}), use_container_width=True)

    col1, col2 = st.columns(2)
    with col1:
        st.subheader("ROC Curves")
        fig, ax = plt.subplots(figsize=(6, 5))
        for name, res in results.items():
            ax.plot(res["fpr"], res["tpr"],
                    label=f"{name} (AUC={res['roc_auc']})",
                    color=COLORS[name], linewidth=2)
        ax.plot([0, 1], [0, 1], "k--")
        ax.set_xlabel("False Positive Rate"); ax.set_ylabel("True Positive Rate")
        ax.legend(fontsize=8)
        st.pyplot(fig); plt.close()

    with col2:
        st.subheader("Confusion Matrix")
        model_sel = st.selectbox("Select model", list(results.keys()))
        fig, ax = plt.subplots(figsize=(5, 4))
        sns.heatmap(results[model_sel]["cm"], annot=True, fmt="d", cmap="Blues",
                    xticklabels=["No Churn", "Churn"],
                    yticklabels=["No Churn", "Churn"], ax=ax)
        ax.set_xlabel("Predicted"); ax.set_ylabel("Actual")
        st.pyplot(fig); plt.close()

    st.subheader(f"Threshold Analysis — {best_name}")
    y_prob_best = results[best_name]["y_prob"]
    precisions, recalls, thresholds = precision_recall_curve(y_test, y_prob_best)
    f1s = 2 * precisions * recalls / (precisions + recalls + 1e-8)
    fig, ax = plt.subplots(figsize=(8, 4))
    ax.plot(thresholds, precisions[:-1], label="Precision", color="#5DCAA5", linewidth=2)
    ax.plot(thresholds, recalls[:-1],    label="Recall",    color="#D85A30", linewidth=2)
    ax.plot(thresholds, f1s[:-1],        label="F1",        color="#7F77DD", linewidth=2)
    best_t = thresholds[np.argmax(f1s[:-1])]
    ax.axvline(x=best_t, color="gray", linestyle="--", label=f"Best threshold={best_t:.2f}")
    ax.set_xlabel("Threshold"); ax.legend(fontsize=9)
    st.pyplot(fig); plt.close()

# ── PAGE 3 – Business Impact ──
elif page == "💰 Business Impact":
    st.title("💰 Business Impact Analysis")

    col1, col2 = st.columns(2)
    cost_fn = col1.number_input("Cost per missed churner ($)", value=250, step=10)
    cost_fp = col2.number_input("Cost per false retention offer ($)", value=15, step=5)

    y_prob_best = results[best_name]["y_prob"]
    baseline_cost = int(y_test.sum()) * cost_fn

    def compute_savings(threshold):
        y_pred = (y_prob_best >= threshold).astype(int)
        fp = int(((y_pred == 1) & (y_test == 0)).sum())
        fn = int(((y_pred == 0) & (y_test == 1)).sum())
        return {"threshold": threshold, "savings": baseline_cost - fp * cost_fp - fn * cost_fn}

    sweep_df = pd.DataFrame([compute_savings(t) for t in np.linspace(0.05, 0.95, 100)])
    best_idx = sweep_df["savings"].idxmax()
    best_t   = sweep_df.loc[best_idx, "threshold"]
    best_savings = sweep_df.loc[best_idx, "savings"]

    c1, c2, c3 = st.columns(3)
    c1.metric("Baseline cost (no model)", f"${baseline_cost:,}")
    c2.metric("Max savings with model",   f"${int(best_savings):,}")
    c3.metric("Optimal threshold",        f"{best_t:.2f}")

    fig, ax = plt.subplots(figsize=(9, 4))
    ax.plot(sweep_df["threshold"], sweep_df["savings"], color="#1D9E75", linewidth=2.5)
    ax.axvline(x=best_t, color="#D85A30", linestyle="--", label=f"Optimal = {best_t:.2f}")
    ax.axhline(y=0, color="gray", linewidth=1)
    ax.fill_between(sweep_df["threshold"], sweep_df["savings"], 0,
                    where=sweep_df["savings"] > 0, alpha=0.15, color="#1D9E75")
    ax.set_xlabel("Decision Threshold"); ax.set_ylabel("Net Savings ($)")
    ax.legend(); ax.yaxis.set_major_formatter(mtick.FuncFormatter(lambda x, _: f"${x:,.0f}"))
    st.pyplot(fig); plt.close()

    st.subheader("Risk Segmentation")
    res_df = X_test.copy().reset_index(drop=True)
    res_df["churn_probability"] = y_prob_best
    res_df["actual_churn"]      = y_test.values
    res_df["risk_tier"] = pd.cut(res_df["churn_probability"],
                                  bins=[0, 0.35, 0.60, 1.0],
                                  labels=["🟢 Low", "🟡 Medium", "🔴 High"])
    tier_summary = res_df.groupby("risk_tier", observed=True).agg(
        Customers=("churn_probability", "count"),
        Avg_Churn_Prob=("churn_probability", "mean"),
        Actual_Churn_Rate=("actual_churn", "mean"),
        Avg_Monthly_Charges=("MonthlyCharges", "mean"),
    ).round(3)
    st.dataframe(tier_summary, use_container_width=True)

# ── PAGE 4 – Predict ──
elif page == "🔬 Predict":
    st.title("🔬 Live Churn Prediction")
    model_sel = st.selectbox("Model", list(results.keys()),
                              index=list(results.keys()).index(best_name))

    col1, col2, col3 = st.columns(3)
    with col1:
        st.subheader("Account Info")
        tenure          = st.slider("Tenure (months)", 0, 72, 12)
        contract        = st.selectbox("Contract", ["Month-to-month", "One year", "Two year"])
        paperless       = st.selectbox("Paperless Billing", ["Yes", "No"])
        payment         = st.selectbox("Payment Method", [
            "Electronic check", "Mailed check",
            "Bank transfer (automatic)", "Credit card (automatic)"])
    with col2:
        st.subheader("Services")
        phone_service    = st.selectbox("Phone Service", ["Yes", "No"])
        multiple_lines   = st.selectbox("Multiple Lines", ["Yes", "No", "No phone service"])
        internet         = st.selectbox("Internet Service", ["DSL", "Fiber optic", "No"])
        online_security  = st.selectbox("Online Security", ["Yes", "No", "No internet service"])
        online_backup    = st.selectbox("Online Backup", ["Yes", "No", "No internet service"])
        device_prot      = st.selectbox("Device Protection", ["Yes", "No", "No internet service"])
        tech_support     = st.selectbox("Tech Support", ["Yes", "No", "No internet service"])
        streaming_tv     = st.selectbox("Streaming TV", ["Yes", "No", "No internet service"])
        streaming_movies = st.selectbox("Streaming Movies", ["Yes", "No", "No internet service"])
    with col3:
        st.subheader("Demographics & Charges")
        gender          = st.selectbox("Gender", ["Male", "Female"])
        partner         = st.selectbox("Partner", ["Yes", "No"])
        dependents      = st.selectbox("Dependents", ["Yes", "No"])
        monthly_charges = st.slider("Monthly Charges ($)", 18.0, 120.0, 65.0)
        total_charges   = st.number_input("Total Charges ($)", value=float(tenure * monthly_charges))

    if st.button("🔍 Predict Churn", type="primary"):
        input_df = pd.DataFrame([{
            "tenure": tenure, "MonthlyCharges": monthly_charges,
            "TotalCharges": total_charges,
            "charges_per_tenure": monthly_charges / (tenure + 1),
            "gender": gender, "Partner": partner, "Dependents": dependents,
            "PhoneService": phone_service, "MultipleLines": multiple_lines,
            "InternetService": internet, "OnlineSecurity": online_security,
            "OnlineBackup": online_backup, "DeviceProtection": device_prot,
            "TechSupport": tech_support, "StreamingTV": streaming_tv,
            "StreamingMovies": streaming_movies, "Contract": contract,
            "PaperlessBilling": paperless, "PaymentMethod": payment,
        }])
        prob = results[model_sel]["pipe"].predict_proba(input_df)[0][1]
        tier = "🔴 High Risk" if prob >= 0.60 else "🟡 Medium Risk" if prob >= 0.35 else "🟢 Low Risk"

        st.markdown("---")
        c1, c2, c3 = st.columns(3)
        c1.metric("Churn Probability", f"{prob:.1%}")
        c2.metric("Risk Tier", tier)
        c3.metric("Model", model_sel)

        fig, ax = plt.subplots(figsize=(5, 2.5))
        ax.barh(["No Churn", "Churn"], [(1-prob)*100, prob*100], color=["#5DCAA5", "#D85A30"])
        ax.set_xlabel("Probability (%)"); ax.set_xlim(0, 100)
        st.pyplot(fig); plt.close()
        st.caption("⚠️ For business decision support only.")