from typing import Any

from numpy.typing import ArrayLike
from sklearn.model_selection import train_test_split


def split_train_test(
    *data: Any,
    test_size: float | None = None,
    train_size: float | None = None,
    random_state: int | None = None,
    shuffle: bool = True,
    stratify: ArrayLike | None = None,
) -> tuple[Any, Any, Any, Any]:
    """
    Split arrays or matrices into random train and test subsets.

    Wrapper for sklearn.model_selection.train_test_split.

    Args:
        *data: data to be split.
        test_size: size of test subset.
        train_size: size of train subset.
        random_state: Seed applied to the data before applying the split.
        shuffle: Whether or not to shuffle the data before splitting.
        stratify: If not None, data is split in a stratified fashion, using this as the class labels.

    Returns:
        X_train, X_test, y_train, y_test

    Thread-safe: This function is stateless and thread-safe.
    """
    X_train, X_test, y_train, y_test = train_test_split(
        *data,
        test_size=test_size,
        train_size=train_size,
        random_state=random_state,
        shuffle=shuffle,
        stratify=stratify,
    )
    return X_train, X_test, y_train, y_test