"""Tests for dgtw_mlflow_utils.metrics.

================================
Covers: compute_retrieval_metrics function.
"""

import pytest
import tensorflow as tf
from dgtw_mlflow_utils.metrics import compute_retrieval_metrics


class TestComputeRetrievalMetrics:
    """Tests for the compute_retrieval_metrics function."""

    def test_perfect_predictions(self):
        """When logits have highest scores on the diagonal, metrics should be perfect."""
        batch_size = 8
        # Identity matrix gives perfect diagonal scores
        logits = tf.eye(batch_size) * 10.0

        metrics = compute_retrieval_metrics(logits, k=3)

        assert float(metrics["recall_3"].numpy()) == pytest.approx(1.0)
        assert float(metrics["mrr_3"].numpy()) == pytest.approx(1.0)
        assert float(metrics["ndcg_3"].numpy()) == pytest.approx(1.0)

    def test_perfect_predictions_large_k(self):
        """Recall@K should be 1.0 when predictions are perfect and K equals batch_size."""
        batch_size = 10
        logits = tf.eye(batch_size) * 10.0

        metrics = compute_retrieval_metrics(logits, k=10)
        assert float(metrics["recall_10"].numpy()) == pytest.approx(1.0)

    def test_random_predictions_recall_bounded(self):
        """With random logits, recall should be between 0 and 1."""
        tf.random.set_seed(42)
        batch_size = 16
        logits = tf.random.normal((batch_size, batch_size))

        metrics = compute_retrieval_metrics(logits, k=3)

        recall = float(metrics["recall_3"].numpy())
        assert 0.0 <= recall <= 1.0

    def test_mrr_bounded(self):
        """MRR should always be between 0 and 1."""
        tf.random.set_seed(42)
        logits = tf.random.normal((8, 8))

        metrics = compute_retrieval_metrics(logits, k=5)
        mrr = float(metrics["mrr_5"].numpy())
        assert 0.0 < mrr <= 1.0

    def test_ndcg_positive(self):
        """NDCG should always be positive."""
        tf.random.set_seed(42)
        logits = tf.random.normal((8, 8))

        metrics = compute_retrieval_metrics(logits, k=5)
        ndcg = float(metrics["ndcg_5"].numpy())
        assert ndcg > 0.0

    def test_output_keys(self):
        """Check that the correct keys are returned."""
        logits = tf.eye(4)
        metrics = compute_retrieval_metrics(logits, k=3)

        assert "recall_3" in metrics
        assert "mrr_3" in metrics
        assert "ndcg_3" in metrics

    def test_key_changes_with_k(self):
        """The recall key should reflect the K value."""
        logits = tf.eye(8)

        m3 = compute_retrieval_metrics(logits, k=3)
        assert "recall_3" in m3

        m5 = compute_retrieval_metrics(logits, k=5)
        assert "recall_5" in m5

    def test_output_types_are_tensors(self):
        """All returned values should be TensorFlow tensors."""
        logits = tf.eye(4)
        metrics = compute_retrieval_metrics(logits, k=3)

        for key, value in metrics.items():
            assert isinstance(value, tf.Tensor), f"{key} is not a Tensor"

    def test_boosted_diagonal_improves_recall(self):
        """Adding a boost to the diagonal should improve recall."""
        tf.random.set_seed(42)
        batch_size = 16
        base_logits = tf.random.normal((batch_size, batch_size))

        metrics_bad = compute_retrieval_metrics(base_logits, k=3)
        boosted = base_logits + tf.eye(batch_size) * 5.0
        metrics_good = compute_retrieval_metrics(boosted, k=3)

        assert float(metrics_good["recall_3"].numpy()) >= float(metrics_bad["recall_3"].numpy())

    def test_small_batch(self):
        """Should work with batch_size=2."""
        logits = tf.constant([[2.0, 0.5], [0.3, 1.8]])
        metrics = compute_retrieval_metrics(logits, k=1)

        assert "recall_1" in metrics
        assert float(metrics["recall_1"].numpy()) == pytest.approx(1.0)


class TestComputeRetrievalMetricsNumpy:
    """Tests for the compute_retrieval_metrics_numpy function."""

    def test_perfect_predictions(self):
        """When logits have highest scores on the diagonal, metrics should be perfect."""
        import numpy as np

        batch_size = 8
        logits = np.eye(batch_size) * 10.0

        from dgtw_mlflow_utils.metrics import compute_retrieval_metrics_numpy

        metrics = compute_retrieval_metrics_numpy(logits, k=3)

        assert metrics["recall_3"] == pytest.approx(1.0)
        assert metrics["mrr_3"] == pytest.approx(1.0)
        assert metrics["ndcg_3"] == pytest.approx(1.0)

    def test_random_predictions_recall_bounded(self):
        """With random logits, recall should be between 0 and 1."""
        import numpy as np

        np.random.seed(42)
        batch_size = 16
        logits = np.random.randn(batch_size, batch_size)

        from dgtw_mlflow_utils.metrics import compute_retrieval_metrics_numpy

        metrics = compute_retrieval_metrics_numpy(logits, k=3)

        assert 0.0 <= metrics["recall_3"] <= 1.0

    def test_output_keys(self):
        """Check that the correct keys are returned."""
        import numpy as np

        logits = np.eye(4)

        from dgtw_mlflow_utils.metrics import compute_retrieval_metrics_numpy

        metrics = compute_retrieval_metrics_numpy(logits, k=3)

        assert "recall_3" in metrics
        assert "mrr_3" in metrics
        assert "ndcg_3" in metrics

    def test_key_changes_with_k(self):
        """The recall key should reflect the K value."""
        import numpy as np

        logits = np.eye(8)

        from dgtw_mlflow_utils.metrics import compute_retrieval_metrics_numpy

        m3 = compute_retrieval_metrics_numpy(logits, k=3)
        assert "recall_3" in m3

        m5 = compute_retrieval_metrics_numpy(logits, k=5)
        assert "recall_5" in m5
