import unittest
from pathlib import Path

import numpy as np

from dataset import _ExpoBaseDataset
from targets import CHECKUP_IDX
from train_util import load_extra_info_types_file


class CheckupSelectionTests(unittest.TestCase):
    @staticmethod
    def _base(extra_info_types):
        dataset = _ExpoBaseDataset.__new__(_ExpoBaseDataset)
        dataset.extra_info_types = list(extra_info_types)
        dataset.event_data = np.asarray(
            [
                [101, 10, CHECKUP_IDX],
                [101, 20, 2],
                [101, 30, 3],
            ],
            dtype=np.float64,
        )
        return dataset

    def test_explicit_empty_extra_info_removes_checkup(self):
        project_root = Path(__file__).resolve().parents[1]
        selected_types = load_extra_info_types_file(
            str(project_root / "extra_info_types_none.txt")
        )
        self.assertEqual(selected_types, [])
        dataset = self._base(selected_types)

        rows = list(dataset._iter_patient_events(impute_no_event_gaps=False))

        self.assertEqual(len(rows), 1)
        eid, times, labels = rows[0]
        self.assertEqual(eid, 101)
        np.testing.assert_array_equal(times, np.asarray([20, 30], dtype=np.float32))
        self.assertNotIn(CHECKUP_IDX, labels.tolist())

    def test_selected_extra_info_keeps_checkup(self):
        dataset = self._base([11])

        rows = list(dataset._iter_patient_events(impute_no_event_gaps=False))

        self.assertEqual(len(rows), 1)
        _, times, labels = rows[0]
        np.testing.assert_array_equal(
            times,
            np.asarray([10, 20, 30], dtype=np.float32),
        )
        self.assertEqual(labels[0], CHECKUP_IDX)


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