app.py

import joblib
import numpy as np
import pandas as pd
import plotly.express as px
import shap
import streamlit as st
from pathlib import Path


st.set_page_config(page_title="Explainable Customer Segmentation", layout="wide")

APP_DIR = Path(__file__).resolve().parent
ARTIFACTS_DIR = APP_DIR / "artifacts"
DATA_DIR = APP_DIR / "data"
DATASET_FILE = DATA_DIR / "dataset.csv"

FEATURES = ["luxury_sales", "fresh_sales", "dry_sales"]
MEAN_ABS_SHAP_COL = "Mean |SHAP Value|"
PAGE_OVERVIEW = "Overview"
PAGE_EXPLORER = "Customer Explorer"
PAGE_XAI = "Why This Customer is in This Cluster (XAI)"
PAGE_GLOBAL = "Global Insights"
PAGE_CITY_PRED = "City Prediction (Probabilistic)"
PAGE_ANOMALY = "Customer Anomaly Detection"

CITY_ARTIFACTS = {
    "city_model": "city_prediction_model.pkl",
    "city_scaler": "city_scaler.pkl",
    "city_label_encoder": "city_label_encoder.pkl",
    "city_anomaly_stats": "city_anomaly_stats.pkl",
}


def _require_columns(df: pd.DataFrame, required: list[str]) -> None:
    missing = [c for c in required if c not in df.columns]
    if missing:
        raise ValueError(
            "dataset.csv is missing required columns: "
            + ", ".join(missing)
            + ".\nExpected at least: "
            + ", ".join(required)
        )


@st.cache_resource
def load_assets():
    gmm = joblib.load(ARTIFACTS_DIR / "customer_gmm.pkl")
    clf = joblib.load(ARTIFACTS_DIR / "cluster_explainer.pkl")
    scaler = joblib.load(ARTIFACTS_DIR / "scaler.pkl")

    city_model = joblib.load(ARTIFACTS_DIR / CITY_ARTIFACTS["city_model"])
    city_scaler = joblib.load(ARTIFACTS_DIR / CITY_ARTIFACTS["city_scaler"])
    city_label_encoder = joblib.load(ARTIFACTS_DIR / CITY_ARTIFACTS["city_label_encoder"])
    city_anomaly_stats = joblib.load(ARTIFACTS_DIR / CITY_ARTIFACTS["city_anomaly_stats"])

    df = pd.read_csv(DATASET_FILE)
    _require_columns(df, ["Customer_ID", "outlet_city", *FEATURES, "Cluster"])

    # Ensure correct dtypes (CSV may store numbers as text)
    for col in FEATURES:
        df[col] = pd.to_numeric(df[col], errors="coerce")
    df["Cluster"] = pd.to_numeric(df["Cluster"], errors="coerce")

    return gmm, clf, scaler, city_model, city_scaler, city_label_encoder, city_anomaly_stats, df


def scale_sales(scaler, luxury: float, fresh: float, dry: float) -> np.ndarray:
    """Return a (1, 3) array of scaled [luxury, fresh, dry].

    Handles both cases:
    - scaler trained on only the 3 sales features
    - scaler trained on a wider feature set (e.g., including one-hot outlet_city)
      by constructing a full feature row and then extracting the scaled sales.
    """

    values = {"luxury_sales": float(luxury), "fresh_sales": float(fresh), "dry_sales": float(dry)}

    feature_names = getattr(scaler, "feature_names_in_", None)
    n_features_in = getattr(scaler, "n_features_in_", None)

    if feature_names is not None:
        row = pd.DataFrame([dict.fromkeys(list(feature_names), 0.0)])
        for k, v in values.items():
            if k in row.columns:
                row.loc[0, k] = v
        scaled_full = scaler.transform(row)
        idx = [
            int(np.nonzero(feature_names == f)[0][0])
            if isinstance(feature_names, np.ndarray)
            else list(feature_names).index(f)
            for f in FEATURES
        ]
        return scaled_full[:, idx]

    if n_features_in is not None and int(n_features_in) != 3:
        raise ValueError(
            "Loaded scaler expects "
            + str(n_features_in)
            + " features, but feature names are not available. "
            "Re-export scaler with feature names (fit on a DataFrame) or fit scaler on only the 3 sales features."
        )

    x = np.array([[values["luxury_sales"], values["fresh_sales"], values["dry_sales"]]], dtype=float)
    return scaler.transform(x)


@st.cache_resource
def build_shap_artifacts(_clf, _scaler, df: pd.DataFrame):
    # Background sample: scaled sales only
    background_raw = df[FEATURES].dropna().sample(n=min(200, len(df)), random_state=42)
    background_scaled = np.vstack(
        [
            scale_sales(_scaler, r["luxury_sales"], r["fresh_sales"], r["dry_sales"])
            for _, r in background_raw.iterrows()
        ]
    )
    explainer = shap.Explainer(_clf, background_scaled)

    # Global sample for importance (kept small to stay responsive)
    global_raw = df[FEATURES].dropna().sample(n=min(500, len(df)), random_state=42)
    global_scaled = np.vstack(
        [
            scale_sales(_scaler, r["luxury_sales"], r["fresh_sales"], r["dry_sales"])
            for _, r in global_raw.iterrows()
        ]
    )

    return explainer, global_scaled


@st.cache_data
def compute_global_importance(_explainer, _clf, global_scaled: np.ndarray, feature_names: list[str]):
    shap_values = _explainer(global_scaled)
    predicted = _clf.predict(global_scaled)

    selected = []
    for i in range(len(global_scaled)):
        class_id = int(predicted[i])
        selected.append(shap_values.values[i, :, class_id])
    selected = np.asarray(selected)

    mean_abs = np.abs(selected).mean(axis=0)
    imp = pd.DataFrame({"Feature": feature_names, MEAN_ABS_SHAP_COL: mean_abs}).sort_values(
        MEAN_ABS_SHAP_COL, ascending=False
    )
    return imp


def render_overview(df: pd.DataFrame):
    st.title("Explainable Customer Segmentation")

    st.markdown(
        """
**What is clustering?**
Clustering groups similar customers together based on their behavior.

**What is GMM?**
Gaussian Mixture Models (GMM) perform probabilistic clustering by modeling data as a mixture of Gaussian distributions.

**What is SHAP?**
SHAP (SHapley Additive exPlanations) explains model predictions by attributing contributions to each feature.

**Academic note (compliance):**
- The **primary segmentation model is the GMM**.
- **SHAP is used for explainability**.
- An **XGBoost surrogate classifier** is used only to enable SHAP explanations of the GMM clusters.
"""
    )

    st.subheader("Cluster profiles (mean spending)")
    profiles = df.groupby("Cluster")[FEATURES].mean().reset_index()
    st.dataframe(profiles, use_container_width=True)


def render_customer_explorer(gmm, clf, scaler, df: pd.DataFrame):
    st.title("Customer Explorer")

    mode = st.radio("Input method", ["Select existing customer", "Manual entry"], horizontal=True)

    luxury = fresh = dry = None
    selected_customer_id = None

    if mode == "Select existing customer":
        selected_customer_id = st.selectbox("Customer_ID", df["Customer_ID"].astype(str).tolist())
        row = df[df["Customer_ID"].astype(str) == str(selected_customer_id)].iloc[0]
        luxury, fresh, dry = float(row["luxury_sales"]), float(row["fresh_sales"]), float(row["dry_sales"])

        st.subheader("Selected customer spending profile")
        st.dataframe(
            pd.DataFrame(
                [{"Customer_ID": selected_customer_id, "luxury_sales": luxury, "fresh_sales": fresh, "dry_sales": dry}]
            ),
            use_container_width=True,
        )
    else:
        c1, c2, c3 = st.columns(3)
        with c1:
            luxury = st.number_input("Luxury Sales", min_value=0.0, value=0.0, step=1.0)
        with c2:
            fresh = st.number_input("Fresh Sales", min_value=0.0, value=0.0, step=1.0)
        with c3:
            dry = st.number_input("Dry Sales", min_value=0.0, value=0.0, step=1.0)

    if luxury is None or fresh is None or dry is None:
        return

    x_scaled = scale_sales(scaler, luxury, fresh, dry)

    cluster = int(gmm.predict(x_scaled)[0])
    probs = clf.predict_proba(x_scaled)[0]

    st.subheader("Prediction")
    st.write(f"Predicted cluster (segment): **Cluster {cluster}**")

    prob_df = pd.DataFrame({"Cluster": [f"Cluster {i}" for i in range(len(probs))], "Probability": probs})
    fig = px.bar(prob_df, x="Cluster", y="Probability", title="Cluster probabilities")
    st.plotly_chart(fig, use_container_width=True)


def render_xai(gmm, clf, scaler, df: pd.DataFrame, explainer):
    st.title("Why This Customer is in This Cluster (XAI)")

    selected_customer_id = st.selectbox("Customer_ID to explain", df["Customer_ID"].astype(str).tolist())
    row = df[df["Customer_ID"].astype(str) == str(selected_customer_id)].iloc[0]
    luxury, fresh, dry = float(row["luxury_sales"]), float(row["fresh_sales"]), float(row["dry_sales"])

    x_scaled = scale_sales(scaler, luxury, fresh, dry)
    pred_cluster = int(gmm.predict(x_scaled)[0])
    probs = clf.predict_proba(x_scaled)[0]

    st.subheader("Customer spending profile")
    st.dataframe(
        pd.DataFrame(
            [{"Customer_ID": selected_customer_id, "luxury_sales": luxury, "fresh_sales": fresh, "dry_sales": dry}]
        ),
        use_container_width=True,
    )

    st.subheader("Predicted cluster")
    st.write(f"Segment: **Cluster {pred_cluster}**")

    shap_values = explainer(x_scaled)

    # Multi-class output: values is (n, features, classes)
    class_id = int(np.argmax(probs))
    values = shap_values.values[0, :, class_id]
    base_values = shap_values.base_values[0, class_id] if np.ndim(shap_values.base_values) == 2 else shap_values.base_values[0]

    explanation = shap.Explanation(
        values=values,
        base_values=base_values,
        data=x_scaled[0],
        feature_names=FEATURES,
    )

    st.subheader("SHAP waterfall plot")
    import matplotlib.pyplot as plt

    plt.figure(figsize=(8, 4))
    shap.plots.waterfall(explanation, show=False)
    st.pyplot(plt.gcf(), clear_figure=True)

    st.subheader("Feature contributions")
    contrib = pd.DataFrame({"Feature": FEATURES, "Contribution": values}).sort_values("Contribution", ascending=False)
    st.dataframe(contrib, use_container_width=True)

    # Simple-language summary
    top = contrib.iloc[0]
    direction = "high" if float(top["Contribution"]) > 0 else "low"
    st.markdown(
        f"This customer belongs to **Cluster {pred_cluster}** mainly because their **{top['Feature']}** is {direction} (relative to the model background), which most strongly pushes the prediction toward this cluster."
    )


def render_global_insights(df: pd.DataFrame, global_importance: pd.DataFrame):
    st.title("Global Insights")

    st.subheader("SHAP global feature importance")
    fig_imp = px.bar(
        global_importance,
        x="Mean |SHAP Value|",
        y="Feature",
        orientation="h",
        title="Global feature importance (mean absolute SHAP)",
    )
    st.plotly_chart(fig_imp, use_container_width=True)

    st.subheader("Customer segments scatter plot")
    fig_scatter = px.scatter(
        df,
        x="luxury_sales",
        y="dry_sales",
        color=df["Cluster"].astype(str),
        title="luxury_sales vs dry_sales colored by cluster",
        labels={"color": "Cluster"},
        opacity=0.7,
    )
    st.plotly_chart(fig_scatter, use_container_width=True)


def _shared_spending_inputs(form_key: str):
    c1, c2, c3 = st.columns(3)
    with c1:
        luxury = st.number_input("Luxury Sales", min_value=0.0, value=0.0, step=1.0, key=f"{form_key}_lux")
    with c2:
        fresh = st.number_input("Fresh Sales", min_value=0.0, value=0.0, step=1.0, key=f"{form_key}_fresh")
    with c3:
        dry = st.number_input("Dry Sales", min_value=0.0, value=0.0, step=1.0, key=f"{form_key}_dry")
    return float(luxury), float(fresh), float(dry)


def _city_academic_expander():
    with st.expander("Academic compliance note"):
        st.markdown(
            "This model uses supervised learning for city prediction and statistical anomaly detection.\n"
            "The primary clustering model remains Gaussian Mixture Model."
        )


def render_city_prediction(city_model, city_scaler, city_label_encoder):
    st.title(PAGE_CITY_PRED)

    with st.form("city_pred_form"):
        luxury, fresh, dry = _shared_spending_inputs("city_pred")
        submitted = st.form_submit_button("Predict City")

    _city_academic_expander()

    if not submitted:
        return

    X = city_scaler.transform([[luxury, fresh, dry]])
    probs = city_model.predict_proba(X)[0]
    cities = city_label_encoder.inverse_transform(np.arange(len(probs)))

    pred_df = pd.DataFrame({"City": cities, "Probability": probs}).sort_values("Probability", ascending=False)
    top5 = pred_df.head(5).copy()
    top5["ProbabilityPct"] = (top5["Probability"] * 100).round(2)

    st.subheader("Top 5 predicted cities")
    fig = px.bar(
        top5.sort_values("Probability", ascending=True),
        x="ProbabilityPct",
        y="City",
        orientation="h",
        title="Top 5 outlet cities with probabilities (%)",
        labels={"ProbabilityPct": "Probability (%)"},
    )
    st.plotly_chart(fig, use_container_width=True)

    st.dataframe(top5[["City", "ProbabilityPct"]].rename(columns={"ProbabilityPct": "Probability (%)"}), use_container_width=True)

    # Natural-language explanation (as required in spec)
    parts = [f"**{row.City} ({row.ProbabilityPct:.2f}%)**" for row in top5.itertuples(index=False)]
    head = ", followed by ".join([parts[0], ", ".join(parts[1:3])]) if len(parts) >= 3 else ", ".join(parts)

    st.markdown(
        f"This customer is most likely associated with {head}.\n"
        "High dry and fresh spending contributed most to this prediction."
    )


def _coerce_anomaly_stats(stats_obj) -> pd.DataFrame:
    if isinstance(stats_obj, pd.DataFrame):
        return stats_obj
    if isinstance(stats_obj, dict):
        return pd.DataFrame(stats_obj)
    raise TypeError("city_anomaly_stats.pkl must be a pandas DataFrame or dict-like object.")


def _extract_city_feature_stats(city_row: pd.Series, feature: str) -> tuple[float, float]:
    # Accept a few common naming conventions
    candidates = [
        (f"{feature}_mean", f"{feature}_std"),
        (f"mean_{feature}", f"std_{feature}"),
        (f"{feature}.mean", f"{feature}.std"),
    ]

    for mean_key, std_key in candidates:
        if mean_key in city_row.index and std_key in city_row.index:
            mean_val = float(city_row[mean_key])
            std_val = float(city_row[std_key])
            return mean_val, std_val

    # Also support a multi-index column layout flattened into tuples
    if isinstance(city_row.index, pd.MultiIndex):
        try:
            mean_val = float(city_row[(feature, "mean")])
            std_val = float(city_row[(feature, "std")])
            return mean_val, std_val
        except Exception:
            pass

    raise KeyError(
        f"Cannot find mean/std columns for '{feature}' in city_anomaly_stats. "
        "Expected columns like '{feature}_mean' and '{feature}_std'."
    )


def render_anomaly_detection(city_anomaly_stats):
    st.title(PAGE_ANOMALY)

    stats_df = _coerce_anomaly_stats(city_anomaly_stats)
    if stats_df.index.name is None and "outlet_city" in stats_df.columns:
        stats_df = stats_df.set_index("outlet_city")

    cities = stats_df.index.astype(str).tolist()
    if not cities:
        st.error("No cities found in city_anomaly_stats.")
        return

    with st.form("anomaly_form"):
        luxury, fresh, dry = _shared_spending_inputs("anomaly")
        selected_city = st.selectbox("City", cities)
        submitted = st.form_submit_button("Check Anomaly")

    _city_academic_expander()

    if not submitted:
        return

    city_row = stats_df.loc[str(selected_city)]
    rows = []
    anomalous_reasons = []

    for feature, value in zip(FEATURES, [luxury, fresh, dry]):
        mean_val, std_val = _extract_city_feature_stats(city_row, feature)
        std_val = float(std_val)
        z = (float(value) - float(mean_val)) / std_val if std_val not in (0.0, np.nan) else np.nan

        is_anom = bool(np.isfinite(z) and abs(z) > 3)
        status = "Anomalous" if is_anom else "Normal"

        if is_anom:
            anomalous_reasons.append(f"{feature} is {abs(z):.1f} standard deviations {'above' if z > 0 else 'below'} the {selected_city} average.")

        rows.append(
            {
                "Feature": feature,
                "Value": float(value),
                "City Mean": float(mean_val),
                "City Std": float(std_val),
                "z-score": float(z) if np.isfinite(z) else np.nan,
                "Status": status,
            }
        )

    result_df = pd.DataFrame(rows)

    any_anom = (result_df["Status"] == "Anomalous").any()
    if any_anom:
        st.error("Anomalous")
    else:
        st.success("Normal")

    def _style_status(row):
        bg = "#ffdddd" if row["Status"] == "Anomalous" else "#ddffdd"
        # Force readable text in dark theme (Streamlit may default to white text)
        return [f"background-color: {bg}; color: #000000" for _ in row]

    st.subheader("Anomaly explanation")
    st.dataframe(result_df.style.apply(_style_status, axis=1), use_container_width=True)

    if any_anom:
        for reason in anomalous_reasons:
            st.write(reason)
    else:
        st.write("This customer’s spending is within the normal range for this city.")


def main():
    gmm, clf, scaler, city_model, city_scaler, city_label_encoder, city_anomaly_stats, df = load_assets()
    explainer, global_scaled = build_shap_artifacts(clf, scaler, df)
    global_importance = compute_global_importance(explainer, clf, global_scaled, FEATURES)

    pages = [
        PAGE_OVERVIEW,
        PAGE_EXPLORER,
        PAGE_XAI,
        PAGE_GLOBAL,
        PAGE_CITY_PRED,
        PAGE_ANOMALY,
    ]

    st.sidebar.markdown("### Menu")
    if "page" not in st.session_state:
        st.session_state.page = pages[0]

    def _set_page(p: str) -> None:
        st.session_state.page = p

    for i, p in enumerate(pages):
        is_active = st.session_state.page == p
        st.sidebar.button(
            p,
            key=f"nav_{i}",
            type="primary" if is_active else "secondary",
            use_container_width=True,
            on_click=_set_page,
            args=(p,),
        )

    page = st.session_state.page

    if page == PAGE_OVERVIEW:
        render_overview(df)
    elif page == PAGE_EXPLORER:
        render_customer_explorer(gmm, clf, scaler, df)
    elif page == PAGE_XAI:
        render_xai(gmm, clf, scaler, df, explainer)
    elif page == PAGE_CITY_PRED:
        render_city_prediction(city_model, city_scaler, city_label_encoder)
    elif page == PAGE_ANOMALY:
        render_anomaly_detection(city_anomaly_stats)
    else:
        render_global_insights(df, global_importance)


if __name__ == "__main__":
    main()