from types import SimpleNamespace

import pytest
import torch

from gnn4colliders.tasks import BinaryClassificationTask, MulticlassClassificationTask


def _batch(labels, weights):
    return SimpleNamespace(
        labels=torch.tensor(labels),
        metadata=SimpleNamespace(weight=torch.tensor(weights)),
    )


def test_binary_normalizes_column_or_vector_targets_and_uses_named_weight():
    task = BinaryClassificationTask()
    batch = _batch([[0], [0], [1], [1]], [1.0, 2.0, 1.0, 3.0])
    logits = torch.tensor([[-1.0], [0.0], [0.5], [1.0]])
    expected_elementwise = torch.nn.functional.binary_cross_entropy_with_logits(
        logits[:, 0], batch.labels[:, 0].float(), reduction="none"
    )
    expected = (
        (expected_elementwise[:2] * torch.tensor([1.0, 2.0])).sum() / 3
        + (expected_elementwise[2:] * torch.tensor([1.0, 3.0])).sum() / 4
    ) / 2
    assert torch.allclose(task.loss(logits, batch), expected)


def test_binary_negative_weights_are_not_absolute_by_default():
    batch = _batch([0, 0, 1, 1], [1.0, -2.0, 1.0, 1.0])
    logits = torch.tensor([[0.0], [2.0], [0.0], [2.0]])
    task = BinaryClassificationTask()
    elementwise = torch.nn.functional.binary_cross_entropy_with_logits(
        logits[:, 0], batch.labels.float(), reduction="none"
    )
    expected = (
        (elementwise[:2] * torch.tensor([1.0, -2.0])).sum() / -1
        + elementwise[2:].mean()
    ) / 2
    assert task.loss(logits, batch).item() == pytest.approx(expected.item())
    assert BinaryClassificationTask(absolute_weights=True).loss(
        logits, batch
    ).item() != pytest.approx(expected.item())


def test_binary_predictions_and_undefined_auc():
    task = BinaryClassificationTask()
    batch = _batch([1, 1], [1.0, 2.0])
    output = task.predictions(torch.tensor([[-2.0], [2.0]]))
    assert output["scores"].shape == (2,)
    assert output["predictions"].tolist() == [False, True]
    assert (
        task.metrics(torch.tensor([[-2.0], [2.0]]), batch)["roc_auc"]
        != task.metrics(torch.tensor([[-2.0], [2.0]]), batch)["roc_auc"]
    )


def test_multiclass_loss_and_metrics():
    task = MulticlassClassificationTask()
    batch = _batch([0, 1, 2], [1.0, 1.0, 1.0])
    logits = torch.eye(3)
    assert task.loss(logits, batch).item() > 0
    output = task.predict(logits)
    assert output["predictions"].tolist() == [0, 1, 2]
    assert task.metrics(logits, batch)["accuracy"] == 1.0


def test_binary_loss_is_invariant_under_joint_event_permutation():
    task = BinaryClassificationTask()
    batch = _batch([0, 1, 0, 1], [1.0, 2.0, 3.0, 4.0])
    logits = torch.tensor([[-2.0], [1.0], [0.5], [-0.5]])
    permutation = torch.tensor([2, 0, 3, 1])
    shuffled = _batch(batch.labels[permutation].tolist(), batch.metadata.weight[permutation].tolist())
    assert task.loss(logits, batch) == pytest.approx(task.loss(logits[permutation], shuffled).item())


def test_multiclass_metrics_are_invariant_under_joint_event_permutation():
    task = MulticlassClassificationTask()
    batch = _batch([0, 1, 2, 1], [1.0, 2.0, 1.0, 3.0])
    logits = torch.tensor([[3.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 4.0], [0.0, 1.0, 0.0]])
    permutation = torch.tensor([3, 0, 2, 1])
    shuffled = _batch(batch.labels[permutation].tolist(), batch.metadata.weight[permutation].tolist())
    assert task.metrics(logits, batch) == task.metrics(logits[permutation], shuffled)