import random


def generate_word2vec_samples(text, num_skips, skip_window):
  """Generates samples that can be used to train a skip-gram word2vec model.
  Code is inspired by the word2vec implementation from:
  https://github.com/tensorflow/tensorflow/blob/master/tensorflow/examples/tutorials/word2vec/word2vec_basic.py

  Args:
    text: The text, where samples should be generated from.
    num_skips: How many samples (current word, context word) should be generated per word.
    skip_window: Size of the context window (skip_window words on the left and on the right
      considered as context.

  Returns:
    Two lists; one containing center words and one for contextual words.
  """
  assert num_skips <= 2 * skip_window
  center_words = []
  context_words = []

  span = 2 * skip_window + 1  # [ skip_window target skip_window ]
  
  # concatenate all sentences in the text
  data = []
  for sentence in text:
    data.extend(sentence)
  
  # contains contextual words for word with index curr_index
  buffer = []
  
  def enqueue_context(i):
    buffer = []
  
    start_left = max(i-skip_window, 0)
    end_right = min(i+skip_window, len(data)-1)
    
    for j in range(start_left, i):
      buffer.append(data[j])
    for j in range(i+1, end_right+1):
      buffer.append(data[j])
      
    return buffer
    
  for i in range(len(data)):
    buffer = enqueue_context(i)
    
    targets_to_avoid = [] # already considered contextual words for current token
    for j in range(min(num_skips, len(buffer))):
      # randomly pick a context word
      target = random.randint(0, len(buffer) - 1)
      while target in targets_to_avoid:
        target = random.randint(0, len(buffer) - 1)
      targets_to_avoid.append(target)
    
      center_words.append(data[i])
      context_words.append(buffer[target])

  return center_words, context_words