import torch
from torch.utils.data import DataLoader, TensorDataset, random_split

features = torch.arange(0, 60, dtype=torch.float32).reshape(20, 3)
labels = torch.arange(20)
dataset = TensorDataset(features, labels)

generator = torch.Generator().manual_seed(20260707)
train_set, valid_set = random_split(dataset, [14, 6], generator=generator)

repeat_train, repeat_valid = random_split(
    dataset,
    [14, 6],
    generator=torch.Generator().manual_seed(20260707),
)

train_indices = list(train_set.indices)
valid_indices = list(valid_set.indices)
repeat_train_indices = list(repeat_train.indices)

print(f"dataset length: {len(dataset)}")
print(f"train length: {len(train_set)}")
print(f"validation length: {len(valid_set)}")
print(f"train indices: {train_indices}")
print(f"validation indices: {valid_indices}")
print(f"overlap count: {len(set(train_indices) & set(valid_indices))}")
print(f"repeat matches: {train_indices == repeat_train_indices}")

train_loader = DataLoader(train_set, batch_size=4, shuffle=True)
valid_loader = DataLoader(valid_set, batch_size=6, shuffle=False)
train_batch, _ = next(iter(train_loader))
_, valid_labels = next(iter(valid_loader))

print(f"train batch shape: {tuple(train_batch.shape)}")
print(f"validation labels: {valid_labels.tolist()}")