# Path: autoencoder/baseline/word2vec.py
# class Word2Vec(object):
#     def __init__(self, dim, min_count=1, sg=1, hs=0, window=5, negative=5, epoches=5):
#         super(Word2Vec, self).__init__()
#         self.dim = dim
#         self.min_count = min_count
#         self.sg = sg
#         self.hs = hs
#         self.window = window
#         self.negative = negative
#         self.epoches = epoches
# 
#     def train(self, corpus):
#         self.model = word2vec.Word2Vec(size=self.dim, min_count=self.min_count,\
#             window=self.window, workers=multiprocessing.cpu_count(), \
#             sg=self.sg, hs=self.hs, negative=self.negative, iter=self.epoches)
#         self.model.build_vocab(corpus())
#         self.model.train(corpus())
# 
#         return self
# 
# def save_w2v(model, outfile):
#     model.save_word2vec_format(outfile, binary=True)
# 
# def load_w2v(mod_file):
#     return word2vec.Word2Vec.load_word2vec_format(mod_file, binary=True)
# 
# Path: autoencoder/baseline/doc_word2vec.py
# def doc_word2vec(model, corpus, vocab, output, avg=True):
#     doc_codes = {}
#     for key, bow in corpus.iteritems():
#         vec = get_doc_codes(model, bow, vocab, avg)
#         doc_codes[key] = vec.tolist()
#     dump_json(doc_codes, output)
# 
#     return doc_codes
# 
# Path: autoencoder/utils/io_utils.py
# def load_json(file):
#     try:
#         with open(file, 'r') as datafile:
#             data = json.load(datafile)
#     except Exception as e:
#         raise e
# 
#     return data
# 
# def dump_json(data, file):
#     try:
#         with open(file, 'w') as datafile:
#             json.dump(data, datafile)
#     except Exception as e:
#         raise e
# 
# def write_file(data, file):
#     try:
#         with open(file, 'w') as datafile:
#             for line in data:
#                 datafile.write(' '.join(line) + '\n')
#     except Exception as e:
#         raise e
# 
# Path: autoencoder/preprocessing/preprocessing.py
# def load_corpus(corpus_path):
#     corpus = load_json(corpus_path)
# 
#     return corpus
# 
# Path: autoencoder/datasets/the20news.py
# class CorpusIter20News(object):
#     def __init__(self, corpus_path, recursive=False, stem=True, with_docname=False):
#         self.stem = stem
#         self.with_docname = with_docname
#         self.files = get_all_files(corpus_path, recursive)
# 
#     def __iter__(self):
#         shuffle(self.files)
#         count = 0
#         for filename in self.files:
#             try:
#                 with open(filename, 'r') as fp:
#                     text = fp.read().lower()
#                     # remove punctuations, stopwords and *unnecessary digits*, stemming
#                     words = tiny_tokenize(text, self.stem, cached_stop_words)
#                     count += 1
#                     if self.with_docname:
#                         parent_name, child_name = os.path.split(filename)
#                         doc_name = os.path.split(parent_name)[-1] + '_' + child_name
#                         yield [words, [doc_name]]
#                     else:
#                         yield words
#             except Exception as e:
#                 raise e
#         print count


import argparse
import timeit
import numpy as np
    import pdb;pdb.set_trace()
from os import path
from autoencoder.baseline.word2vec import Word2Vec, save_w2v, load_w2v
from autoencoder.baseline.doc_word2vec import doc_word2vec
from autoencoder.utils.io_utils import load_json, dump_json, write_file
from autoencoder.preprocessing.preprocessing import load_corpus
from autoencoder.datasets.the20news import CorpusIter20News
'''
from __future__ import absolute_import

# from autoencoder.datasets.reuters import CorpusIterReuters
# from autoencoder.datasets.movie_review_data import CorpusIterMRD
# from autoencoder.datasets.wiki10plus import CorpusIterWiki10plus


def train(args):
    vocab = load_json(args.vocab)
    # import pdb;pdb.set_trace()
    # load corpus
    corpus = CorpusIter20News(args.corpus[0], recursive=True, stem=True, with_docname=False)
    # corpus = CorpusIterMRD(args.corpus[0], load_json(args.docnames), stem=True, with_docname=False)
    # corpus = CorpusIterWiki10plus(args.corpus[0], load_json(args.docnames), stem=True, with_docname=False)
    # corpus = CorpusIterReuters(args.corpus, load_json(args.docnames), with_docname=False)
    # print len([1 for x in corpus])
    corpus_iter = lambda: ([word for word in sentence if word in vocab] for sentence in corpus)
    w2v = Word2Vec(args.n_dim, window=args.window_size, \
        negative=args.negative, epoches=args.n_epoch)

    start = timeit.default_timer()
    w2v.train(corpus_iter)
    print 'runtime: %ss' % (timeit.default_timer() - start)

    save_w2v(w2v.model, args.save_model)

def test(args):
    corpus = load_corpus(args.corpus[0])
    docs, vocab_dict = corpus['docs'], corpus['vocab']

# Path: autoencoder/baseline/word2vec.py
# class Word2Vec(object):
#     def __init__(self, dim, min_count=1, sg=1, hs=0, window=5, negative=5, epoches=5):
#         super(Word2Vec, self).__init__()
#         self.dim = dim
#         self.min_count = min_count
#         self.sg = sg
#         self.hs = hs
#         self.window = window
#         self.negative = negative
#         self.epoches = epoches
# 
#     def train(self, corpus):
#         self.model = word2vec.Word2Vec(size=self.dim, min_count=self.min_count,\
#             window=self.window, workers=multiprocessing.cpu_count(), \
#             sg=self.sg, hs=self.hs, negative=self.negative, iter=self.epoches)
#         self.model.build_vocab(corpus())
#         self.model.train(corpus())
# 
#         return self
# 
# def save_w2v(model, outfile):
#     model.save_word2vec_format(outfile, binary=True)
# 
# def load_w2v(mod_file):
#     return word2vec.Word2Vec.load_word2vec_format(mod_file, binary=True)
# 
# Path: autoencoder/baseline/doc_word2vec.py
# def doc_word2vec(model, corpus, vocab, output, avg=True):
#     doc_codes = {}
#     for key, bow in corpus.iteritems():
#         vec = get_doc_codes(model, bow, vocab, avg)
#         doc_codes[key] = vec.tolist()
#     dump_json(doc_codes, output)
# 
#     return doc_codes
# 
# Path: autoencoder/utils/io_utils.py
# def load_json(file):
#     try:
#         with open(file, 'r') as datafile:
#             data = json.load(datafile)
#     except Exception as e:
#         raise e
# 
#     return data
# 
# def dump_json(data, file):
#     try:
#         with open(file, 'w') as datafile:
#             json.dump(data, datafile)
#     except Exception as e:
#         raise e
# 
# def write_file(data, file):
#     try:
#         with open(file, 'w') as datafile:
#             for line in data:
#                 datafile.write(' '.join(line) + '\n')
#     except Exception as e:
#         raise e
# 
# Path: autoencoder/preprocessing/preprocessing.py
# def load_corpus(corpus_path):
#     corpus = load_json(corpus_path)
# 
#     return corpus
# 
# Path: autoencoder/datasets/the20news.py
# class CorpusIter20News(object):
#     def __init__(self, corpus_path, recursive=False, stem=True, with_docname=False):
#         self.stem = stem
#         self.with_docname = with_docname
#         self.files = get_all_files(corpus_path, recursive)
# 
#     def __iter__(self):
#         shuffle(self.files)
#         count = 0
#         for filename in self.files:
#             try:
#                 with open(filename, 'r') as fp:
#                     text = fp.read().lower()
#                     # remove punctuations, stopwords and *unnecessary digits*, stemming
#                     words = tiny_tokenize(text, self.stem, cached_stop_words)
#                     count += 1
#                     if self.with_docname:
#                         parent_name, child_name = os.path.split(filename)
#                         doc_name = os.path.split(parent_name)[-1] + '_' + child_name
#                         yield [words, [doc_name]]
#                     else:
#                         yield words
#             except Exception as e:
#                 raise e
#         print count


# Path: run_w2v.py
import argparse
import timeit
import numpy as np
    import pdb;pdb.set_trace()
from os import path
from autoencoder.baseline.word2vec import Word2Vec, save_w2v, load_w2v
from autoencoder.baseline.doc_word2vec import doc_word2vec
from autoencoder.utils.io_utils import load_json, dump_json, write_file
from autoencoder.preprocessing.preprocessing import load_corpus
from autoencoder.datasets.the20news import CorpusIter20News

'''
from __future__ import absolute_import

# from autoencoder.datasets.reuters import CorpusIterReuters
# from autoencoder.datasets.movie_review_data import CorpusIterMRD
# from autoencoder.datasets.wiki10plus import CorpusIterWiki10plus


def train(args):
    vocab = load_json(args.vocab)
    # import pdb;pdb.set_trace()
    # load corpus
    corpus = CorpusIter20News(args.corpus[0], recursive=True, stem=True, with_docname=False)
    # corpus = CorpusIterMRD(args.corpus[0], load_json(args.docnames), stem=True, with_docname=False)
    # corpus = CorpusIterWiki10plus(args.corpus[0], load_json(args.docnames), stem=True, with_docname=False)
    # corpus = CorpusIterReuters(args.corpus, load_json(args.docnames), with_docname=False)
    # print len([1 for x in corpus])
    corpus_iter = lambda: ([word for word in sentence if word in vocab] for sentence in corpus)
    w2v = Word2Vec(args.n_dim, window=args.window_size, \
        negative=args.negative, epoches=args.n_epoch)

    start = timeit.default_timer()
    w2v.train(corpus_iter)
    print 'runtime: %ss' % (timeit.default_timer() - start)

    save_w2v(w2v.model, args.save_model)

def test(args):
    corpus = load_corpus(args.corpus[0])
    docs, vocab_dict = corpus['docs'], corpus['vocab']

    doc_codes = doc_word2vec(docs, revdict(vocab_dict), args.load_model, args.output, avg=True)