NLP/lab1/stylometry/evaluation.py

"""Метрики качества классификации и их агрегация."""

from __future__ import annotations

from typing import Iterable, Sequence

import numpy as np
import pandas as pd
from sklearn.metrics import (
    accuracy_score,
    balanced_accuracy_score,
    cohen_kappa_score,
    confusion_matrix,
    f1_score,
    log_loss,
    matthews_corrcoef,
    precision_recall_fscore_support,
    top_k_accuracy_score,
)


def _labels_or_default(y_true: Sequence, y_pred: Sequence, labels: Sequence | None) -> np.ndarray:
    if labels is not None:
        return np.asarray(labels)
    return np.unique(np.concatenate([np.asarray(y_true), np.asarray(y_pred)]))


def brier_score_multiclass(y_true: Sequence, y_score: np.ndarray, labels: Sequence) -> float:
    """Многоклассовая оценка Брайера: средний квадрат отклонения вероятностей."""
    label_to_index = {label: index for index, label in enumerate(labels)}
    y_true = np.asarray(y_true)
    one_hot = np.zeros_like(y_score, dtype=float)
    for row, label in enumerate(y_true):
        one_hot[row, label_to_index[label]] = 1.0
    return float(np.mean(np.sum((y_score - one_hot) ** 2, axis=1)))


def align_scores(
    y_score: np.ndarray,
    model_classes: Sequence,
    labels: Sequence,
    fill_value: float = 0.0,
) -> np.ndarray:
    """Приводит столбцы вероятностей/решений к единому порядку классов."""
    labels = list(labels)
    model_classes = list(model_classes)
    result = np.full((y_score.shape[0], len(labels)), fill_value, dtype=float)
    for source_index, label in enumerate(model_classes):
        if label in labels:
            result[:, labels.index(label)] = y_score[:, source_index]
    return result


def classification_metrics(
    y_true: Sequence,
    y_pred: Sequence,
    y_score: np.ndarray | None = None,
    labels: Sequence | None = None,
    top_k: Iterable[int] = (3, 5),
) -> dict[str, float]:
    """Считает набор метрик для одного разбиения."""
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)
    labels_array = _labels_or_default(y_true, y_pred, labels)
    n_classes = len(labels_array)

    result: dict[str, float] = {
        "n_samples": int(len(y_true)),
        "n_classes": int(n_classes),
        "accuracy": float(accuracy_score(y_true, y_pred)),
        "balanced_accuracy": float(
            balanced_accuracy_score(y_true, y_pred)
        ),
        "macro_f1": float(
            f1_score(y_true, y_pred, labels=labels_array, average="macro", zero_division=0)
        ),
        "weighted_f1": float(
            f1_score(y_true, y_pred, labels=labels_array, average="weighted", zero_division=0)
        ),
        "micro_f1": float(
            f1_score(y_true, y_pred, labels=labels_array, average="micro", zero_division=0)
        ),
        "mcc": float(matthews_corrcoef(y_true, y_pred)),
        "cohen_kappa": float(cohen_kappa_score(y_true, y_pred, labels=labels_array)),
    }

    if y_score is not None:
        y_score = np.asarray(y_score, dtype=float)
        if y_score.ndim == 2 and y_score.shape[1] == n_classes:
            for k in top_k:
                if 1 < k < n_classes:
                    result[f"top{k}_accuracy"] = float(
                        top_k_accuracy_score(y_true, y_score, k=k, labels=labels_array)
                    )
            row_sums = y_score.sum(axis=1)
            if np.all(y_score >= 0) and np.allclose(row_sums, 1.0, atol=1e-5):
                result["log_loss"] = float(log_loss(y_true, y_score, labels=labels_array))
                result["brier_score"] = brier_score_multiclass(y_true, y_score, labels_array)

    return result


def per_class_report(
    y_true: Sequence,
    y_pred: Sequence,
    labels: Sequence,
) -> pd.DataFrame:
    """Точность, полнота и F-мера по каждому классу."""
    precision, recall, f1, support = precision_recall_fscore_support(
        y_true, y_pred, labels=labels, zero_division=0
    )
    return pd.DataFrame(
        {
            "class": list(labels),
            "precision": precision,
            "recall": recall,
            "f1": f1,
            "support": support,
        }
    )


def confusion_matrix_frame(
    y_true: Sequence,
    y_pred: Sequence,
    labels: Sequence,
    normalize: str | None = None,
) -> pd.DataFrame:
    """Матрица ошибок в виде таблицы с подписями классов."""
    matrix = confusion_matrix(y_true, y_pred, labels=labels, normalize=normalize)
    return pd.DataFrame(matrix, index=list(labels), columns=list(labels))


def top_confusions(
    y_true: Sequence,
    y_pred: Sequence,
    labels: Sequence,
    top_n: int = 30,
) -> pd.DataFrame:
    """Самые частые пары «истинный класс — предсказанный класс»."""
    frame = confusion_matrix_frame(y_true, y_pred, labels)
    rows = []
    for true_label in frame.index:
        for predicted_label in frame.columns:
            if true_label == predicted_label:
                continue
            count = int(frame.loc[true_label, predicted_label])
            if count > 0:
                rows.append(
                    {
                        "true_class": true_label,
                        "predicted_class": predicted_label,
                        "count": count,
                        "share_of_true_class": count / max(int(frame.loc[true_label].sum()), 1),
                    }
                )
    result = pd.DataFrame(rows)
    if result.empty:
        return result
    return result.sort_values(["count", "share_of_true_class"], ascending=False).head(top_n)


def bootstrap_ci(
    y_true: Sequence,
    y_pred: Sequence,
    labels: Sequence,
    metric: str = "macro_f1",
    n_boot: int = 1000,
    seed: int = 42,
) -> dict[str, float]:
    """Перцентильный доверительный интервал метрики по репликам."""
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)
    labels_array = np.asarray(labels)
    rng = np.random.default_rng(seed)

    def _metric(true_values: np.ndarray, predicted_values: np.ndarray) -> float:
        if metric == "macro_f1":
            return float(
                f1_score(
                    true_values,
                    predicted_values,
                    labels=labels_array,
                    average="macro",
                    zero_division=0,
                )
            )
        if metric == "accuracy":
            return float(accuracy_score(true_values, predicted_values))
        if metric == "balanced_accuracy":
            return float(balanced_accuracy_score(true_values, predicted_values))
        if metric == "mcc":
            return float(matthews_corrcoef(true_values, predicted_values))
        raise ValueError(f"Неизвестная метрика для bootstrap: {metric}")

    values = []
    n_samples = len(y_true)
    for _ in range(n_boot):
        indices = rng.integers(0, n_samples, n_samples)
        values.append(_metric(y_true[indices], y_pred[indices]))
    values_array = np.asarray(values, dtype=float)
    return {
        "mean": float(np.mean(values_array)),
        "lower": float(np.percentile(values_array, 2.5)),
        "upper": float(np.percentile(values_array, 97.5)),
        "n_boot": int(n_boot),
    }


def metrics_summary(
    metrics: pd.DataFrame,
    columns: Iterable[str] | None = None,
) -> pd.DataFrame:
    """Среднее и стандартное отклонение метрик по фолдам/повторам."""
    if columns is None:
        columns = [
            column
            for column in metrics.columns
            if column not in {"repeat", "fold", "n_train", "n_test", "n_samples", "n_classes"}
        ]
    rows = []
    for column in columns:
        if column not in metrics.columns:
            continue
        values = pd.to_numeric(metrics[column], errors="coerce").dropna()
        if values.empty:
            continue
        rows.append(
            {
                "metric": column,
                "mean": float(values.mean()),
                "std": float(values.std(ddof=1)) if len(values) > 1 else 0.0,
                "min": float(values.min()),
                "max": float(values.max()),
                "n": int(len(values)),
            }
        )
    return pd.DataFrame(rows)


def summarize_pooled_predictions(
    oof: pd.DataFrame,
    labels: Sequence,
    scores_by_repeat: dict[int, np.ndarray] | None = None,
    top_k: Iterable[int] = (3, 5),
) -> tuple[pd.DataFrame, dict[str, float]]:
    """Считает метрики по pooled out-of-fold предсказаниям.

    Возвращает таблицу метрик по повторам и итоговую bootstrap-оценку по
    первому повтору (для доверительного интервала).
    """
    rows = []
    for repeat, group in oof.groupby("repeat"):
        scores = None
        if scores_by_repeat is not None and repeat in scores_by_repeat:
            scores = scores_by_repeat[repeat]
        metrics = classification_metrics(
            group["y_true"].values,
            group["y_pred"].values,
            y_score=scores,
            labels=labels,
            top_k=top_k,
        )
        metrics["repeat"] = int(repeat)
        rows.append(metrics)
    per_repeat = pd.DataFrame(rows)

    first_repeat = oof[oof["repeat"] == oof["repeat"].min()]
    bootstrap = {
        metric: bootstrap_ci(first_repeat["y_true"].values, first_repeat["y_pred"].values, labels, metric)
        for metric in ("macro_f1", "accuracy", "balanced_accuracy", "mcc")
    }
    return per_repeat, bootstrap