#!/usr/bin/python

import argparse
import csv

from sortedcontainers import SortedSet

FREQ_THRESHOLD = 500
NEWLINE = '\n'
SPACE = " "

def get_bigrams(sent):
  bigrams = []
  for i in range(0, len(sent)-1):
    bigrams.append(sent[i] + SPACE + sent[i+1])
  return bigrams

def main(args):
  dict = {}
  stopwords = SortedSet([])
  with open(args.input_examples, 'r') as input_examples_file:
    reader = csv.DictReader(input_examples_file, delimiter='\t',
                            fieldnames=['question', 'answer'])
    for row in reader:
      question = row['question']
      q_words = question.split(" ")
      for word in q_words:
        freq = dict.get(word, 0)
        dict[word] = freq+1
      #also consider common bigram phrases as stopwords, eg. the story
      for word in get_bigrams(q_words):
        freq = dict.get(word, 0)
        dict[word] = freq+1

  if args.kb_docs:
    with open(args.kb_docs, 'r') as kb_docs_file:
      reader = csv.DictReader(kb_docs_file, delimiter='|', fieldnames=["entity_name", "fieldname", "content"])
      for row in reader:
        content = row["content"]
        content_words = content.split(" ")
        for word in content_words:
          freq = dict.get(word, 0)
          dict[word] = freq+1

  with open(args.output, 'w') as output_file:
    writer = csv.DictWriter(output_file, delimiter='\t', fieldnames=['stopword', 'frequency'])
    for word in dict.keys():
      if dict[word] > FREQ_THRESHOLD:
        stopwords.add(word)
    for stopword in stopwords:
      writer.writerow({'stopword': stopword, 'frequency': dict[stopword]})


if __name__ == "__main__":
  parser = argparse.ArgumentParser(description='Specify arguments')
  parser.add_argument('--input_examples', help='the raw qa pairs', required=True)
  parser.add_argument('--kb_docs', help='the clean_wiki-entities_kb_doc.txt file generated by clean_kb.py',
                      required=False)
  parser.add_argument('--output', help='the stopwords file', required=True)
  args = parser.parse_args()
  main(args)