import unittest

import numpy as np

from . import metric


def to_numpy_list(x: list[list[list[int]]]) -> list[list[np.ndarray]]:
    return [[np.array(z) for z in y] for y in x]


class TestEosMethods(unittest.TestCase):
    def helper(self, data_arr, x, y) -> None:
        expected = {
            "EoS Token Pesent": x,
            "EoS Padding": y,
        }
        for data in data_arr:
            res = metric.end_of_sentence(data)
            self.assertEqual(res, expected)

    def test_false_false(self):
        data_arr = to_numpy_list(
            [
                [
                    [1, 2, 3],
                    [1, 2, 4],
                ],
                [
                    [1, 2, 0],
                    [0, 2, 0],
                ],
            ]
        )
        self.helper(data_arr, False, False)

    def test_true_false(self):
        data_arr = to_numpy_list(
            [
                [
                    [1, 2, 0],
                    [1, 2, 0],
                ],
            ]
        )
        self.helper(data_arr, True, False)

    def test_true_true(self):
        data_arr = to_numpy_list(
            [
                [
                    [1, 2, 0],
                    [1, 0, 0],
                ],
            ]
        )
        self.helper(data_arr, True, True)

    def test_true(self):
        pass


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