from typing import OrderedDict


def splitData(set_x, set_y, split_size=1000):
        """
        this method gets set_x and set_y and splits them into two sets

        Parameters:
        -------------
                set_x: is expected to be a tensor or matrix, which will
                       be split by its rows (across samples)
                set_y: an OrderedDict.
                    For each element in the OrderedDict,
                    it the values of the key are split
                    with the same split_size as the ones in the set_x

                split_size: the size of the second set to be created
        """

        num_samples = len(set_x)
        assert num_samples > split_size

        set_x_1 = set_x[:-split_size]
        set_x_2 = set_x[-split_size:]

        set_y_1 = OrderedDict()
        set_y_2 = OrderedDict()
        if set_y:
            for key in set_y.keys():
                assert len(set_y[key]) == num_samples
                set_y_1[key] = set_y[key][:-split_size]
                set_y_2[key] = set_y[key][-split_size:]

        return [set_x_1, set_y_1, set_x_2, set_y_2]