import torch
from typing import Tuple


def _split_crossval(xy: torch.Tensor, crossval_count: int, crossval_index: int) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Generates a split of the given dataset along the first dimension for cross-validation.

    :param xy: The data that should be split. The split will be generated acros dimension 0.
    :param crossval_count: The number of splits in total
    :param crossval_index: The index of the split that should be generated (0 <= crossval_index < crossval_count)
    :return: A tuple of (training data, validation data)
    """
    n = xy.shape[0]
    split_size = n // crossval_count
    val_start = crossval_index * split_size
    val_end = (crossval_index + 1) * split_size
    train1_start = 0 if crossval_index == 0 else (crossval_index - 1) * split_size
    train1_end = 0 if crossval_index == 0 else val_start
    train2_start = val_end if crossval_index < (crossval_count - 1) else 0
    train2_end = n if crossval_index < (crossval_count - 1) else 0
    val = xy[val_start:val_end]
    train = torch.concat([xy[train1_start:train1_end], xy[train2_start:train2_end]])
    return (train, val)