"""K-fold Stratified Cross-Validation utility for binary cyberbullying detection.

Generic over model categories: the caller passes a ``fit_predict_fn`` callback
that does its own featurization (TF-IDF for Trad ML, FastText for Hybrid DL,
raw text for Transformer) and returns predictions + probabilities on the held-out
fold. The utility handles the split, evaluates with ``compute_metrics`` +
``compute_fairness``, and aggregates mean ± std across folds.

Each outer fold's training data is further split into a small validation slice
that the callback can use for early stopping (transformer / hybrid DL only -
Trad ML callbacks just ignore the val argument).

Used by:
- ``notebooks/colab/transformers.ipynb``  (3 transformer backbones × 5 folds)
- ``notebooks/colab/traditional_ml.ipynb`` (3 trad ML × 5 folds)
- ``notebooks/colab/hybrid_dl.ipynb``     (3 hybrid DL × 5 folds)
"""

from __future__ import annotations

import logging
from typing import Callable, Sequence

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

from src.evaluation import compute_fairness, compute_metrics

logger = logging.getLogger(__name__)


FitPredictFn = Callable[
    [list, list, list, list, list],          # X_tr, y_tr, X_va, y_va, X_te
    tuple[np.ndarray, np.ndarray],            # preds, probas (n, 2)
]


def kfold_cv(
    texts: Sequence[str],
    labels: Sequence[int],
    fit_predict_fn: FitPredictFn,
    *,
    n_splits: int = 5,
    seed: int = 42,
    val_size: float = 0.1,
    model_name: str = "model",
) -> dict:
    """Run K-fold stratified CV. Returns a dict with per-fold records + mean/std.

    Parameters
    ----------
    texts, labels
        The full dataset (after preprocessing if any; the callback may further
        transform per fold, but text-level preprocessing should be applied
        before calling).
    fit_predict_fn
        Callable ``(X_tr, y_tr, X_va, y_va, X_te) -> (preds, probas)``. The
        callback owns featurization, model construction, fitting (optionally
        using X_va/y_va for early stopping), and prediction on X_te.
    n_splits
        Number of folds. Default 5 (literature standard for this dataset size).
    seed
        Shuffle seed for ``StratifiedKFold`` + the inner train/val split.
    val_size
        Fraction of each fold's training data to hold out for the validation
        slice. Default 0.1 = ~140 samples per fold of ~1395 train.
    model_name
        Logging tag.

    Returns
    -------
    dict with keys:
        ``n_splits``, ``seed``, ``val_size``, ``per_fold`` (list of dicts -
        each fold's metrics + fairness gaps), ``mean``, ``std`` (scalar
        aggregates).
    """
    X = list(texts)
    y = list(labels)
    if len(X) != len(y):
        raise ValueError(f"texts/labels length mismatch: {len(X)} vs {len(y)}")

    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
    per_fold: list[dict] = []

    for fold_idx, (train_idx, test_idx) in enumerate(skf.split(X, y), start=1):
        X_tr_full = [X[i] for i in train_idx]
        y_tr_full = [y[i] for i in train_idx]
        X_te = [X[i] for i in test_idx]
        y_te = [y[i] for i in test_idx]

        # Inner split for val (for early stopping). Stratify on y_tr_full.
        X_tr, X_va, y_tr, y_va = train_test_split(
            X_tr_full, y_tr_full,
            test_size=val_size,
            stratify=y_tr_full,
            random_state=seed,
        )

        logger.info(
            "fold %d/%d: train=%d val=%d test=%d",
            fold_idx, n_splits, len(X_tr), len(X_va), len(X_te),
        )

        preds, probas = fit_predict_fn(X_tr, y_tr, X_va, y_va, X_te)
        m = compute_metrics(y_te, preds, probas)
        f = compute_fairness(y_te, preds)

        per_fold.append({
            "fold": fold_idx,
            **m,
            "accuracy_gap": f["accuracy_gap"],
            "equal_opportunity_gap": f["equal_opportunity_gap"],
            "demographic_parity_gap": f["demographic_parity_gap"],
            "precision_gap": f["precision_gap"],
            "recall_class_0": f["per_class_accuracy"]["0"],
            "recall_class_1": f["per_class_accuracy"]["1"],
        })

    df_folds = pd.DataFrame(per_fold)
    scalar_cols = [
        c for c in df_folds.columns
        if c != "fold" and pd.api.types.is_numeric_dtype(df_folds[c])
    ]
    mean = {c: float(df_folds[c].mean()) for c in scalar_cols}
    std = {c: float(df_folds[c].std()) for c in scalar_cols}

    return {
        "model": model_name,
        "n_splits": n_splits,
        "seed": seed,
        "val_size": val_size,
        "per_fold": per_fold,
        "mean": mean,
        "std": std,
    }


def aggregate_cv_dataframe(cv_result: dict) -> pd.DataFrame:
    """Convenience: 2-row DataFrame (mean, std) over the scalar metrics."""
    rows = [
        {"agg": "mean", **cv_result["mean"]},
        {"agg": "std",  **cv_result["std"]},
    ]
    return pd.DataFrame(rows)