from optparse import OptionParser

import re
import os
import sys

import numpy as np

from ..util import dirs
from ..util import file_handling as fh
from ..preprocessing import data_splitting as ds
from ..feature_extractors.vocabulary_with_counts import VocabWithCounts

def 
():
    usage = "%prog project"
    parser = OptionParser(usage=usage)
    parser.add_option('-v', dest='vocab_size', default=10000,
                      help='Vocabulary size (most frequent words): default=%default')
    parser.add_option('--seed', dest='seed', default=42,
                      help='Random seed: default=%default')
    #parser.add_option('--boolarg', action="store_true", dest="boolarg", default=False,
    #                  help='Keyword argument: default=%default')

    (options, args) = parser.parse_args()
    project_name = args[0]

    dirs.make_base_dir(project_name)

    vocab_size = int(options.vocab_size)

    suffixes = {"'s", "n't"}
    pronouns = {"i", 'you', 'he', 'his', 'she', 'her', 'hers', 'it', 'its', 'we', 'you', 'your', 'they', 'them', 'their'}
    determiners = {'a', 'an', 'the', 'this', 'that', 'these', 'those'}
    prepositions = {'at', 'by', 'for', 'from', 'in', 'into', 'of', 'on', 'than', 'to', 'with'}
    transitional = {'and', 'also', 'as', 'but', 'if', 'or',  'then'}
    common_verbs = {'are', 'be', 'been', 'had', 'has', 'have', 'is', 'said', 'was', 'were'}
    stopwords = suffixes.union(pronouns).union(determiners).union(prepositions).union(transitional).union(common_verbs)
    print "Removing %d stopwords:" % len(stopwords)
    for s in stopwords:
        print s

    # set random seed
    np.random.seed(int(options.seed))

    # read in data
    dirs.make_base_dir(project_name)
    sentences = fh.read_json(dirs.get_processed_text_file())
    all_documents = sentences.keys()
    documents = list(set(all_documents))

    # create a vocabulary and fill it with the tokenized documents
    tokenized, vocab = tokenize(sentences, documents, stopwords=stopwords)

    print "Most common words in corpus:"
    most_common = vocab.most_common(50)
    most_common.sort()
    for v in most_common:
        print v

    # set vocabulary size and prune tokens
    print "Pruning vocabulary"
    vocab.prune(n_words=vocab_size)


    n_words = 0
    for k in documents:
        tokens = [t for t in tokenized[k] if t in vocab.token2index]
        n_words += len(tokens)
        tokenized[k] = tokens

    n_documents = len(documents)
    n_vocab = len(vocab)

    print n_documents, "documents"
    print n_vocab, "word types"
    print n_words, "word tokens"


    # create the count matrices
    vocab_assignments = np.zeros(n_words, dtype=int)     # vocab index of the ith word
    #topic_assignments = np.zeros(n_words, dtype=int)      # topic of the ith word
    doc_assignments = np.zeros(n_words, dtype=int)  # document of the ith word

    count = 0
    for d_i, d in enumerate(documents):
        tokens = tokenized[d]
        for t in tokens:
            v_index = vocab.get_index(t)
            assert v_index >= 0
            #w_topic = np.random.randint(n_topics)

            vocab_assignments[count] = v_index
            #topic_assignments[count] = w_topic
            doc_assignments[count] = d_i

            #topic_counts[w_topic] += 1
            #vocab_topics[v_index, w_topic] += 1
            #doc_topics[d_i, w_topic] += 1

            count += 1

    assert count == n_words

    output_filename = os.path.join(dirs.lda_dir, 'word_num.json')
    fh.write_to_json(list(vocab_assignments), output_filename, sort_keys=False)

    output_filename = os.path.join(dirs.lda_dir, 'word_doc.json')
    fh.write_to_json(list(doc_assignments), output_filename, sort_keys=False)

    output_filename = os.path.join(dirs.lda_dir, 'vocab.json')
    fh.write_to_json(vocab.index2token, output_filename, sort_keys=False)

    output_filename = os.path.join(dirs.lda_dir, 'documents.json')
    fh.write_to_json(documents, output_filename, sort_keys=False)

    # just exit after writing data


def tokenize(sentences, documents_to_tokenize, stopwords=set()):
    print "Tokenizing"
    vocab = VocabWithCounts('', add_oov=False)
    tokenized = {}
    for k in documents_to_tokenize:
        text = sentences[k].lower()
        text = re.sub('\d', '#', text)
        tokens = text.split()
        tokens = [t for t in tokens if re.search('[a-zA-Z]', t)]
        tokens = [t for t in tokens if t not in stopwords]
        vocab.add_tokens(tokens)
        tokenized[k] = tokens
    return tokenized, vocab


if __name__ == '__main__':
    main()

main