"""Классическая дельта Барроуза и её косинусный вариант.

Дельта Барроуза — стандартный метод стилометрии. Он использует частоты
наиболее частотных слов, z-нормирование по обучающей выборке и расстояние до
центроида класса. Мы реализуем два варианта: манхэттенское расстояние
(классическая дельта) и косинусное расстояние (Cosine Delta).
"""

from __future__ import annotations

import re
from collections import Counter
from typing import Any, Sequence

import numpy as np
import pandas as pd
from scipy.special import softmax
from scipy.spatial.distance import cdist
from sklearn.base import BaseEstimator, ClassifierMixin, TransformerMixin

_TOKEN_RE = re.compile(r"[A-Za-z]+(?:'[A-Za-z]+)?")


def _texts(X: Any) -> list[str]:
    if isinstance(X, pd.DataFrame):
        if "text" not in X.columns:
            raise KeyError("Для Delta нужна колонка 'text'.")
        return X["text"].astype(str).tolist()
    return [str(value) for value in X]


def _tokens(text: str, lowercase: bool = True) -> list[str]:
    tokens = _TOKEN_RE.findall(text)
    return [token.lower() for token in tokens] if lowercase else tokens


class MFWVectorizer(BaseEstimator, TransformerMixin):
    """Векторизатор по наиболее частотным словам (Most Frequent Words)."""

    def __init__(
        self,
        n_features: int = 100,
        lowercase: bool = True,
        zscore: bool = True,
    ) -> None:
        self.n_features = n_features
        self.lowercase = lowercase
        self.zscore = zscore

    def fit(self, X: Any, y=None) -> "MFWVectorizer":
        texts = _texts(X)
        counter: Counter[str] = Counter()
        for text in texts:
            counter.update(_tokens(text, self.lowercase))
        self.vocabulary_ = [word for word, _ in counter.most_common(self.n_features)]
        if not self.vocabulary_:
            raise ValueError("Не удалось выделить слова для Delta.")
        index = {word: position for position, word in enumerate(self.vocabulary_)}
        raw = np.zeros((len(texts), len(self.vocabulary_)), dtype=float)
        for row, text in enumerate(texts):
            tokens = _tokens(text, self.lowercase)
            counts = Counter(tokens)
            total = max(len(tokens), 1)
            for word, count in counts.items():
                position = index.get(word)
                if position is not None:
                    raw[row, position] = count / total
        self.mean_ = raw.mean(axis=0)
        std = raw.std(axis=0, ddof=0)
        self.std_ = np.where(std == 0, 1.0, std)
        return self

    def transform(self, X: Any) -> np.ndarray:
        texts = _texts(X)
        index = {word: position for position, word in enumerate(self.vocabulary_)}
        raw = np.zeros((len(texts), len(self.vocabulary_)), dtype=float)
        for row, text in enumerate(texts):
            tokens = _tokens(text, self.lowercase)
            counts = Counter(tokens)
            total = max(len(tokens), 1)
            for word, count in counts.items():
                position = index.get(word)
                if position is not None:
                    raw[row, position] = count / total
        if self.zscore:
            return (raw - self.mean_) / self.std_
        return raw

    def get_feature_names_out(self, input_features=None) -> np.ndarray:
        return np.asarray(self.vocabulary_, dtype=object)


class DeltaClassifier(BaseEstimator, ClassifierMixin):
    """Классификатор на основе расстояния до центроида класса."""

    def __init__(
        self,
        n_features: int = 100,
        metric: str = "manhattan",
        zscore: bool = True,
        lowercase: bool = True,
    ) -> None:
        self.n_features = n_features
        self.metric = metric
        self.zscore = zscore
        self.lowercase = lowercase

    def _make_vectorizer(self) -> MFWVectorizer:
        return MFWVectorizer(
            n_features=self.n_features,
            lowercase=self.lowercase,
            zscore=self.zscore,
        )

    def fit(self, X: Any, y: Sequence) -> "DeltaClassifier":
        y = np.asarray(y)
        self.vectorizer_ = self._make_vectorizer()
        matrix = self.vectorizer_.fit_transform(X)
        self.classes_ = np.unique(y)
        self.centroids_ = np.vstack(
            [matrix[y == label].mean(axis=0) for label in self.classes_]
        )
        return self

    def _distances(self, X: Any) -> np.ndarray:
        matrix = self.vectorizer_.transform(X)
        if self.metric == "manhattan":
            return cdist(matrix, self.centroids_, metric="cityblock") / max(matrix.shape[1], 1)
        if self.metric == "euclidean":
            return cdist(matrix, self.centroids_, metric="euclidean")
        if self.metric == "cosine":
            return cdist(matrix, self.centroids_, metric="cosine")
        raise ValueError(f"Неизвестная метрика Delta: {self.metric}")

    def decision_function(self, X: Any) -> np.ndarray:
        return -self._distances(X)

    def predict(self, X: Any) -> np.ndarray:
        distances = self._distances(X)
        return self.classes_[np.argmin(distances, axis=1)]

    def predict_proba(self, X: Any) -> np.ndarray:
        distances = self._distances(X)
        scale = np.maximum(distances.std(axis=1, keepdims=True), 1e-6)
        return softmax(-distances / scale, axis=1)

    def get_feature_names_out(self) -> np.ndarray:
        return self.vectorizer_.get_feature_names_out()