import argparse
import datetime
from os import makedirs, getcwd
from os.path import isdir, exists, join
from project import corpus, knn_corpus, lda_corpus, word2vec_corpus

algorithms = {"lda": lda_corpus.LDACorpus,
              "knn": knn_corpus.KNNCorpus,
              "w2v": word2vec_corpus.W2VCorpus}

base_dir = getcwd()
output_loc = base_dir + "/%s.corpus_out"
dictionary_loc = output_loc + "/%scorpus.dict"
corpus_loc = output_loc + "/%scorpus.mm"
log_file = output_loc + "/runtimes.log"
sup_file_loc = output_loc + "/%d.%s"


def get_file_distributions(directory):
    distributions = dict()
    total = 0
    log = join(directory, "paperstats.log")
    with open(log) as file:
        for line in file:
            name, count = line.split("\t")
            distributions[name] = float(count)
            total = total + float(count)
    distributions["total"] = total
    return distributions


def run_algo(directory, ints, algorithm):
    # Build corpus from largest int and directory (Check valid directory)
    max_corpus = max(ints)
    output_dir = output_loc % algorithm
    if not exists(output_dir):
        makedirs(output_dir)
    log = open(log_file % algorithm, 'a+')
    distributions = get_file_distributions(directory)

    # For each int in the param list then apply the corpus algorithm using a sliced corpus
    for size in ints:
        corpus_file = corpus_loc % (algorithm, "%d_" % size)
        corp_dict = dictionary_loc % (algorithm, "%d_" % size)
        start_time = datetime.datetime.now()
        if algorithm != "w2v":
            # Build corpus of size max_corpus and save to be reused
            current_corpus = corpus.Corpus(directory=directory, max_docs=max_corpus, distributions=distributions)
            current_corpus.save(corp_dict, corpus_file)
            directory = None
        end_time = datetime.datetime.now()
        current_corpus_build_time = end_time - start_time
        log.write("%s_corpus_build_time %d:\t%s\n" % (algorithm, size, current_corpus_build_time))
        test_corpus = algorithms[algorithm](directory=directory, dictionary=corp_dict, corpus=corpus_file, max_docs=size)

        # Save any mm, index files, etc to a directory so they can be used again.
        dict_loc = dictionary_loc % (algorithm, size)
        corp_loc = corpus_loc % (algorithm, size)
        sup_loc = sup_file_loc % (algorithm, size, algorithm)
        test_corpus.save(dictionary_file=dict_loc, corpus_file=corp_loc, sup_file=sup_loc)

        # Log temporal time
        log.write("%s %d train time:\t" % (algorithm, size) + str(test_corpus.get_train_time()) + "\n")
    log.close()


def main():
    parser = argparse.ArgumentParser(description='Build bow corpus on the arxiv corpus.')
    parser.add_argument('integers', metavar='N', type=int, nargs='+', help='size values for the corpus')
    parser.add_argument('directory', help='directory for the arxiv txt files')
    parser.add_argument('algorithm', help='algorithm to apply to the corpus', choices=algorithms)
    args = parser.parse_args()
    if isdir(args.directory):
        run_algo(args.directory, args.integers, args.algorithm)
    else:
        print "Directory argument should be a valid directory"


if __name__ == "__main__":
    main()

@@ -450,16 +450,56 @@
 pus.mm%22%0A
+file_logs = output_loc + %22/%25sfiles.log%22%0A
 log_file
@@ -1394,16 +1394,50 @@
 n ints:%0A
+        new_directory = directory%0A
         
@@ -1544,32 +1544,143 @@
 , %22%25d_%22 %25 size)%0A
+        sup_loc = sup_file_loc %25 (algorithm, size, algorithm)%0A        file_log = file_logs %25 (algorithm, size)%0A
         start_ti
@@ -1880,26 +1880,20 @@
 ax_docs=
-max_corpus
+size
 , distri
@@ -1947,16 +1947,32 @@
 us.save(
+dictionary_file=
 corp_dic
@@ -1985,16 +1985,47 @@
 pus_file
+=corpus_file, sup_file=file_log
 )%0A      
@@ -2030,16 +2030,20 @@
         
+new_
 director
@@ -2297,32 +2297,36 @@
 ithm%5D(directory=
+new_
 directory, dicti
@@ -2379,348 +2379,37 @@
 size
-)%0A%0A        # Save any mm, index files, etc to a directory so they can be used again.%0A        dict_loc = dictionary_loc %25 (algorithm, size)%0A        corp_loc = corpus_loc %25 (algorithm, size)%0A        sup_loc = sup_file_loc %25 (algorithm, size, algorithm)%0A        test_corpus.save(dictionary_file=dict_loc, corpus_file=corp_loc, sup_file=sup_loc
+, distributions=distributions
 )%0A%0A 

Fix error where directory var was overwritten