from statistics import mean, stdev
from unittest.mock import MagicMock

import pytest
from datasets import Dataset

from data_measurements.measurements import TextLengths


@pytest.fixture
def mock_statistics(monkeypatch):
    mock_mean = MagicMock()
    mock_mean.return_value = "mean"
    mock_stdev = MagicMock()
    mock_stdev.return_value = "stdev"

    monkeypatch.setattr("data_measurements.measurements.text_lengths.mean", mock_mean)
    monkeypatch.setattr("data_measurements.measurements.text_lengths.stdev", mock_stdev)

    return {"mean": mock_mean, "stdev": mock_stdev}


@pytest.fixture
def dataset():
    return Dataset.from_list(
        [
            {"text": "Hello world"},
            {"text": "What is up"},
            {"text": "Kitty Cat"},
        ]
    )


def test_text_lengths_initialize(dummy_tokenizer):
    TextLengths(tokenizer=dummy_tokenizer, feature=None)


def test_text_lengths_run_mock_stats(mock_statistics, dummy_tokenizer, dataset):
    text_lengths = TextLengths(tokenizer=dummy_tokenizer, feature="text")
    results = text_lengths.measure(dataset)

    assert results.average_instance_length == "mean"
    assert results.standard_dev_instance_length == "stdev"
    assert results.num_instance_lengths == 2

    mock_statistics["mean"].assert_called_once()
    mock_statistics["stdev"].assert_called_once()


def test_text_lengths_run(dataset, dummy_tokenizer):
    text_lengths = TextLengths(tokenizer=dummy_tokenizer, feature="text")
    results = text_lengths.measure(dataset)

    assert results.average_instance_length == mean([2, 3, 2])
    assert results.standard_dev_instance_length == stdev([2, 3, 2])
    assert results.num_instance_lengths == 2


def test_text_lengths_figure(dataset, dummy_tokenizer):
    # TODO: Some kind of assertion?
    text_lengths = TextLengths(tokenizer=dummy_tokenizer, feature="text")
    results = text_lengths.measure(dataset)

    results.to_figure()
