import torch


def species_split(dataset, train_valid_species_id_list=[3702, 6239, 511145,
                                                        7227, 10090, 4932, 7955],
                  test_species_id_list=[9606]):
    """
    Split dataset based on species_id attribute
    :param dataset:
    :param train_valid_species_id_list:
    :param test_species_id_list:
    :return: train_valid dataset, test dataset
    """
    # NB: pytorch geometric dataset object can be indexed using slices or
    # byte tensors. We will use byte tensors here

    train_valid_byte_tensor = torch.zeros(len(dataset), dtype=torch.uint8)
    for id in train_valid_species_id_list:
        train_valid_byte_tensor += (dataset.data.species_id == id)

    test_species_byte_tensor = torch.zeros(len(dataset), dtype=torch.uint8)
    for id in test_species_id_list:
        test_species_byte_tensor += (dataset.data.species_id == id)

    assert ((train_valid_byte_tensor + test_species_byte_tensor) == 1).all()

    train_valid_dataset = dataset[train_valid_byte_tensor]
    test_valid_dataset = dataset[test_species_byte_tensor]

    return train_valid_dataset, test_valid_dataset