import pickle
import random
from torch_geometric.loader import DataLoader

def load_data(tcga_path, split: float, batch_size: int):
    with open(tcga_path, 'rb') as f:
        tcga_graphs = pickle.load(f)
    random.shuffle(tcga_graphs)
    num_train = int(split * len(tcga_graphs))
    train = tcga_graphs[:num_train]
    test  = tcga_graphs[num_train:]
    train_loader = DataLoader(train, batch_size=batch_size, shuffle=True)
    test_loader  = DataLoader(test,  batch_size=batch_size, shuffle=False)
    return train_loader, test_loader
