import random
import torch


def _create_cl_data(data):
  """Create data batches for contrastive learning.

  Args:
    data (torch.tensor): (batch_size, sequence_length, 2)
    
  Outputs:
    x1 (torch.tensor): (batch_size, sequence_length_A, 2)
    x2 (torch.tensor): (batch_size, sequence_length_B, 2)
    y (torch.tensor): (batch_size, ) binary labels s.t. if x1 and x2 are from the 
      same document (sequence), it is 1; otherwise, it's 0
  """
  n, s, _ = data.shape
  assert n >= 2

  # Random input sequence length
  s1 = random.randrange(int(s * 0.2), int(s * 0.4))
  s2 = random.randrange(int(s * 0.2), int(s * 0.4))

  x1 = torch.zeros((n, s1, 2)).to(data.device)
  x2 = torch.zeros((n, s2, 2)).to(data.device)
  y = torch.zeros(n).to(data.device)

  for i in range(n):
    # get x1
    try:
      x1_start = random.randrange(0, s - s1)
    except:
      x1_start = 0
    x1[i, :, :] = data[i, x1_start:x1_start + s1, :]

    if random.random() > 0.5:
      # Get x2 from the same sequence
      j = i
      y[i] = 1
    else:
      # Get x2 from different sequence
      j = i
      y[i] = 0
      while j == i:
        j = random.randrange(0, n)

    try:
      x2_start = random.randrange(0, s - s2)
    except:
      x2_start = 0

    x2[i, :, :] = data[j, x2_start:x2_start + s2, :]

  return x1.float(), x2.float(), y.float()