"""Тесты дельты Барроуза."""

from __future__ import annotations

import numpy as np
import pandas as pd

from stylometry.delta import DeltaClassifier, MFWVectorizer


def _toy_frame() -> pd.DataFrame:
    texts = [
        "the cat sat on the mat",
        "the cat sat on the mat again",
        "the dog sat on the log",
        "the dog sat on the log again",
        "a cat and a dog",
        "a cat and a dog again",
    ]
    labels = ["cat", "cat", "dog", "dog", "cat", "dog"]
    return pd.DataFrame({"text": texts, "character": labels})


def test_mfw_vectorizer_shape() -> None:
    frame = _toy_frame()
    vectorizer = MFWVectorizer(n_features=5)
    matrix = vectorizer.fit_transform(frame)
    assert matrix.shape == (6, 5)
    assert np.isfinite(matrix).all()


def test_delta_fit_predict() -> None:
    frame = _toy_frame()
    classifier = DeltaClassifier(n_features=5, metric="manhattan")
    classifier.fit(frame, frame["character"].values)
    predictions = classifier.predict(frame)
    assert len(predictions) == len(frame)
    assert set(predictions) <= {"cat", "dog"}
    probabilities = classifier.predict_proba(frame)
    assert probabilities.shape == (len(frame), 2)
    assert np.allclose(probabilities.sum(axis=1), 1.0)