src/customer_intelligence/ml/churn.py

"""XGBoost churn prediction model."""

from datetime import date

import numpy as np
import pandas as pd
import xgboost as xgb
from sklearn.metrics import (
    average_precision_score,
    f1_score,
    precision_score,
    recall_score,
    roc_auc_score,
)

CHURN_THRESHOLD_DAYS = 180
CHURN_PROBA_CUTOFF = 0.5

FEATURE_COLS = [
    # NOTE: recency_days intentionally excluded — it IS the churn target (leakage)
    "frequency",
    "monetary",
    "avg_order_value",
    "avg_review_score",
    "avg_days_to_delivery",
    "late_delivery_rate",
    "avg_delivery_delay",
    "complaint_rate",
    "tenure_days",
]


class ChurnModel:
    def __init__(self) -> None:
        self.model: xgb.XGBClassifier | None = None
        self.feature_cols = FEATURE_COLS

    def define_churn_target(
        self, rfm_df: pd.DataFrame, reference_date: date | None = None
    ) -> pd.Series:
        """Label a customer as churned if recency > CHURN_THRESHOLD_DAYS."""
        if reference_date is None:
            reference_date = date.today()
        churned = (rfm_df["recency_days"] > CHURN_THRESHOLD_DAYS).astype(int)
        churned.index = rfm_df["customer_unique_id"]
        return churned

    def build_feature_matrix(
        self,
        rfm_df: pd.DataFrame,
        clv_df: pd.DataFrame,
        delivery_df: pd.DataFrame,
    ) -> tuple[pd.DataFrame, pd.Series]:
        """Merge all feature sources and return (X, y)."""
        df = (
            rfm_df[["customer_unique_id", "recency_days", "frequency", "monetary"]]
            .merge(
                clv_df[[
                    "customer_unique_id",
                    "avg_order_value",
                    "avg_review_score",
                    "avg_days_to_delivery",
                    "late_delivery_rate",
                    "tenure_days",
                ]],
                on="customer_unique_id",
                how="inner",
            )
            .merge(
                delivery_df[["customer_unique_id", "avg_delivery_delay", "complaint_rate"]],
                on="customer_unique_id",
                how="left",
            )
        )

        df["avg_delivery_delay"] = df["avg_delivery_delay"].fillna(0)
        df["complaint_rate"] = df["complaint_rate"].fillna(0)
        df = df.dropna(subset=self.feature_cols)

        # Churn target: no purchase in last CHURN_THRESHOLD_DAYS days
        y = (df["recency_days"] > CHURN_THRESHOLD_DAYS).astype(int).values

        return df, y

    def fit(self, X: np.ndarray, y: np.ndarray) -> "ChurnModel":
        neg, pos = (y == 0).sum(), (y == 1).sum()
        scale_pos = neg / pos if pos > 0 else 1.0

        self.model = xgb.XGBClassifier(
            n_estimators=200,
            max_depth=5,
            learning_rate=0.05,
            scale_pos_weight=scale_pos,
            random_state=42,
            eval_metric="auc",
            verbosity=0,
        )
        self.model.fit(X, y)
        return self

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        if self.model is None:
            raise RuntimeError("Model not fitted.")
        return self.model.predict_proba(X)[:, 1]

    def predict(self, X: np.ndarray, threshold: float = CHURN_PROBA_CUTOFF) -> np.ndarray:
        return (self.predict_proba(X) >= threshold).astype(int)

    def evaluate(self, X: np.ndarray, y: np.ndarray) -> dict:
        proba = self.predict_proba(X)
        labels = (proba >= CHURN_PROBA_CUTOFF).astype(int)
        return {
            "auc_roc": float(roc_auc_score(y, proba)),
            "avg_precision": float(average_precision_score(y, proba)),
            "f1": float(f1_score(y, labels, zero_division=0)),
            "precision": float(precision_score(y, labels, zero_division=0)),
            "recall": float(recall_score(y, labels, zero_division=0)),
            "churn_rate": float(y.mean()),
        }