train_eval.py

"""Train/test split, models, metrics, ROC (no leakage: simple random split per person-row)."""
from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import matplotlib.pyplot as plt
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    f1_score,
    precision_score,
    recall_score,
    roc_auc_score,
    roc_curve,
)
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

import config


def _feature_columns(panel: pd.DataFrame, id_col: str, target: str) -> List[str]:
    return [c for c in panel.columns if c not in (id_col, target, "cesd_sum")]


def prepare_xy(
    panel: pd.DataFrame,
    target: str = "y_depressed",
    id_col: str = config.ID_COL,
) -> Tuple[pd.DataFrame, pd.Series, List[str]]:
    feat_cols = _feature_columns(panel, id_col, target)
    if not feat_cols:
        raise ValueError("No feature columns in panel.")
    X = panel[feat_cols]
    y = panel[target]
    mask = y.notna()
    X = X.loc[mask]
    y = y.loc[mask]
    return X, y, feat_cols


def make_logistic_pipeline() -> Pipeline:
    return Pipeline(
        [
            ("imputer", SimpleImputer(strategy="median")),
            ("scaler", StandardScaler()),
            (
                "clf",
                LogisticRegression(
                    max_iter=2000,
                    class_weight="balanced",
                    random_state=config.RANDOM_SEED,
                ),
            ),
        ]
    )


def make_forest_pipeline() -> Pipeline:
    return Pipeline(
        [
            ("imputer", SimpleImputer(strategy="median")),
            (
                "clf",
                RandomForestClassifier(
                    n_estimators=200,
                    max_depth=8,
                    class_weight="balanced",
                    random_state=config.RANDOM_SEED,
                    n_jobs=-1,
                ),
            ),
        ]
    )


def evaluate_models(
    X: pd.DataFrame,
    y: pd.Series,
    test_size: float = 0.2,
    out_dir: Optional[Path] = None,
) -> Dict[str, Any]:
    """Fit logistic regression and random forest; return metrics dict."""
    out_dir = Path(out_dir or config.OUTPUT_DIR)
    out_dir.mkdir(parents=True, exist_ok=True)

    if y.nunique() < 2:
        raise ValueError("Target has a single class; cannot train classifier.")

    X_train, X_test, y_train, y_test = train_test_split(
        X,
        y,
        test_size=test_size,
        random_state=config.RANDOM_SEED,
        stratify=y,
    )

    models = {
        "logistic_regression": make_logistic_pipeline(),
        "random_forest": make_forest_pipeline(),
    }

    results: Dict[str, Any] = {}
    for name, pipe in models.items():
        pipe.fit(X_train, y_train)
        proba = pipe.predict_proba(X_test)[:, 1]
        pred = pipe.predict(X_test)

        results[name] = {
            "accuracy": float(accuracy_score(y_test, pred)),
            "precision": float(precision_score(y_test, pred, zero_division=0)),
            "recall": float(recall_score(y_test, pred, zero_division=0)),
            "f1": float(f1_score(y_test, pred, zero_division=0)),
            "roc_auc": float(roc_auc_score(y_test, proba)),
            "confusion_matrix": confusion_matrix(y_test, pred).tolist(),
            "classification_report": classification_report(
                y_test, pred, zero_division=0
            ),
        }

        # ROC plot
        fpr, tpr, _ = roc_curve(y_test, proba)
        plt.figure(figsize=(6, 5))
        plt.plot(fpr, tpr, label=f"{name} AUC={results[name]['roc_auc']:.3f}")
        plt.plot([0, 1], [0, 1], "k--", alpha=0.4)
        plt.xlabel("False positive rate")
        plt.ylabel("True positive rate")
        plt.title(f"ROC — {name}")
        plt.legend(loc="lower right")
        plt.tight_layout()
        roc_path = out_dir / f"roc_{name}.png"
        plt.savefig(roc_path, dpi=120)
        plt.close()
        results[name]["roc_path"] = str(roc_path)

    return results


def metrics_to_csv(results: Dict[str, Any], path: Path) -> None:
    rows = []
    for name, r in results.items():
        rows.append(
            {
                "model": name,
                "accuracy": r["accuracy"],
                "precision": r["precision"],
                "recall": r["recall"],
                "f1": r["f1"],
                "roc_auc": r["roc_auc"],
            }
        )
    pd.DataFrame(rows).to_csv(path, index=False)