import random


def split_train_test(sample_names, train_rate=0.75, seed=0):
    """
    Split data into a train and a test sets keeping replicates within the same set
    :param sample_names: list of sample names
    :param train_rate: percentage of training samples
    :param seed: random seed
    :return: lists of train and test sample indices
    """
    # Set random seed
    random.seed(seed)

    # Find replicate segments
    replicate_ranges = {}
    for i, name in enumerate(sample_names):
        repl_nb = int(name.split('_')[-1][1:])
        n = len(replicate_ranges)
        if repl_nb == 1:  # First replicate sample
            replicate_ranges[n] = {'start': i, 'end': i}
        else:
            replicate_ranges[n - 1]['end'] = i

    # Permute unique samples
    unique_sample_idxs = list(replicate_ranges.keys())
    random.shuffle(unique_sample_idxs)

    # Split data
    nb_unique = len(unique_sample_idxs)
    split_point = int(train_rate * nb_unique)
    unique_train = unique_sample_idxs[:split_point]

    # Recover replicates
    train_idxs = []
    test_idxs = []
    for i in range(nb_unique):
        for j in range(replicate_ranges[i]['start'], replicate_ranges[i]['end'] + 1):
            if i in unique_train:
                train_idxs.append(j)
            else:
                test_idxs.append(j)

    assert len(set(train_idxs + test_idxs)) == len(sample_names)
    return train_idxs, test_idxs