from typing import Union
from typing import Counter
from typing import Dict
from typing import List
from typing import Tuple
import csv
import math
import random


def create_training_examples(data_path: str, counter: Union[Counter, Dict[str, int]],
                             context_size: int) -> List[Tuple[str, str]]:
  """Creates training example while subsampling frequent words.
  Assumes that text is already tokenized."""
  total_token_count = sum(counter.values())
  examples = []
  window_size = context_size // 2
  with open(data_path) as f:
    reader = csv.reader(f)
    next(reader)  # skip header
    for row in reader:
      sentence = row[0].split()
      kept_words = []
      for word in sentence:
        word_count = counter.get(word, 1)
        word_fraction = word_count / total_token_count
        prob_of_keeping_word = (math.sqrt(word_fraction / 0.001) + 1) * (0.001 / word_fraction)
        if random.random() < prob_of_keeping_word:
          # Randomly dropout words that with a probability related to a terms frequency
          kept_words.append(word)
      for i, target in enumerate(sentence):
        words_to_the_left = sentence[max(0, i - window_size):i]
        words_to_the_right = sentence[i + 1:i + window_size + 1]
        for context in words_to_the_left + words_to_the_right:
          examples.append((target, context))
  return examples