import pytest
import torch
from torch import nn

from scripts.export_classifier_openvino import YoloClassifierAdapter, class_order, preprocessing


class FixedClassifier(nn.Module):
    def forward(self, x):
        return torch.tensor([1.0, 3.0, -1.0]).expand(x.shape[0], -1)


def test_yolo_contract_and_class_reordering():
    adapter = YoloClassifierAdapter(FixedClassifier(), 128, 336, [2, 0, 1])
    result = adapter(torch.rand(1, 3, 128, 128))
    assert result.shape == (1, 7, 336)
    torch.testing.assert_close(result[0, :4, 0], torch.tensor([64., 64., 128., 128.]))
    expected = torch.tensor([1., 3., -1.]).softmax(0)[[2, 0, 1]]
    torch.testing.assert_close(result[0, 4:, 0], expected)
    assert result[0, 4:, 0].argmax().item() == 2
    assert torch.count_nonzero(result[:, :, 1:]) == 0


def test_normalization_matches_training_for_each_image():
    adapter = YoloClassifierAdapter(FixedClassifier(), 128, 1, [0, 1, 2])
    images = torch.stack((torch.rand(3, 128, 128), torch.ones(3, 128, 128)))
    expected = torch.stack([(x - x.mean()) / (x.std() + 1e-5) for x in images])
    torch.testing.assert_close(adapter.normalize(images), expected)
    assert torch.isfinite(adapter(images)).all()


def test_class_sets_must_match():
    assert class_order(["normal", "mosquito"], ["mosquito", "normal"]) == [1, 0]
    with pytest.raises(ValueError, match="Class sets differ"):
        class_order(["normal", "mosquito"], ["normal", "blackdot"])
    with pytest.raises(ValueError, match="Duplicate"):
        class_order(["normal", "normal"], ["normal"])


def test_background_suppression_uses_reordered_output_index():
    # Source class 1 wins, but it is output class 2 after reordering.
    adapter = YoloClassifierAdapter(FixedClassifier(), 128, 336, [2, 0, 1], background_index=2)
    result = adapter(torch.rand(1, 3, 128, 128))
    assert result.shape == (1, 7, 336)
    assert torch.count_nonzero(result[:, 4:, :]) == 0
    # The existing C++ decoder returns zero candidates at any positive threshold.
    assert not (result[0, 4:, :].max(dim=0).values >= 0.25).any()


def test_foreground_class_indices_and_scores_are_preserved():
    foreground = YoloClassifierAdapter(FixedClassifier(), 128, 336, [0, 1, 2],
                                      background_index=0)
    legacy = YoloClassifierAdapter(FixedClassifier(), 128, 336, [0, 1, 2])
    x = torch.rand(1, 3, 128, 128)
    torch.testing.assert_close(foreground(x), legacy(x))
    assert foreground(x)[0, 4:, 0].argmax() == 1


def test_background_gate_is_per_image_in_batch():
    class BatchClassifier(nn.Module):
        def forward(self, x):
            return x[:, :, 0, 0]

    x = torch.zeros(2, 3, 4, 4)
    x[0, 0, 0, 0] = 3  # normal
    x[1, 1, 0, 0] = 3  # blackdot
    adapter = YoloClassifierAdapter(BatchClassifier(), 4, 10, [0, 1, 2],
                                    eps=None, background_index=0)
    result = adapter(x)
    assert not torch.count_nonzero(result[0, 4:, :])
    assert result[1, 4:, 0].argmax() == 1
    assert result[1, 5, 0] > 0.9


def test_refuses_resampling_and_unknown_preprocessing():
    config = {"data": {"val_transforms": [
        {"name": "Resize", "size": [160, 160]}, {"name": "ToTensor"},
        {"name": "PerImageStandardize"}]}}
    with pytest.raises(ValueError, match="Training size differs"):
        preprocessing(config, 128)
    config["data"]["val_transforms"] = [{"name": "Normalize"}]
    with pytest.raises(ValueError, match="Unsupported"):
        preprocessing(config, 128)