import unittest

import numpy as np

from app.inference.convnext_onnx import softmax, to_probabilities


class ClassifierProbabilityTests(unittest.TestCase):
    def test_existing_probabilities_are_not_softmaxed_again(self):
        output = np.array([0.01, 0.09, 0.90], dtype=np.float32)
        probabilities = to_probabilities(output)
        np.testing.assert_allclose(probabilities, output, atol=1e-6)

    def test_logits_are_converted_with_softmax(self):
        logits = np.array([-1.0, 0.0, 3.0], dtype=np.float32)
        np.testing.assert_allclose(
            to_probabilities(logits),
            softmax(logits),
            atol=1e-6,
        )


if __name__ == "__main__":
    unittest.main()