NLP/lab1/stylometry/plotting.py

"""Построение графиков для отчёта и презентации."""

from __future__ import annotations

from pathlib import Path
from typing import Sequence

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402
from sklearn.decomposition import PCA  # noqa: E402
from sklearn.manifold import TSNE  # noqa: E402


def _save(fig: plt.Figure, path: str | Path) -> Path:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.tight_layout()
    fig.savefig(path, dpi=200, bbox_inches="tight")
    plt.close(fig)
    return path


def plot_class_distribution(
    frame: pd.DataFrame,
    character_column: str = "character",
    path: str | Path = "results/figures/class_distribution.png",
) -> Path:
    counts = frame[character_column].value_counts().sort_values()
    fig, ax = plt.subplots(figsize=(8, max(6, len(counts) * 0.25)))
    ax.barh(counts.index, counts.values)
    ax.set_xlabel("Число реплик")
    ax.set_ylabel("Персонаж")
    ax.set_title("Распределение реплик по персонажам")
    return _save(fig, path)


def plot_scree(
    explained_variance_ratio: Sequence[float],
    path: str | Path = "results/figures/lsa_scree.png",
    title: str = "Доля объяснённой дисперсии компонент LSA",
) -> Path:
    values = np.asarray(explained_variance_ratio, dtype=float)
    cumulative = np.cumsum(values)
    fig, ax = plt.subplots(figsize=(7, 4.5))
    ax.plot(np.arange(1, len(values) + 1), cumulative, marker="o", markersize=3)
    ax.set_xlabel("Число компонент")
    ax.set_ylabel("Накопленная доля дисперсии")
    ax.set_title(title)
    ax.grid(alpha=0.3)
    return _save(fig, path)


def plot_confusion_matrix(
    matrix: pd.DataFrame,
    path: str | Path = "results/figures/confusion_matrix.png",
    title: str = "Матрица ошибок",
    annotate: bool | None = None,
) -> Path:
    values = matrix.to_numpy(dtype=float)
    n_classes = values.shape[0]
    if annotate is None:
        annotate = n_classes <= 20
    fig, ax = plt.subplots(figsize=(max(8, n_classes * 0.35), max(7, n_classes * 0.32)))
    image = ax.imshow(values, cmap="Blues", aspect="auto")
    ax.set_xticks(np.arange(n_classes))
    ax.set_yticks(np.arange(n_classes))
    ax.set_xticklabels(matrix.columns, rotation=90, fontsize=7)
    ax.set_yticklabels(matrix.index, fontsize=7)
    ax.set_xlabel("Предсказанный класс")
    ax.set_ylabel("Истинный класс")
    ax.set_title(title)
    if annotate:
        for row in range(n_classes):
            for column in range(n_classes):
                value = values[row, column]
                if value:
                    ax.text(
                        column, row, f"{value:.2f}" if value < 1 else f"{int(value)}",
                        ha="center", va="center", fontsize=6,
                        color="white" if value > values.max() * 0.5 else "black",
                    )
    fig.colorbar(image, ax=ax, fraction=0.046, pad=0.04)
    return _save(fig, path)


def plot_pca(
    matrix: np.ndarray,
    labels: Sequence[str],
    path: str | Path = "results/figures/pca.png",
    title: str = "PCA признакового пространства",
) -> Path:
    values = np.asarray(matrix, dtype=float)
    if values.shape[0] < 2:
        raise ValueError("Для PCA нужно хотя бы два объекта.")
    coordinates = PCA(n_components=2, random_state=42).fit_transform(values)
    fig, ax = plt.subplots(figsize=(9, 7))
    unique_labels = sorted(set(labels))
    cmap = plt.get_cmap("tab20")
    color_map = {label: cmap(index % 20) for index, label in enumerate(unique_labels)}
    for label in unique_labels:
        mask = np.asarray(labels) == label
        ax.scatter(
            coordinates[mask, 0], coordinates[mask, 1],
            s=12, alpha=0.7, label=label, color=color_map[label],
        )
    ax.set_xlabel("Первая главная компонента")
    ax.set_ylabel("Вторая главная компонента")
    ax.set_title(title)
    ax.legend(fontsize=6, ncol=2, loc="best", frameon=False)
    return _save(fig, path)


def plot_tsne(
    matrix: np.ndarray,
    labels: Sequence[str],
    path: str | Path = "results/figures/tsne.png",
    title: str = "t-SNE признакового пространства",
    perplexity: float = 30.0,
) -> Path:
    values = np.asarray(matrix, dtype=float)
    if values.shape[0] < 10:
        raise ValueError("Для t-SNE нужно хотя бы десять объектов.")
    perplexity = min(perplexity, max(5.0, values.shape[0] / 4.0))
    coordinates = TSNE(
        n_components=2,
        perplexity=perplexity,
        init="pca",
        learning_rate="auto",
        random_state=42,
    ).fit_transform(values)
    fig, ax = plt.subplots(figsize=(9, 7))
    unique_labels = sorted(set(labels))
    cmap = plt.get_cmap("tab20")
    color_map = {label: cmap(index % 20) for index, label in enumerate(unique_labels)}
    for label in unique_labels:
        mask = np.asarray(labels) == label
        ax.scatter(
            coordinates[mask, 0], coordinates[mask, 1],
            s=10, alpha=0.7, label=label, color=color_map[label],
        )
    ax.set_xlabel("t-SNE 1")
    ax.set_ylabel("t-SNE 2")
    ax.set_title(title)
    ax.legend(fontsize=6, ncol=2, loc="best", frameon=False)
    return _save(fig, path)


def plot_feature_importance(
    names: Sequence[str],
    values: Sequence[float],
    path: str | Path = "results/figures/feature_importance.png",
    title: str = "Важность признаков",
    top_n: int = 25,
) -> Path:
    frame = pd.DataFrame({"feature": list(names), "value": np.asarray(values, dtype=float)})
    frame["abs_value"] = frame["value"].abs()
    frame = frame.sort_values("abs_value", ascending=False).head(top_n).iloc[::-1]
    fig, ax = plt.subplots(figsize=(9, max(5, len(frame) * 0.28)))
    ax.barh(frame["feature"], frame["value"])
    ax.set_xlabel("Вклад признака")
    ax.set_ylabel("Признак")
    ax.set_title(title)
    return _save(fig, path)