"""
Stratified K-fold cross-validation for the stable production pipeline.
"""

from __future__ import annotations

from typing import Any, Callable

import numpy as np
import pandas as pd
from sklearn.metrics import f1_score, roc_auc_score
from sklearn.model_selection import StratifiedKFold

from src.utils.logger import get_logger

logger = get_logger(__name__)


def stratified_kfold_cv(
    X: pd.Series,
    y: pd.Series,
    *,
    n_splits: int,
    random_state: int,
    fit_predict_fn: Callable[[pd.Series, pd.Series, pd.Series, pd.Series], dict[str, float]],
) -> dict[str, Any]:
    """
    Run stratified K-fold CV with a caller-provided fit/eval hook.

    ``fit_predict_fn`` receives (X_tr, y_tr, X_val, y_val) and returns
    per-fold metrics including at least ``f1_weighted``.
    """
    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=random_state)
    fold_metrics: list[dict[str, float]] = []

    for fold, (tr_idx, val_idx) in enumerate(skf.split(X, y)):
        X_tr = X.iloc[tr_idx].reset_index(drop=True)
        y_tr = y.iloc[tr_idx].reset_index(drop=True)
        X_val = X.iloc[val_idx].reset_index(drop=True)
        y_val = y.iloc[val_idx].reset_index(drop=True)

        m = fit_predict_fn(X_tr, y_tr, X_val, y_val)
        m["fold"] = fold
        fold_metrics.append(m)
        logger.info(
            f"CV fold {fold + 1}/{n_splits} — F1={m['f1_weighted']:.4f} "
            f"gap={m.get('train_val_gap', 0):.4f}"
        )

    f1s = [m["f1_weighted"] for m in fold_metrics]
    gaps = [m.get("train_val_gap", 0.0) for m in fold_metrics]
    rocs = [m.get("roc_auc", np.nan) for m in fold_metrics]

    summary = {
        "n_splits": n_splits,
        "f1_mean": round(float(np.mean(f1s)), 4),
        "f1_std": round(float(np.std(f1s)), 4),
        "f1_min": round(float(np.min(f1s)), 4),
        "f1_max": round(float(np.max(f1s)), 4),
        "gap_mean": round(float(np.mean(gaps)), 4),
        "gap_std": round(float(np.std(gaps)), 4),
        "gap_max": round(float(np.max(gaps)), 4),
        "roc_auc_mean": round(float(np.nanmean(rocs)), 4),
        "folds": fold_metrics,
    }
    stable = summary["f1_std"] < 0.05 and summary["gap_max"] < 0.05
    summary["stable_across_folds"] = stable
    return summary


def evaluate_lr_fold(
    lr_model_factory,
    X_tr: pd.Series,
    y_tr: pd.Series,
    X_val: pd.Series,
    y_val: pd.Series,
    *,
    augment_fn=None,
    cfg: dict | None = None,
    seed: int = 42,
) -> dict[str, float]:
    """Fit LR on (optionally augmented) fold train; score fold val."""
    if augment_fn and cfg is not None:
        X_fit, y_fit = augment_fn(X_tr, y_tr, cfg, seed=seed)
    else:
        X_fit, y_fit = X_tr, y_tr

    model = lr_model_factory()
    model.fit(X_fit, y_fit)

    y_val_arr = y_val.astype(int).values
    preds_val = model.predict(X_val)
    preds_train = model.predict(X_fit)
    probs_val = model.predict_proba(X_val)[:, 1]

    f1_val = float(f1_score(y_val_arr, preds_val, average="weighted", zero_division=0))
    f1_train = float(
        f1_score(y_fit.astype(int), preds_train, average="weighted", zero_division=0)
    )
    gap = abs(f1_train - f1_val)

    return {
        "f1_weighted": f1_val,
        "f1_train": f1_train,
        "train_val_gap": round(gap, 4),
        "train_val_gap_pp": round(gap * 100, 2),
        "roc_auc": round(float(roc_auc_score(y_val_arr, probs_val)), 4),
    }