"""Tests for the evaluation metrics on small known fixtures."""

from __future__ import annotations

import numpy as np
import pytest

from protean_research.evaluation.metrics import (
    classification_metrics,
    confusion_matrix,
    rare_class_metrics,
)


def test_confusion_matrix_counts():
    matrix = confusion_matrix([0, 0, 1, 1], [0, 1, 1, 1], num_classes=2)
    assert matrix.tolist() == [[1, 1], [0, 2]]


def test_classification_metrics_known_fixture():
    y_true = [0, 0, 1, 1, 2, 2]
    y_pred = [0, 0, 1, 2, 2, 1]
    metrics = classification_metrics(y_true, y_pred, num_classes=3)
    assert metrics["accuracy"] == pytest.approx(4 / 6)
    assert metrics["macro_f1"] == pytest.approx((1.0 + 0.5 + 0.5) / 3)
    per_class = {row["class_id"]: row for row in metrics["per_class"]}
    assert per_class[0]["f1"] == pytest.approx(1.0)
    assert per_class[1]["precision"] == pytest.approx(0.5)
    assert per_class[2]["recall"] == pytest.approx(0.5)


def test_rare_class_metrics_restricts_to_given_classes():
    y_true = [0, 0, 1, 1, 2, 2]
    y_pred = [0, 0, 1, 2, 2, 1]
    rare = rare_class_metrics(y_true, y_pred, num_classes=3, rare_class_ids=[1, 2])
    assert rare["rare_class_macro_f1"] == pytest.approx(0.5)
    assert rare["rare_class_recall"] == pytest.approx(0.5)
    assert rare["rare_support"] == 4


def test_rare_class_metrics_skips_absent_classes():
    rare = rare_class_metrics([0, 0], [0, 0], num_classes=3, rare_class_ids=[2])
    assert rare["rare_class_macro_f1"] == 0.0
    assert rare["rare_support"] == 0