import itertools
import math

import pytest
import torch

from gnn4colliders.tasks import (
    BinaryClassificationTask,
    MulticlassClassificationTask,
    macro_scores,
    per_class_weighted_mean,
    weighted_binary_auc,
)


def _pairwise_auc(labels, scores, weights):
    total = correct = 0.0
    for i, j in itertools.product(range(len(labels)), repeat=2):
        if labels[i] and not labels[j]:
            pair = weights[i] * weights[j]
            total += pair
            if scores[i] > scores[j]:
                correct += pair
            elif scores[i] == scores[j]:
                correct += 0.5 * pair
    return correct / total


def test_sorted_auc_matches_pairwise_definition_with_ties():
    generator = torch.Generator().manual_seed(0)
    labels = torch.rand(60, generator=generator) > 0.4
    scores = torch.randint(0, 8, (60,), generator=generator).float()
    weights = torch.rand(60, generator=generator)
    assert weighted_binary_auc(labels, scores, weights) == pytest.approx(
        _pairwise_auc(labels.tolist(), scores.tolist(), weights.tolist())
    )
    assert math.isnan(weighted_binary_auc(labels & False, scores, weights))


def test_loss_is_the_mean_of_per_class_weighted_means():
    loss = torch.tensor([1.0, 2.0, 3.0, 4.0])
    labels = torch.tensor([0, 0, 2, 2])
    weights = torch.tensor([1.0, 3.0, 1.0, 1.0])
    expected = ((1 * 1 + 3 * 2) / 4 + (3 + 4) / 2) / 2
    assert float(per_class_weighted_mean(loss, labels, weights, 3)) == pytest.approx(
        expected
    )


def test_macro_scores():
    scores = macro_scores(torch.tensor([0, 0, 1, 1]), torch.tensor([0, 1, 1, 1]), 2)
    assert scores["precision"] == pytest.approx((1 + 2 / 3) / 2)
    assert scores["recall"] == pytest.approx((0.5 + 1) / 2)


def test_multiclass_metrics_exclude_non_positive_weights():
    task = MulticlassClassificationTask()
    logits = torch.tensor([[2.0, 0.0], [0.0, 2.0], [2.0, 0.0]])
    labels = torch.tensor([0, 1, 1])
    metrics = task.metrics(logits, labels, torch.tensor([1.0, 1.0, -1.0]))
    assert metrics["auc"] == 1.0
    assert metrics["accuracy"] == pytest.approx(2 / 3)


def test_binary_task_balances_classes_and_uses_absolute_weights():
    task = BinaryClassificationTask(absolute_weights=True)
    logits = torch.tensor([[0.0], [0.0], [0.0]])
    labels = torch.tensor([1, 0, 0])
    loss = task.loss(logits, labels, torch.tensor([-1.0, 1.0, 2.0]))
    assert float(loss) == pytest.approx(float(torch.log(torch.tensor(2.0))))
    assert task.metrics(torch.tensor([3.0, -1.0, 0.0]), labels, torch.ones(3))[
        "auc"
    ] == pytest.approx(1.0)