"""Tests for evaluation metrics."""

from __future__ import annotations

import pytest
import torch

from obliteratus.evaluation.metrics import accuracy, f1_score_metric, perplexity


class TestPerplexity:
    def test_perfect_prediction(self):
        # Create logits that strongly predict the correct next token
        vocab_size = 10
        seq_len = 5
        batch_size = 1

        labels = torch.tensor([[0, 1, 2, 3, 4]])
        logits = torch.full((batch_size, seq_len, vocab_size), -100.0)
        # Set high logit for the correct next token
        for t in range(seq_len - 1):
            logits[0, t, labels[0, t + 1]] = 100.0

        ppl = perplexity(logits, labels)
        assert ppl < 2.0, f"Expected near-1 perplexity, got {ppl}"

    def test_random_prediction_higher(self):
        vocab_size = 100
        seq_len = 20
        batch_size = 2

        torch.manual_seed(42)
        logits = torch.randn(batch_size, seq_len, vocab_size)
        labels = torch.randint(0, vocab_size, (batch_size, seq_len))

        ppl = perplexity(logits, labels)
        assert ppl > 10, f"Random logits should yield high perplexity, got {ppl}"

    def test_uniform_logits_equal_vocabulary_size(self):
        logits = torch.zeros(2, 4, 7)
        labels = torch.tensor([[0, 1, 2, 3], [3, 4, 5, 6]])
        assert perplexity(logits, labels) == pytest.approx(7.0)

    @pytest.mark.parametrize(
        ("logits", "labels", "message"),
        [
            (torch.zeros(2, 3), torch.zeros(2, 3, dtype=torch.long), "3D"),
            (torch.zeros(2, 3, 4), torch.zeros(2, 3, 1, dtype=torch.long), "2D"),
            (torch.zeros(2, 3, 4), torch.zeros(2, 2, dtype=torch.long), "does not match"),
            (torch.zeros(2, 1, 4), torch.zeros(2, 1, dtype=torch.long), "two sequence"),
        ],
    )
    def test_rejects_invalid_shapes(self, logits, labels, message):
        with pytest.raises(ValueError, match=message):
            perplexity(logits, labels)

    def test_rejects_unavailable_and_invalid_targets(self):
        logits = torch.zeros(1, 3, 4)
        with pytest.raises(ValueError, match="every target"):
            perplexity(logits, torch.full((1, 3), -100))
        with pytest.raises(ValueError, match="outside"):
            perplexity(logits, torch.tensor([[0, 1, 9]]))
        logits[0, 0, 0] = float("nan")
        with pytest.raises(ValueError, match="finite"):
            perplexity(logits, torch.tensor([[0, 1, 2]]))

    def test_rejects_non_integer_labels(self):
        with pytest.raises(ValueError, match="integer token IDs"):
            perplexity(torch.zeros(1, 3, 4), torch.zeros(1, 3))


class TestAccuracy:
    def test_perfect(self):
        assert accuracy([1, 2, 3], [1, 2, 3]) == 1.0

    def test_zero(self):
        assert accuracy([1, 2, 3], [4, 5, 6]) == 0.0

    def test_partial(self):
        assert accuracy([1, 2, 3, 4], [1, 2, 0, 0]) == 0.5

    def test_empty(self):
        assert accuracy([], []) == 0.0

    def test_rejects_length_mismatch_instead_of_truncating(self):
        with pytest.raises(ValueError, match="equal length"):
            accuracy([1, 2], [1])


class TestF1:
    def test_perfect(self):
        assert f1_score_metric([0, 1, 0, 1], [0, 1, 0, 1]) == 1.0

    def test_zero(self):
        score = f1_score_metric([0, 0, 0, 0], [1, 1, 1, 1])
        assert score == 0.0

    def test_empty(self):
        assert f1_score_metric([], []) == 0.0

    def test_rejects_length_mismatch(self):
        with pytest.raises(ValueError, match="equal length"):
            f1_score_metric([1, 2], [1])
