import gensim
import nltk
from gensim.utils import simple_preprocess
from gensim.parsing.preprocessing import STOPWORDS
from nltk.stem import WordNetLemmatizer, SnowballStemmer
from nltk.stem.porter import *
import ssl
from pathlib import Path
import pandas as pd
import numpy as np
from sys import argv
import json
import random
import csv

# lemmatizing and stemming
def lem_stem(stemmer, text):
    return stemmer.stem(WordNetLemmatizer().lemmatize(text,pos='v'))

def preprocess(stemmer,text):
    result = []
    for token in gensim.utils.simple_preprocess(text):
        if token not in STOPWORDS and len(token) >= 3:
            result.append(lem_stem(stemmer, token))
    return result

def read_process_data(filepath):
    try:
        _create_unverified_https_context = ssl._create_unverified_context
    except AttributeError:
        pass
    else:
        ssl._create_default_https_context = _create_unverified_https_context
    # Read and preprocess data
    stemmer = SnowballStemmer('english')
    nltk.download('wordnet')
    documentsx = pd.read_csv(filepath, error_bad_lines=False)
    processed_docs = documentsx['abstract'].map(lambda x: preprocess(stemmer,x))
    dictionary = gensim.corpora.Dictionary(processed_docs)
    dictionary.filter_extremes(no_below=10, no_above=0.5)
    bow_corpus = [dictionary.doc2bow(doc) for doc in processed_docs]
    dictionary[1] # neccesary to generate id2token otherwise it will be empty
    id2token = dictionary.id2token

    # Process data
    num_doc = len(bow_corpus)
    num_vocab = len(dictionary)
    documents = {}
    for i, doc in enumerate(bow_corpus):
        tmp_doc = []
        for word in doc:
            id = word[0]
            freq = word[1]
            if(freq==0):
                continue
            while(freq!=0):
                tmp_doc.append(id)
                freq -=1
        random.shuffle(tmp_doc)
        documents[i] = tmp_doc
    return [documents, id2token]

def write_data(outdir, ident, documents,id2token):
    with open(outdir + '/' + ident + '_docs.csv', 'w') as csv_file:
        writer = csv.writer(csv_file, delimiter=' ')
        for key, value in documents.items():
            writer.writerow([int(key)] + value)
    with open(outdir + '/' + ident + '_id2token.csv', 'w') as csv_file:
        writer = csv.writer(csv_file, delimiter=' ')
        for key, value in id2token.items():
            writer.writerow([key, value])

if __name__ == '__main__':
    filepath = 'papers2017.csv'
    outdir = '_tmp'
    ident = 'neurips'

    while argv:
        if argv[0] == '-filepath':
            filepath = argv[1]
        elif argv[0] == '-outdir':
            outdir = argv[1]
        elif argv[0] == '-ident':
            ident = argv[1]
        argv = argv[1:]
    Path(outdir).mkdir(parents=True,exist_ok=True)
    [documents,id2token] = read_process_data(filepath)
    write_data(outdir,ident,documents,id2token)


