"""Схемы кросс-валидации и сбор out-of-fold предсказаний."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Iterable, Sequence

import numpy as np
import pandas as pd
from sklearn.base import clone
from sklearn.model_selection import StratifiedKFold

from .evaluation import (
    align_scores,
    classification_metrics,
    metrics_summary,
    summarize_pooled_predictions,
)
from .pipelines import make_grid_search


@dataclass
class CVResult:
    """Результат повторной стратифицированной кросс-валидации."""

    fold_metrics: pd.DataFrame
    oof: pd.DataFrame
    labels: np.ndarray
    scores_by_repeat: dict[int, np.ndarray] = field(default_factory=dict)
    n_splits: int = 5
    n_repeats: int = 1
    model_name: str = ""

    def summary(self) -> pd.DataFrame:
        return metrics_summary(self.fold_metrics)

    def pooled_summary(self) -> tuple[pd.DataFrame, dict[str, float]]:
        return summarize_pooled_predictions(
            self.oof, self.labels, self.scores_by_repeat
        )


@dataclass
class NestedCVResult(CVResult):
    """Результат вложенной кросс-валидации с настройкой гиперпараметров."""

    best_params: list[dict[str, Any]] = field(default_factory=list)


def _predict_scores(pipeline: Any, X: pd.DataFrame) -> np.ndarray | None:
    if hasattr(pipeline, "predict_proba"):
        try:
            return np.asarray(pipeline.predict_proba(X), dtype=float)
        except Exception:
            pass
    if hasattr(pipeline, "decision_function"):
        try:
            scores = np.asarray(pipeline.decision_function(X), dtype=float)
            if scores.ndim == 1:
                scores = np.column_stack([-scores, scores])
            return scores
        except Exception:
            pass
    return None


def _model_classes(pipeline: Any) -> np.ndarray:
    if hasattr(pipeline, "named_steps") and "model" in pipeline.named_steps:
        return np.asarray(pipeline.named_steps["model"].classes_)
    return np.asarray(pipeline.classes_)


def validate_folds(
    y: Sequence,
    splits: Iterable[tuple[np.ndarray, np.ndarray]],
    labels: Sequence,
) -> None:
    """Проверяет, что каждый класс есть и в обучении, и в тесте."""
    y = np.asarray(y)
    labels_set = set(labels)
    for fold_index, (train_idx, test_idx) in enumerate(splits):
        train_labels = set(y[train_idx])
        test_labels = set(y[test_idx])
        missing_train = labels_set - train_labels
        missing_test = labels_set - test_labels
        if missing_train or missing_test:
            raise ValueError(
                f"Фолд {fold_index}: в обучении нет классов {sorted(missing_train)}, "
                f"в тесте нет классов {sorted(missing_test)}. "
                "Увеличьте размер выборки или уменьшите число фолдов."
            )


def _line_ids(X: pd.DataFrame, indices: np.ndarray) -> np.ndarray:
    if "line_id" in X.columns:
        return X.iloc[indices]["line_id"].to_numpy()
    return np.asarray(indices)


def _order_scores_like_oof(X: pd.DataFrame, rows: list[dict[str, Any]]) -> np.ndarray:
    """Возвращает порядок строк матрицы score-ов, соответствующий OOF-таблице."""
    if "line_id" in X.columns:
        positions = {
            int(value): index for index, value in enumerate(X["line_id"].to_numpy())
        }
        return np.asarray([positions[int(row["line_id"])] for row in rows], dtype=int)
    return np.asarray([int(row["line_id"]) for row in rows], dtype=int)


def run_repeated_cv(
    pipeline: Any,
    X: pd.DataFrame,
    y: Sequence,
    n_splits: int = 5,
    n_repeats: int = 5,
    seed: int = 42,
    model_name: str = "",
    verbose: bool = True,
) -> CVResult:
    """Повторная стратифицированная кросс-валидация.

    Все преобразования внутри ``pipeline`` обучаются только на обучающем
    блоке. Out-of-fold предсказания собираются для каждого повтора.
    """
    y = np.asarray(y)
    labels = np.unique(y)
    n_samples = len(y)
    fold_rows: list[dict[str, Any]] = []
    oof_rows: list[dict[str, Any]] = []
    scores_by_repeat: dict[int, np.ndarray] = {}

    for repeat in range(n_repeats):
        splitter = StratifiedKFold(
            n_splits=n_splits, shuffle=True, random_state=seed + repeat
        )
        splits = list(splitter.split(X, y))
        validate_folds(y, splits, labels)
        scores_oof = np.full((n_samples, len(labels)), np.nan, dtype=float)
        oof_start = len(oof_rows)

        for fold, (train_idx, test_idx) in enumerate(splits, start=1):
            model = clone(pipeline)
            model.fit(X.iloc[train_idx], y[train_idx])
            y_pred = model.predict(X.iloc[test_idx])
            raw_scores = _predict_scores(model, X.iloc[test_idx])
            if raw_scores is not None:
                scores = align_scores(
                    raw_scores, _model_classes(model), labels, fill_value=0.0
                )
                scores_oof[test_idx] = scores
            else:
                scores = None

            metrics = classification_metrics(
                y[test_idx], y_pred, y_score=scores, labels=labels
            )
            metrics.update(
                {
                    "repeat": repeat,
                    "fold": fold,
                    "n_train": int(len(train_idx)),
                    "n_test": int(len(test_idx)),
                }
            )
            fold_rows.append(metrics)

            for row_index, sample_index in enumerate(test_idx):
                oof_rows.append(
                    {
                        "repeat": repeat,
                        "fold": fold,
                        "line_id": int(_line_ids(X, np.asarray([sample_index]))[0]),
                        "character": str(X.iloc[sample_index].get("character", "")),
                        "y_true": y[sample_index],
                        "y_pred": y_pred[row_index],
                    }
                )
            if verbose:
                print(
                    f"[cv] повтор {repeat + 1}/{n_repeats}, "
                    f"фолд {fold}/{n_splits}: macro-F1={metrics['macro_f1']:.4f}"
                )

        if not np.isnan(scores_oof).any():
            repeat_rows = sorted(oof_rows[oof_start:], key=lambda row: int(row["line_id"]))
            order = _order_scores_like_oof(X, repeat_rows)
            scores_by_repeat[repeat] = scores_oof[order]

    oof = pd.DataFrame(oof_rows)
    if not oof.empty:
        oof = oof.sort_values(["repeat", "line_id"]).reset_index(drop=True)
    return CVResult(
        fold_metrics=pd.DataFrame(fold_rows),
        oof=oof,
        labels=labels,
        scores_by_repeat=scores_by_repeat,
        n_splits=n_splits,
        n_repeats=n_repeats,
        model_name=model_name,
    )


def run_nested_cv(
    pipeline: Any,
    X: pd.DataFrame,
    y: Sequence,
    model_name: str,
    outer_splits: int = 5,
    inner_splits: int = 3,
    seed: int = 42,
    verbose: bool = True,
) -> NestedCVResult:
    """Вложенная кросс-валидация: настройка гиперпараметров внутри outer-фолдов."""
    y = np.asarray(y)
    labels = np.unique(y)
    n_samples = len(y)
    outer = StratifiedKFold(
        n_splits=outer_splits, shuffle=True, random_state=seed
    )
    splits = list(outer.split(X, y))
    validate_folds(y, splits, labels)

    fold_rows: list[dict[str, Any]] = []
    oof_rows: list[dict[str, Any]] = []
    best_params: list[dict[str, Any]] = []
    scores_oof = np.full((n_samples, len(labels)), np.nan, dtype=float)

    for fold, (train_idx, test_idx) in enumerate(splits, start=1):
        search = make_grid_search(
            pipeline, model_name, cv=inner_splits, scoring="f1_macro", n_jobs=-1
        )
        search.fit(X.iloc[train_idx], y[train_idx])
        model = search.best_estimator_
        y_pred = model.predict(X.iloc[test_idx])
        raw_scores = _predict_scores(model, X.iloc[test_idx])
        if raw_scores is not None:
            scores = align_scores(
                raw_scores, _model_classes(model), labels, fill_value=0.0
            )
            scores_oof[test_idx] = scores
        else:
            scores = None

        metrics = classification_metrics(
            y[test_idx], y_pred, y_score=scores, labels=labels
        )
        metrics.update(
            {
                "repeat": 0,
                "fold": fold,
                "n_train": int(len(train_idx)),
                "n_test": int(len(test_idx)),
            }
        )
        fold_rows.append(metrics)
        best_params.append(dict(search.best_params_))
        for row_index, sample_index in enumerate(test_idx):
            oof_rows.append(
                {
                    "repeat": 0,
                    "fold": fold,
                    "line_id": int(_line_ids(X, np.asarray([sample_index]))[0]),
                    "character": str(X.iloc[sample_index].get("character", "")),
                    "y_true": y[sample_index],
                    "y_pred": y_pred[row_index],
                }
            )
        if verbose:
            print(
                f"[nested-cv] фолд {fold}/{outer_splits}: "
                f"macro-F1={metrics['macro_f1']:.4f}, "
                f"best={search.best_params_}"
            )

    scores_by_repeat = {}
    if not np.isnan(scores_oof).any():
        ordered_rows = sorted(oof_rows, key=lambda row: int(row["line_id"]))
        order = _order_scores_like_oof(X, ordered_rows)
        scores_by_repeat[0] = scores_oof[order]

    oof = pd.DataFrame(oof_rows)
    if not oof.empty:
        oof = oof.sort_values(["repeat", "line_id"]).reset_index(drop=True)
    return NestedCVResult(
        fold_metrics=pd.DataFrame(fold_rows),
        oof=oof,
        labels=labels,
        scores_by_repeat=scores_by_repeat,
        n_splits=outer_splits,
        n_repeats=1,
        model_name=model_name,
        best_params=best_params,
    )