model/train.py

"""End-to-end training pipeline (CLASS-05, D-REPRO-04, Pattern 6).

Run: `make train` (delegates to `uv run python -m model.train`).
Produces: artifacts/{classifier.joblib, anomaly.joblib, eval_metrics.json}
          artifacts/plots/{confusion_matrix, reliability_raw,
                            reliability_calibrated, lead_time_cdf}.png

NaN policy (canonical, deferred from 02-04 deviation #3):
  - Classifier path: LightGBM handles NaN natively — no preprocessing.
  - Anomaly TRAINING: filter rows with any NaN feature (no window-layout
    constraint at training time; preserves IForest fit semantics).
  - Anomaly EVAL / score_lead_times: MUST preserve window-major layout
    (Pitfall 11); impute NaN to 0.0 in-place (zero-fill). With ~20% NaN
    concentrated in the disconnect-event windows where operational nulls
    are emitted, zero-fill keeps the IForest score signal directionally
    correct (zero values are typically "less anomalous" against a healthy
    baseline so missed-window detections are conservative — surfaced via
    per_class_miss_rate).
"""
from __future__ import annotations

import json
from pathlib import Path

import joblib
import numpy as np

from model.eval import build_eval_metrics
from model.features import (
    CLASSES,
    load_anomaly_features,
    load_split,
)
from model.inference import apply_mask_and_renormalize
from model.normal_split import generate_normal_split
from model.plots import (
    plot_confusion_matrix,
    plot_lead_time_cdfs,
    plot_reliability_grid,
)
from model.seeds import phase2_seeds
from model.train_anomaly import (
    calibrate_threshold,
    per_class_miss_rate,
    score_lead_times,
    train_anomaly,
)
from model.train_classifier import (
    train_calibrated_classifier,
    train_raw_classifier,
)

SCHEMA_VERSION: str = "1.0.0"


def _impute_nan_zero(X: np.ndarray) -> np.ndarray:
    """Zero-fill NaN entries in-place-safe; preserves shape (Pitfall 11 layout)."""
    X = X.copy()
    X[~np.isfinite(X)] = 0.0
    return X


def main() -> None:
    artifacts = Path("artifacts")
    plots_dir = artifacts / "plots"
    plots_dir.mkdir(parents=True, exist_ok=True)

    seeds = phase2_seeds()  # D-REPRO-02

    # ----- 1. Train calibrated classifier -----
    # LightGBM handles NaN natively (Pitfall: don't preprocess; let LGBM split on missing).
    X_train, y_train, _ = load_split(Path("data/train.parquet"))
    print(f"[train] classifier: {len(X_train)} rows × {X_train.shape[1]} features")
    classifier = train_calibrated_classifier(
        X_train, y_train,
        classifier_seed=seeds["classifier_train"],
        cv_seed=seeds["classifier_cv"],
    )
    joblib.dump(classifier, artifacts / "classifier.joblib", compress=3)

    # ----- 2. Raw (non-calibrated) baseline for the dual reliability grid -----
    # CalibratedClassifierCV does not expose the underlying full-train softmax;
    # cheapest route is a parallel LGBMClassifier fit on the same train + seed.
    raw_clf = train_raw_classifier(
        X_train, y_train, classifier_seed=seeds["classifier_train"]
    )

    # ----- 3. Anomaly detector + threshold calibration on normal-only split -----
    X_anom_train, _, _ = load_anomaly_features(Path("data/train.parquet"))
    # Filter NaN rows for IForest training (no window-layout constraint here).
    finite_mask = np.isfinite(X_anom_train).all(axis=1)
    X_anom_train_clean = X_anom_train[finite_mask]
    print(
        f"[train] anomaly: {len(X_anom_train_clean)} of {len(X_anom_train)} "
        f"finite rows ({100.0 * finite_mask.mean():.1f}%)"
    )
    iforest = train_anomaly(X_anom_train_clean, random_state=seeds["anomaly_train"])

    normal_path = Path("data/normal.parquet")
    if not normal_path.exists():
        generate_normal_split(seeds["normal_split_synth"], normal_path)
    X_normal, _, _ = load_anomaly_features(normal_path)
    # Normal split is healthy-shape and should have no NaN (asserted by 02-04 tests).
    threshold = calibrate_threshold(iforest, X_normal, percentile=95.0)

    joblib.dump(
        {"detector": iforest, "threshold": threshold},
        artifacts / "anomaly.joblib", compress=3,
    )

    # ----- 4. Eval the classifier -----
    X_eval, y_eval, _ = load_split(Path("data/eval.parquet"))
    calibrated_proba = classifier.predict_proba(X_eval)
    raw_proba = raw_clf.predict_proba(X_eval)

    # ----- 4a. Apply per-row mask-and-renormalize using each row's network_mode -----
    # We need the original `network_mode` column from the eval Parquet; load it
    # alongside features for the by-mode F1 + the post-mask argmax.
    import pyarrow.parquet as pq

    eval_tbl = pq.read_table(Path("data/eval.parquet"), columns=["network_mode"])
    network_mode_per_row = np.array(eval_tbl["network_mode"].to_pylist())

    n_eval = len(X_eval)
    masked_proba = np.zeros_like(calibrated_proba)
    for i in range(n_eval):
        masked_proba[i] = apply_mask_and_renormalize(
            calibrated_proba[i], network_mode_per_row[i]
        )
    y_pred_after_mask = np.argmax(masked_proba, axis=1)

    # ----- 5. Eval the anomaly detector — per-class lead-time CDFs -----
    # Window-major layout REQUIRED (Pitfall 11) — zero-fill NaN, do NOT row-filter.
    X_anom_eval, y_anom_eval, ts_eval = load_anomaly_features(Path("data/eval.parquet"))
    X_anom_eval = _impute_nan_zero(X_anom_eval)
    lead_times = score_lead_times(
        iforest, X_anom_eval, y_anom_eval, ts_eval, threshold
    )
    miss_rates = per_class_miss_rate(iforest, X_anom_eval, y_anom_eval, threshold)

    # ----- 6. Plots (D-CAL-06 dual grid; D-ANOM-03 CDFs) -----
    plot_confusion_matrix(
        y_eval, y_pred_after_mask, str(plots_dir / "confusion_matrix.png")
    )
    plot_reliability_grid(
        y_eval, raw_proba, str(plots_dir / "reliability_raw.png"),
        "Reliability — raw LightGBM softmax",
    )
    plot_reliability_grid(
        y_eval, calibrated_proba, str(plots_dir / "reliability_calibrated.png"),
        "Reliability — isotonic-calibrated",
    )
    plot_lead_time_cdfs(lead_times, str(plots_dir / "lead_time_cdf.png"))

    # ----- 7. Build + write eval_metrics.json (Pattern 11) -----
    metrics = build_eval_metrics(
        y_eval=y_eval,
        calibrated_proba=calibrated_proba,
        y_pred_after_mask=y_pred_after_mask,
        network_mode_per_row=network_mode_per_row,
        anomaly_threshold=threshold,
        per_class_lead_times=lead_times,
        per_class_miss_rates=miss_rates,
        schema_version=SCHEMA_VERSION,
    )
    (artifacts / "eval_metrics.json").write_text(
        json.dumps(metrics, indent=2, sort_keys=True)
    )
    print(f"wrote {artifacts / 'eval_metrics.json'}")
    print(f"  CLASSES = {CLASSES}")
    print(f"  macro_f1                       = {metrics['macro_f1']:.4f}")
    print(f"  ece_mean                       = {metrics['ece_mean']:.4f}")
    print(
        f"  lead_time_aggregate_median_s   = "
        f"{metrics['anomaly']['lead_time_aggregate_median_s']:.2f}"
    )


if __name__ == "__main__":
    main()