import pytest
import torch
from torch import nn

from scripts.export_classifier_openvino import ClassifierProbabilities, ordered_names, 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_probabilities_keep_training_class_order_and_sum_to_one():
    result = ClassifierProbabilities(FixedClassifier())(torch.rand(2, 3, 128, 128))
    assert result.shape == (2, 3)
    torch.testing.assert_close(result[0], torch.tensor([1., 3., -1.]).softmax(0))
    torch.testing.assert_close(result.sum(1), torch.ones(2))
    assert result.argmax(1).tolist() == [1, 1]


def test_normal_probability_is_not_suppressed():
    # Suppose class 1 is normal: its probability is returned like any other class.
    result = ClassifierProbabilities(FixedClassifier())(torch.zeros(1, 3, 128, 128))
    assert result[0, 1] > .8
    assert (result > 0).all()


def test_standardization_matches_training_per_image():
    exported = ClassifierProbabilities(FixedClassifier(), eps=1e-5)
    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(exported.normalize(images), expected)
    assert torch.isfinite(exported(images)).all()


def test_classes_validate_index_order_and_duplicates():
    assert ordered_names({1: "blackdot", 0: "normal"}) == ["normal", "blackdot"]
    with pytest.raises(ValueError, match="Duplicate"):
        ordered_names(["normal", "normal"])
    with pytest.raises(ValueError, match="contiguous"):
        ordered_names({1: "normal"})


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)