import unittest

import pandas as pd

from classification_utils import class_position, combined_confidence, map_category_label


class ClassificationUtilsTests(unittest.TestCase):
    def test_known_numeric_categories_are_mapped(self):
        self.assertEqual(map_category_label(0), "Contract")
        self.assertEqual(map_category_label("15"), "Construction Works")
        self.assertEqual(map_category_label(6.0), "Occupational Health and Safety")

    def test_free_text_and_unknown_categories_are_preserved(self):
        self.assertEqual(map_category_label("risk"), "risk")
        self.assertEqual(map_category_label(99), "99")
        self.assertEqual(map_category_label(None), "")

    def test_combined_confidence_is_joint_probability(self):
        party = pd.Series([0.4, 0.9, None, 1.2])
        stakeholder = pd.Series([0.95, 0.8, 0.7, 1.0])

        result = combined_confidence(party, stakeholder)

        expected = [0.38, 0.72, 0.0, 1.0]
        for actual, wanted in zip(result, expected):
            self.assertAlmostEqual(actual, wanted)

    def test_low_confidence_in_either_stage_keeps_joint_confidence_low(self):
        party = pd.Series([0.95])
        stakeholder = pd.Series([0.4])

        self.assertLess(combined_confidence(party, stakeholder).iloc[0], 0.6)

    def test_probability_column_uses_class_order_not_label_value(self):
        classes = [10, 20, 30]

        self.assertEqual(class_position(classes, 20), 1)
        self.assertIsNone(class_position(classes, 99))


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