import unittest

import numpy as np

from export_death_burden import (
    death_probability,
    find_death_column,
    validate_horizons,
)


class DeathBurdenTests(unittest.TestCase):
    def test_death_column_is_selected_by_code_and_outcome_type(self) -> None:
        source = {
            "tokens/column": np.asarray([0, 1, 2], dtype=np.int64),
            "tokens/token_id": np.asarray([3, 4, 5], dtype=np.int64),
            "tokens/label_code": np.asarray([b"A00", b"I10", b"Death"]),
            "tokens/label_text": np.asarray(
                [b"A00 cholera", b"I10 hypertension", b"Death"]
            ),
            "tokens/outcome_type": np.asarray(
                [b"disease", b"disease", b"death"]
            ),
        }

        self.assertEqual(find_death_column(source), 2)

    def test_death_probability_combines_shape_and_scale(self) -> None:
        result = death_probability(
            shape=np.asarray([2.0, 1.0], dtype=np.float32),
            scale=np.asarray([10.0, 4.0], dtype=np.float32),
            horizons=np.asarray([5.0, 10.0], dtype=np.float64),
        )
        expected = np.asarray(
            [
                [1.0 - np.exp(-0.25), 1.0 - np.exp(-1.0)],
                [1.0 - np.exp(-1.25), 1.0 - np.exp(-2.5)],
            ],
            dtype=np.float32,
        )
        np.testing.assert_allclose(result, expected, rtol=1e-6, atol=1e-7)

    def test_invalid_parameters_produce_nan(self) -> None:
        result = death_probability(
            shape=np.asarray([1.0, -1.0], dtype=np.float32),
            scale=np.asarray([np.nan, 2.0], dtype=np.float32),
            horizons=np.asarray([5.0], dtype=np.float64),
        )
        self.assertTrue(np.isnan(result).all())

    def test_horizons_must_be_positive_and_unique(self) -> None:
        with self.assertRaises(ValueError):
            validate_horizons([0.0, 5.0])
        with self.assertRaises(ValueError):
            validate_horizons([5.0, 5.0])


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