run_pipeline.py

"""
End-to-end pipeline: load → features → train → evaluate → explain → save.

Run with: python run_pipeline.py
"""
import logging

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

from src.config import CHARTS_DIR, METRICS_DIR, RANDOM_STATE, TEST_SIZE
from src.data_pipeline import encode_features, load_raw_data, run_quality_checks
from src.evaluation import (
    build_comparison_table,
    evaluate_model,
    plot_cost_vs_threshold,
    plot_roc_curves,
)
from src.explainability import (
    compute_shap_values,
    get_feature_importances,
    plot_feature_importance,
    segment_churn_analysis,
)
from src.feature_engineering import (
    SEGMENTATION_ONLY_FEATURES,
    add_behavioral_features,
)
from src.modeling import train_all_models
from src.utils import ensure_dirs, save_csv, save_json, setup_logging


def main():
    setup_logging()
    logger = logging.getLogger(__name__)
    ensure_dirs(CHARTS_DIR, METRICS_DIR)

    # ── 1. Data ──────────────────────────────────────────────────────────────
    logger.info("=== PHASE 1: Data Pipeline ===")
    raw_df = load_raw_data()
    quality_report = run_quality_checks(raw_df)
    save_json(quality_report, METRICS_DIR / "data_quality_report.json")

    # Apply feature engineering before encoding (stateless transforms only)
    raw_df_fe = add_behavioral_features(raw_df)

    # Encode (binary plans, drop charge cols, drop state/area)
    encoded = encode_features(raw_df_fe)

    # Separate segmentation-only features before split
    X_full = encoded.drop(columns=["Churn"])
    y = encoded["Churn"]

    # Model-safe feature columns (exclude segmentation-only)
    model_cols = [c for c in X_full.columns if c not in SEGMENTATION_ONLY_FEATURES]
    X_model = X_full[model_cols]

    X_train, X_test, y_train, y_test = train_test_split(
        X_model, y, test_size=TEST_SIZE, random_state=RANDOM_STATE, stratify=y
    )
    logger.info(
        "Train: %d | Test: %d | Churn rate: %.1f%%",
        len(X_train), len(X_test), y_train.mean() * 100,
    )

    # ── 2. Training ───────────────────────────────────────────────────────────
    logger.info("=== PHASE 2: Model Training ===")
    trained_models = train_all_models(X_train, y_train)

    # ── 3. Evaluation ─────────────────────────────────────────────────────────
    logger.info("=== PHASE 3: Evaluation ===")
    all_results = {}
    for name, result in trained_models.items():
        model = result["model"]
        metrics = evaluate_model(model, X_test, y_test)
        all_results[name] = {
            "metrics": metrics,
            "best_params": result["best_params"],
            "purpose": result["purpose"],
            "tradeoffs": result["tradeoffs"],
        }
        plot_cost_vs_threshold(model, X_test, y_test, name)

    comparison_df = build_comparison_table(all_results)
    save_csv(
        comparison_df.drop(columns=["Purpose", "Trade-offs"]),
        METRICS_DIR / "model_comparison.csv",
    )
    plot_roc_curves(all_results, y_test)

    # Save metrics JSON (strip y_prob arrays)
    metrics_export = {}
    for name, r in all_results.items():
        metrics_export[name] = {
            k: v for k, v in r["metrics"].items() if k != "y_prob"
        }
        metrics_export[name]["best_params"] = r["best_params"]
    save_json(metrics_export, METRICS_DIR / "all_model_metrics.json")

    logger.info(
        "\n%s",
        comparison_df[[
            "Model", "Precision", "Recall", "F1",
            "ROC-AUC", "PR-AUC", "Business Cost ($)", "Churn Prevented/1000",
        ]].to_string(index=False),
    )

    # ── 4. Explainability (best model by PR-AUC) ──────────────────────────────
    logger.info("=== PHASE 4: Explainability ===")
    best_name = max(all_results, key=lambda n: all_results[n]["metrics"]["pr_auc"])
    best_model = trained_models[best_name]["model"]
    logger.info("Best model by PR-AUC: %s", best_name)

    feature_names = list(X_train.columns)
    importance_df = get_feature_importances(best_model, feature_names)
    plot_feature_importance(importance_df, best_name)
    save_csv(importance_df, METRICS_DIR / f"feature_importance_{best_name}.csv")

    shap_df = compute_shap_values(best_model, X_train, X_test)
    if shap_df is not None:
        save_csv(shap_df, METRICS_DIR / "shap_importance.csv")

    # Segment analysis: add back segmentation-only features for the test indices
    X_test_seg = X_test.copy()
    for col in ["high_service_caller", "intl_plan_underuser"]:
        if col in X_full.columns:
            X_test_seg[col] = X_full.loc[X_test.index, col].values

    y_prob = np.array(all_results[best_name]["metrics"]["y_prob"])
    segment_df = segment_churn_analysis(X_test_seg, y_test, y_prob)
    save_csv(segment_df, METRICS_DIR / "segment_churn_analysis.csv")
    logger.info("\nSegment Analysis:\n%s", segment_df.to_string(index=False))

    logger.info("=== PIPELINE COMPLETE ===")
    logger.info("Charts saved to: %s", CHARTS_DIR)
    logger.info("Metrics saved to: %s", METRICS_DIR)


if __name__ == "__main__":
    main()