import torch


def train_test_split(dataset, params):
    """Grabs random Omniglot samples and generates test samples from same class.

    The random seed is taken from params.sampler_seed, the test_shift is which sample
            to grab as a test. If it ends up being a different class, the sampler is
            walked back until the class is same, and the sample is different. 
        
    Args:
        dataset: (Dataset) Sampler from Omniglot dataset.
        params: (json dict) Params.json file. 

    Returns:
        train_dataloader, test_dataloader: (tuple) Containing matched train test pairs.
    """
    train_dataset = []
    test_dataset = []

    # Random seed from params file. 
    torch.manual_seed(params.sampler_seed)

    # Create batch_size random indices from dataset. 
    #       Subtract params.test_shift so that we don't pick a random sample
    #       so close to the end of the set that it looks for a test pair in 
    #       the blackness of 'index out of range'.  
    idxs = torch.randint(len(dataset) - params.test_shift, (1, params.batch_size))

    # Make sure one of them is our control.
    idxs[0, 0] = 19

    for i, idx in enumerate(idxs[0]):
        shift_idx = params.test_shift
        train_sample, train_lbl = dataset[idx]
        test_sample, test_lbl = dataset[idx + shift_idx]

        # Make sure labels are the same, and it is not the same sample. 
        while (train_lbl != test_lbl) or (torch.equal(train_sample, test_sample)):
            test_sample, test_lbl = dataset[idx + shift_idx]
            shift_idx -= 1

        train_dataset.append(train_sample)
        test_dataset.append(test_sample)
        #=====MONITORING=====#

        # Uncomment to see train_samples or change selection to test_sample.
        # utils.animate_weights(train_sample, auto=True)

        #=====END MONITORING=====#

    train_dataloader = torch.stack(train_dataset)
    train_dataloader.unsqueeze_(1)
    test_dataloader = torch.stack(test_dataset)
    test_dataloader.unsqueeze_(1)

    return train_dataloader, test_dataloader