from __future__ import annotations

import unittest

import numpy as np

from prepare_egospeak_floor_eval import (
    NUM_CLASSES,
    NUM_HORIZONS,
    align_miniroad_prediction,
)


class EgoSpeakFloorAlignmentTests(unittest.TestCase):
    def test_source_forecasts_are_aligned_to_target_rows(self) -> None:
        target_steps = 14
        raw = np.zeros(
            (target_steps - NUM_HORIZONS, NUM_HORIZONS, NUM_CLASSES),
            dtype=np.float32,
        )
        for source_row in range(len(raw)):
            for horizon_index in range(NUM_HORIZONS):
                raw[source_row, horizon_index, 1] = (
                    100 * source_row + horizon_index
                ) / 1000

        aligned, valid = align_miniroad_prediction(raw, target_steps)
        self.assertEqual(aligned.shape, (target_steps, NUM_CLASSES, NUM_HORIZONS))
        self.assertEqual(valid.shape, (target_steps, NUM_HORIZONS))
        for source_row in range(len(raw)):
            for horizon_index in range(NUM_HORIZONS):
                target_row = source_row + horizon_index + 1
                self.assertEqual(
                    aligned[target_row, 1, horizon_index],
                    raw[source_row, horizon_index, 1],
                )
                self.assertTrue(valid[target_row, horizon_index])
        for horizon_index in range(NUM_HORIZONS):
            self.assertFalse(valid[: horizon_index + 1, horizon_index].any())

    def test_shape_mismatch_is_rejected(self) -> None:
        with self.assertRaisesRegex(ValueError, "Expected MiniROADA prediction"):
            align_miniroad_prediction(
                np.zeros((4, NUM_CLASSES, NUM_HORIZONS), dtype=np.float32),
                14,
            )


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