src/customer_intelligence/ml/train.py

"""Training orchestrator: feature computation → model fit → MLflow logging → registry."""

import mlflow
import mlflow.sklearn
import mlflow.xgboost
from sklearn.model_selection import train_test_split
from sqlalchemy import Engine

from customer_intelligence.config import settings
from customer_intelligence.features.clv import compute_clv_features
from customer_intelligence.features.delivery import compute_delivery_features
from customer_intelligence.features.rfm import compute_rfm
from customer_intelligence.ml.churn import ChurnModel
from customer_intelligence.ml.registry import register_model, write_scores_to_warehouse
from customer_intelligence.ml.segmentation import CustomerSegmentationModel


def run_segmentation_training(
    engine: Engine,
    n_clusters: int = 5,
    mlflow_tracking_uri: str | None = None,
) -> str:
    """Train KMeans segmentation, log to MLflow, register as Production."""
    if mlflow_tracking_uri:
        mlflow.set_tracking_uri(mlflow_tracking_uri)
    else:
        mlflow.set_tracking_uri(settings.mlflow_tracking_uri)

    mlflow.set_experiment("customer_segmentation")

    print("Computing RFM features...")
    rfm_df = compute_rfm(engine)
    print("Computing CLV features...")
    clv_df = compute_clv_features(engine)

    model = CustomerSegmentationModel(n_clusters=n_clusters)
    index_df, X = model.build_feature_matrix(rfm_df, clv_df)

    print(f"Training KMeans (n_clusters={n_clusters}) on {len(X):,} customers...")

    with mlflow.start_run(run_name="kmeans_segmentation") as run:
        model.fit(X)
        metrics = model.evaluate(X)
        labels = model.predict(X)

        mlflow.log_params({"n_clusters": n_clusters, "algorithm": "kmeans"})
        mlflow.log_metrics(metrics)
        mlflow.sklearn.log_model(model.pipeline, "segmentation_pipeline")

        print(f"  silhouette={metrics['silhouette_score']:.4f}, inertia={metrics['inertia']:.0f}")

        run_id = run.info.run_id

    register_model(run_id, "customer_segmentation", "segmentation_pipeline")

    # Write segment labels back to warehouse
    index_df = index_df.copy()
    index_df["segment_label"] = labels
    scores = index_df[["customer_unique_id", "segment_label"]]

    # Also add RFM segment and CLV score
    scores = scores.merge(rfm_df[["customer_unique_id", "rfm_segment"]], on="customer_unique_id", how="left")
    scores = scores.merge(clv_df[["customer_unique_id", "predicted_clv"]], on="customer_unique_id", how="left")
    scores = scores.rename(columns={"predicted_clv": "clv_score"})

    print(f"Writing scores to warehouse for {len(scores):,} customers...")
    write_scores_to_warehouse(engine, scores)

    return run_id


def run_churn_training(
    engine: Engine,
    mlflow_tracking_uri: str | None = None,
) -> str:
    """Train XGBoost churn classifier, log to MLflow, register as Production."""
    if mlflow_tracking_uri:
        mlflow.set_tracking_uri(mlflow_tracking_uri)
    else:
        mlflow.set_tracking_uri(settings.mlflow_tracking_uri)

    mlflow.set_experiment("churn_prediction")

    print("Computing features for churn model...")
    rfm_df = compute_rfm(engine)
    clv_df = compute_clv_features(engine)
    delivery_df = compute_delivery_features(engine)

    churn_model = ChurnModel()
    feature_df, y = churn_model.build_feature_matrix(rfm_df, clv_df, delivery_df)
    X = feature_df[churn_model.feature_cols].values

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )

    print(f"Training XGBoost churn model on {len(X_train):,} samples...")

    with mlflow.start_run(run_name="xgboost_churn") as run:
        churn_model.fit(X_train, y_train)
        train_metrics = churn_model.evaluate(X_train, y_train)
        test_metrics = churn_model.evaluate(X_test, y_test)

        mlflow.log_params(
            {
                "algorithm": "xgboost",
                "n_estimators": 200,
                "max_depth": 5,
                "learning_rate": 0.05,
                "churn_threshold_days": 180,
                "features": churn_model.feature_cols,
            }
        )
        mlflow.log_metrics({f"train_{k}": v for k, v in train_metrics.items()})
        mlflow.log_metrics({f"test_{k}": v for k, v in test_metrics.items()})
        mlflow.xgboost.log_model(churn_model.model, "churn_model")

        print(
            f"  test AUC-ROC={test_metrics['auc_roc']:.4f}, "
            f"F1={test_metrics['f1']:.4f}, "
            f"churn_rate={test_metrics['churn_rate']:.4f}"
        )

        run_id = run.info.run_id

    register_model(run_id, "churn_classifier", "churn_model")

    # Write churn probabilities back to warehouse for all customers
    all_proba = churn_model.predict_proba(X)
    scores = feature_df[["customer_unique_id"]].copy()
    scores["churn_probability"] = all_proba

    print(f"Writing churn scores to warehouse for {len(scores):,} customers...")
    write_scores_to_warehouse(engine, scores)

    return run_id