import torch
import random


def create_sample_batch(speaker_data, batch_size, vector_dim):
    """
    Return torch tensors ((input1, input2), target) for siamese
    network training (for mutual information).

    Constructs each batch to have roughly equal amount of target
    and non-target samples to keep training balanced
    """
    input1_tensor = torch.zeros(batch_size, vector_dim).float()
    input2_tensor = torch.zeros(batch_size, vector_dim).float()
    target_tensor = torch.zeros(batch_size, 1).float()

    speaker_ids = list(speaker_data.keys())
    for i in range(batch_size):
        # Throw a coin if we add target or non-target sample
        input1 = None
        input2 = None
        
        target = None
        if random.random() < 0.5:
            # Target sample:
            #   Pick speaker, and then pick two separate vectors
            #   from that speaker
            random_speaker_id = random.choice(speaker_ids)
            speaker_vectors = speaker_data[random_speaker_id]
            # random.sample takes two unique vectors 
            input1, input2 = random.sample(speaker_vectors, 2)
            target = 1
        else:
            # Non-target sample
            #  Pick two different speakers and one vector from
            #  each
            speaker1, speaker2 = random.sample(speaker_ids, 2)
            input1 = random.choice(speaker_data[speaker1])
            input2 = random.choice(speaker_data[speaker2])
            target = 0

        # Put sampled vectors to batch
        input1_tensor[i] = torch.from_numpy(input1).float()
        input2_tensor[i] = torch.from_numpy(input2).float()
        target_tensor[i, 0] = target

    input1_tensor = input1_tensor.cuda()
    input2_tensor = input2_tensor.cuda()
    target_tensor = target_tensor.cuda()

    return ((input1_tensor, input2_tensor), target_tensor)