import argparse

import lockfile

from daemon import DaemonContext

from gensim.models import Doc2Vec

from gensim.models.doc2vec import LabeledSentence

from loader import get_data

class LabeledDocIterator(object):

    def __init__(self, patient_list, categories, status):

        self.patient_list = patient_list 

        self.category = categories

        self.status = status

    def __iter__(self):

        for i in self.patient_list:

            p = get_data([i])[0]

            self.status.write(p['NEW_EMPI'] + '\n')

            for category in categories:

                if category in p:

                    for idx, doc in enumerate(p[category]):

                        tag = p['NEW_EMPI'] + '_' + category + '_' + str(idx) + '\n'

                        yield LabeledSentence(words=doc['free_text'].split(), tags=[tag])

def train_doc2vec_model(categories, n_patients, output_file, status_file, dm):

    with open(status_file, 'w') as status:

        it = LabeledDocIterator(range(n_patients), categories, status)

        model = Doc2Vec(size=300, window=10, dm=dm, min_count=5, workers=11,alpha=0.025, min_alpha=0.025) # use fixed learning rate

        model.build_vocab(it)

        for epoch in range(10):

            message = ("***********Training Epoch: " + str(epoch)

                       + ("***********") + '\n')

            print(message)

            status.write(message)

            model.train(it)

            model.alpha -= 0.002 # decrease the learning rate

            model.min_alpha = model.alpha # fix the learning rate, no decay

            model.train(it)

        # Save the model

        model.save(output_file)

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument("output_file")

    parser.add_argument("n_patients")

    parser.add_argument("categories")

    # Switches between Distributed Memory and Distributed Bag of Words Model

    parser.add_argument("dm")

    args = parser.parse_args()

    status_file = args.output_file + '.status'

    categories = args.categories.split(',')

    base = '/home/ubuntu/josh_project'

    context = DaemonContext(

        working_directory=base,

        umask=0o002,

        pidfile=lockfile.FileLock(base + 'doc2vec_trainer.pid'),

    )

    with context:

        train_doc2vec_model(categories, int(args.n_patients), 

                        args.output_file, status_file, int(args.dm))