"""
Minimal script: run the DL 5-fold CV exactly as train_deep_classifiers.py
does, but also save the raw out-of-fold probability array per model to
models/dl_oof_probs.npz, so ensemble_model.py can build a real ensemble
without retraining DL models a second time.

Usage:
    python -m app.training.dump_dl_oof
"""

from __future__ import annotations

import sys
import time
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from sklearn.metrics import roc_auc_score
from sklearn.model_selection import StratifiedKFold

from app.training.train_deep_classifiers import (
    DeepMLP, Conv1DClassifier, ResidualMLP, AttentionMLP,
    load_data, train_one_fold, set_seed, SEED, N_FOLDS, DEVICE,
)

MODELS_DIR = Path(__file__).resolve().parents[2] / "models"
FEATURES_CSV = Path("D:/CrownCode/DataSet/features.csv")

MODEL_CLASSES = {
    "Deep MLP (512-256-128-64)": DeepMLP,
    "1D-CNN": Conv1DClassifier,
    "Residual MLP (3 blocks)": ResidualMLP,
    "Attention MLP": AttentionMLP,
}


def main() -> None:
    print(f"Device: {DEVICE}", flush=True)
    X, y, feature_cols = load_data(FEATURES_CSV)
    print(f"Samples: {len(y)}, Features: {X.shape[1]}", flush=True)

    cv = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=SEED)
    fold_assignments = list(cv.split(X, y))

    oof: dict[str, np.ndarray] = {}
    for name, cls in MODEL_CLASSES.items():
        print(f"\n{'='*60}\n  {name}\n{'='*60}", flush=True)
        t0 = time.time()
        all_probs = np.zeros(len(y))
        for fold, (train_idx, val_idx) in enumerate(fold_assignments):
            set_seed(SEED + fold)
            model = cls(X.shape[1])
            fold_t0 = time.time()
            auc, probs = train_one_fold(model, X[train_idx], y[train_idx], X[val_idx], y[val_idx])
            all_probs[val_idx] = probs
            print(f"    Fold {fold+1}/{N_FOLDS}: AUC={auc:.4f} ({time.time()-fold_t0:.1f}s)", flush=True)
        oof[name] = all_probs
        total_auc = roc_auc_score(y, all_probs)
        print(f"  => Overall OOF AUC={total_auc:.4f}  ({time.time()-t0:.1f}s total)", flush=True)

    out_path = MODELS_DIR / "dl_oof_probs.npz"
    np.savez(out_path, y=y, **{name.replace(" ", "_").replace("(", "").replace(")", "").replace("-", "_"): probs
                                for name, probs in oof.items()})
    print(f"\nSaved: {out_path}", flush=True)
    print("Keys:", list(np.load(out_path).keys()))


if __name__ == "__main__":
    main()