import os
import sys
import json
import numpy as np
import argparse
import datetime
import cPickle
import copy
from pprint import pprint

import theano
import theano.tensor as T

from data_reader import *
import lasagne as la
from lasagne.utils import floatX
from saveable_la_model import SaveableModel

sys.path.append(os.path.abspath("./util"))
from pycocoevalcap.bleu.bleu import Bleu
from evaluator import SentenceEvaluator
from theano.sandbox.rng_mrg import MRG_RandomStreams as RandomStreams
from lasagne.random import get_rng
from tokenizer import TokenSequences, TextNormalizer, TextNormalizerOptions, TextTokenizer
from eval_romance import SentenceEvaluator

from style_remover import remove_style2, remove_unknown, VocabMap, VocabMapGlove, remove_stopwords_sents, remove_style_text

from gen_img_tags_fast import GenImgTags

from nltk.corpus import stopwords
sw = stopwords.words('english')
sw.extend(['two', 'nt', 're', 'll', ',', '.', 'm', 'd', 've', 'nah', 'one', 'UNKNOWN'])

mp = {}
ap = argparse.ArgumentParser()
ap.add_argument("-o", "--outfile", default="./models/def_model_saved", help="where to save the model file")
ap.add_argument("-j", "--json", help="where to save generated output as json")
ap.add_argument("-i", "--infile", help="the input model file")
ap.add_argument("-I", "--infile2", help="the second input model file")
ap.add_argument("--infile_js", help="the joint space input file")
ap.add_argument("--gpu", action="store_true", help="use gpu params")
ap.add_argument("--force_nondet", action="store_true", help="force the use of non-deterministic layers during generation: for debugging")
ap.add_argument("--dataset", choices = ["captions", "romance", "both", "poems", "both_poems"], default="captions")
ap.add_argument("--remove_order", action="store_true")
ap.add_argument("--num_test_batches", type=int, default=1)
ap.add_argument("--large", action="store_true")
ap.add_argument("--sample", action="store_true")
ap.add_argument("--num_sample", type=int, default=1)
ap.add_argument("--preattn", action="store_true")
ap.add_argument("--read_imp", action="store_true")
ap.add_argument("--use_js", action="store_true", help="filter output using the joint space embeddings")
ap.add_argument("--small", action="store_true")
ap.add_argument("--similar_words", choices = ["none", "glove"], default="none")
ap.add_argument("--lemmatize", action="store_true", help="remove word morphology information")
ap.add_argument("--vocabsize", type=int, default=10000)
ap.add_argument("--seqlen", type=int, default=16)
ap.add_argument("--use_pos", action="store_true", help="add pos tags to input text")
ap.add_argument("--force_feedback", action="store_true", help="when decoding force a small fraction of output to input")
ap.add_argument("--drop_in", type=float, default=0.0, help="fraction of input words to be dropped out")
ap.add_argument("--prepend_type", choices = ["captions", "romance"], default="romance")
ap.add_argument("--gen_both", action="store_true", help="Use both generators and try to weight them off")
ap.add_argument("--alternate_tradeoff", type=float, default=1.0, help="how much weight to put on the inverse of the alternate probability")
ap.add_argument("--random_removal", action="store_true", help="randomly delete a fraction of the input data, usefull for learning which words we can remove")
ap.add_argument("--extended_pos", action="store_true", help="use the extended pos list for style removal")
ap.add_argument("--extended_pos_wverbs", action="store_true", help="use the extended pos list with added verbs")
ap.add_argument("--cap_from_img", action="store_true", help="use the terms generated by the image model rather than gt")
ap.add_argument("--framenet", action="store_true", help="use framenet verb conversion")
ap.add_argument("mode", choices = ["train", "test", "test_img"])
args = ap.parse_args()

if not args.small and (os.path.exists("use_gpu.config") or args.gpu):
    mp["EMB_SIZE"] = 353
    mp["HIDDEN_SIZE"] = 192
    mp["BATCH_SIZE"] = 128
    mp["SPACY_THREADS"] = 12
    mp["SPACY_BATCH"] = 50000
    max_load = 0
else:
    mp["EMB_SIZE"] = 128
    mp["HIDDEN_SIZE"] = 64
    mp["BATCH_SIZE"] = 128
    mp["SPACY_THREADS"] = 2
    mp["SPACY_BATCH"] = 1000
    max_load = 0
    max_load = 4000
if args.large and not args.small:
    mp["EMB_SIZE"] = 512
    mp["HIDDEN_SIZE"] = 512

#max_load = 4000

mp["WORD_DROPOUT"] = 0.0
mp["WORD_DROPOUT_IN"] = args.drop_in

mp["SEQ_LEN"] = args.seqlen
mp["IN_VOCAB"] = args.vocabsize
mp["OUT_VOCAB"] = args.vocabsize
mp["FORCE_FEEDBACK"] = args.force_feedback
mp["GRAD_CLIP"] = 5.0
mp["LEARNING_RATE"] = 0.001
mp["DROPOUT_RATE"] = 0.5
mp["REGULARIZE_WEIGHT"] = 1e-6

class ZeroLayer(la.layers.Layer):
    def get_output_for(self, input, **kwargs):
        return input * la.utils.as_theano_expression(0.0)

class EncDecModel(SaveableModel):
    def __init__(self):
        super(EncDecModel, self).__init__()
        self.senteval = SentenceEvaluator()
        self.model_is_built=False

    def build(self, in_vocab_size, out_vocab_size, word_count, gum_temp=1.0):
        self.model_is_built=True

        self.in_vocab_size = in_vocab_size
        self.out_vocab_size = out_vocab_size
        self.word_count = word_count
        self.var["temp"] = theano.shared(np.float32(gum_temp), "temp")
        self.var["enc_in"] = T.imatrix()
        self.net["l_enc_in"] = la.layers.InputLayer((mp["BATCH_SIZE"], mp["SEQ_LEN"]), input_var = self.var["enc_in"])
        self.var["dec_in"] = T.imatrix()
        self.net["l_dec_in"] = la.layers.InputLayer((mp["BATCH_SIZE"], mp["SEQ_LEN"]), input_var = self.var["dec_in"])
        self.var["enc_mask_in"] = T.matrix()
        self.net["l_enc_mask_in"] = la.layers.InputLayer((mp["BATCH_SIZE"], mp["SEQ_LEN"]), input_var = self.var["enc_mask_in"])

        # -------Encoder -----
        self.net["l_enc_emb_in"] = la.layers.EmbeddingLayer(self.net["l_enc_in"], input_size=self.in_vocab_size, output_size=mp["EMB_SIZE"])

        self.net["l_enc_gru_fwd"] = la.layers.GRULayer(self.net["l_enc_emb_in"], num_units=mp["HIDDEN_SIZE"]/2, mask_input=self.net["l_enc_mask_in"], 
            grad_clipping=floatX(mp["GRAD_CLIP"]), learn_init=True, only_return_final=False)
        self.net["l_enc_gru_bak"] = la.layers.GRULayer(self.net["l_enc_emb_in"], num_units=mp["HIDDEN_SIZE"]/2, mask_input=self.net["l_enc_mask_in"], 
            grad_clipping=floatX(mp["GRAD_CLIP"]), learn_init=True, only_return_final=False, backwards=True)
        self.net["l_enc_gru_full"] = la.layers.ConcatLayer([self.net["l_enc_gru_fwd"], self.net["l_enc_gru_bak"]], axis=2)

        self.net["l_enc_gru"] = la.layers.SliceLayer(self.net["l_enc_gru_full"], mp["SEQ_LEN"]-1, axis=1)

        # ------Transform-----
        hidden_input = la.layers.dropout(self.net["l_enc_gru"], 0.3)

        # ------Decoder-------

        if args.preattn:
            self.net["l_dec_gru"] = la.layers.GRULayerESMGumPreAttn(self.net["l_dec_in"], num_units=mp["HIDDEN_SIZE"], vocab_size=self.out_vocab_size, 
                    bsm=self._process_word_count(word_count), hid_init=hidden_input, 
                    grad_clipping=floatX(mp["GRAD_CLIP"]), learn_init=True, teacher_force=True, temp=self.var["temp"],
                    encoder_hout=la.layers.dropout(self.net["l_enc_gru_full"], 0.2), enc_in_mask=self.net["l_enc_mask_in"], use_mlp_attn=True)
        else:
            self.net["l_dec_gru"] = la.layers.GRULayerESMGumAttn(self.net["l_dec_in"], num_units=mp["HIDDEN_SIZE"], vocab_size=self.out_vocab_size, 
                    bsm=self._process_word_count(word_count), hid_init=hidden_input, 
                    grad_clipping=floatX(mp["GRAD_CLIP"]), learn_init=True, teacher_force=True, temp=self.var["temp"],
                    encoder_hout=la.layers.dropout(self.net["l_enc_gru_full"], 0.2), enc_in_mask=self.net["l_enc_mask_in"], use_mlp_attn=True)

        self.net["l_dec_gru_rs"] = la.layers.ReshapeLayer(self.net["l_dec_gru"], (-1, self.out_vocab_size))
        self.net["l_out"] = la.layers.NonlinearityLayer(self.net["l_dec_gru_rs"], nonlinearity=la.nonlinearities.softmax)

        # for updating the gumbel softmax temperature
        temp_in = T.scalar(name="temp_in")
        #self.update_temp = theano.function([temp_in], self.var["temp"], updates={self.var["temp"] : temp_in})

    def build_trainer(self):
        self.var["dec_mask_out"] = T.matrix()
        self.var["dec_out"] = T.imatrix()

        to_regularize = [self.net["l_out"]]#, self.net["l_enc_dense_in"]]#, self.net["l_dec_dense_in"]]

        self.expr["out"] = la.layers.get_output(self.net["l_out"], deterministic=False)

        self.expr["reg"] = la.regularization.regularize_layer_params(to_regularize, la.regularization.l2)
        self.expr["cce"] = la.objectives.categorical_crossentropy(self.expr["out"], self.var["dec_out"].ravel())
        self.expr["cost"] = la.objectives.aggregate(self.expr["cce"], weights=self.var["dec_mask_out"].ravel(), mode="normalized_sum") 
        
        self.expr["loss"] = self.expr["cost"] + (self.expr["reg"] * mp["REGULARIZE_WEIGHT"])

        all_params = la.layers.get_all_params(self.net["l_out"], trainable=True)
        updates = la.updates.adam(self.expr["loss"], all_params, floatX(mp["LEARNING_RATE"]))

        vars_in = [self.var["enc_in"], self.var["dec_in"], self.var["enc_mask_in"], self.var["dec_out"], self.var["dec_mask_out"]]
        vars_out = [self.expr["cost"], self.expr["loss"]]

        print vars_in
        print vars_out

        train = theano.function(vars_in, vars_out, updates=updates)

        self.th_train = train
        return train

    def build_generator(self):
        self.expr["out"] = la.layers.get_output(self.net["l_out"], deterministic=True and (not args.force_nondet))

        vars_in = [self.var["enc_in"], self.var["dec_in"], self.var["enc_mask_in"]]
        generate = theano.function(vars_in, self.expr["out"])
        self.th_generate = generate
        return generate

    # -----------------------Model Run------------------

    def set_train_data(self, data):
        self.data = data

    def set_test_data(self, test_data, alternate_enc_in = None):
        self.test_data = test_data
        self.alternate_enc_in = alternate_enc_in

    def word_dropout(self, seq, mask, token, drop_frac = mp["WORD_DROPOUT"]):
        if drop_frac > 0.0:
            chosen = np.random.random(seq.shape) < drop_frac
            chosen = np.logical_and(chosen, mask > 0.5)

            # at least one word is always available
            c_count = np.sum(chosen, axis=1)
            m_count = np.sum(mask, axis=1)
            for i in xrange(c_count.shape[0]):
                if c_count[i] == m_count[i]:
                    chosen[i, 0] = 0

            seq[chosen] = token
        return seq

    def word_remove(self, seq, mask, drop_frac = mp["WORD_DROPOUT"]):
        if drop_frac > 0.0:
            chosen = np.random.random(seq.shape) < drop_frac
            chosen = np.logical_and(chosen, mask > 0.5)

            # at least one word is always available
            c_count = np.sum(chosen, axis=1)
            m_count = np.sum(mask, axis=1)
            for i in xrange(c_count.shape[0]):
                if c_count[i] == m_count[i]:
                    chosen[i, 0] = 0
                sub = 0
                for j,m in enumerate(chosen[i]):
                    if m:
                        seq[i, j-sub:-1] = seq[i, j-sub+1:]
                        seq[i, -1] = 0
                        sub += 1
        return seq


    def word_insert(self, seq, mask, token, ins_frac = 0.15):
        if ins_frac > 0.0:
            for i in xrange(seq.shape[0]):
                new_row = []
                new_row_mask = []
                for j in xrange(seq.shape[1]):
                    if np.random.rand() < ins_frac:
                        new_row.append(token)
                        new_row_mask.append(mask[i,j])
                    new_row.append(seq[i,j])
                    new_row_mask.append(mask[i,j])
                seq[i, :] = np.array(new_row, dtype=np.int32)[:seq.shape[1]]
                mask[i, :] = np.array(new_row_mask, dtype=theano.config.floatX)[:mask.shape[1]]
        return seq, mask
    
    def train(self, data, save_name = None, save_object = {}):
        num_examples = data.dec_in.get_num_seqs()
        num_batches = num_examples / mp["BATCH_SIZE"]
        smooth_loss = 0

        if args.similar_words == "none": 
            vmap = VocabMap(data.enc_in.tokenizer) 
        elif args.similar_words == "glove":
            vmap = VocabMapGlove(data.enc_in.tokenizer, temp=0.5)

        for epoch in xrange(20):
            data.shuffle()
            #self.generate()
            self.val_score(num_batches=20, epoch=epoch)
            for i in xrange(num_batches):
                s = i * mp["BATCH_SIZE"]
                e = (i+1) * mp["BATCH_SIZE"]
                enc_in = np.array(data.enc_in.get_padded_seqs(s, e))
                enc_in_mask = np.array(data.enc_in.get_masks(s, e, dtype=theano.config.floatX))
                if args.random_removal:
                    enc_in = self.word_remove(enc_in, enc_in_mask, drop_frac=0.7)
                #enc_in = self.word_dropout(enc_in, enc_in_mask, 
                #        token = data.enc_in.tokenizer.special_token_index[WORD_DROP_TOKEN], drop_frac=mp["WORD_DROPOUT_IN"])

                enc_in = vmap.map_tokens(enc_in)

                dec_in = np.array(data.dec_in.get_padded_seqs(s, e))
                dec_out = np.array(data.dec_out.get_padded_seqs(s, e))
                dec_out_mask = np.array(data.dec_out.get_masks(s, e, dtype=theano.config.floatX))

                dec_in_mask = data.dec_in.get_masks(s, e, dtype=theano.config.floatX)
                #if args.random_removal:
                    #dec_in = self.word_remove(dec_in, dec_in_mask, 0.7)
                #dec_in = self.word_dropout(dec_in, dec_in_mask, 
                #        token = data.dec_in.tokenizer.special_token_index[WORD_DROP_TOKEN])

                #rand_teacher_force = np.random.randint(2, size=dec_in.shape, dtype=np.int32)*2 - 1
                if epoch == 0:
                    force_prob = 0.99
                elif epoch == 1:
                    force_prob = 0.9
                elif epoch == 2 or epoch == 3:
                    force_prob = 0.8
                else:
                    force_prob = 0.8
                    
                if mp["FORCE_FEEDBACK"]:
                    rand_teacher_force = np.array(np.random.choice([-1, 1], size=dec_in.shape, p = np.array([1.0 - force_prob, force_prob])), dtype=np.int32)
                else:
                    rand_teacher_force = np.ones(dec_in.shape, dtype=np.int32)
                rvals = self.th_train(enc_in, dec_in * rand_teacher_force, enc_in_mask, dec_out, dec_out_mask)
                cur_loss = rvals[0]

                if epoch == 0 and i == 0:
                    smooth_loss = cur_loss
                else:
                    smooth_loss = smooth_loss * 0.95 + cur_loss*0.05

                epoch_frac = epoch + (i / float(num_batches))

                print rvals
                print "Epoch %f, Loss %f, Smooth Loss %f" % (epoch_frac, cur_loss, smooth_loss)
                print ""

                if i % 1000 == 0:
                    self.generate()
                    self.val_score(num_batches=4, epoch=epoch_frac)
                    save_object = self.save_to_object(save_object)
                    filename = save_name + "%02d_e.pik" % epoch
                    cPickle.dump(save_object, open(filename, "wb"), protocol=2)
                    temp = np.float32(max(0.5, np.exp(-0.07*epoch_frac)))
                    #print "Udated temp, from:", self.update_temp(temp), "to:", temp

            if save_name is not None:
                save_object = self.save_to_object(save_object)
                filename = save_name + "%02d_e.pik" % epoch
                cPickle.dump(save_object, open(filename, "wb"), protocol=2)

    def val_score(self, s_start = 0, num_batches=2, epoch = 0.0):
        bs = mp["BATCH_SIZE"]
        self.senteval.clear()
        num_examples = self.test_data.dec_in.get_num_seqs()
        max_num_batches = num_examples / bs
        for i in xrange(min(num_batches, max_num_batches)):
            s = s_start+bs*i
            e = s_start+bs*(i+1)
            in_all, gt_txt, gen_txt, fnames = self.generate(s=s, allow_unk=False)
            for g, f in zip(gen_txt, fnames):
                self.senteval.add_gen_sent(f, " ".join(g))
                #print g
                #sys.exit(0)

            for g, f in zip(gt_txt, fnames):
                self.senteval.add_gt_sent(f," ".join(g))
        print "Epoch %0.4f BS: %d LM_Score: %0.4f" % (epoch, num_batches, self.senteval.get_lm_score())
        print "Epoch %0.4f BS: %d Bleu_Score: %s" % (epoch, num_batches, str(self.senteval.get_bleu_score()))


    def _process_word_count(self, word_count):
        wc = np.array(word_count, dtype=theano.config.floatX)
        wc = np.log(wc/np.sum(wc))
        wc -= np.median(wc)
        return wc

    def generate(self, s = 0, allow_unk=True, enc_in=None, enc_in_mask=None, enc_txt=None, filenames=None, 
            no_feedback=False, sample=False, force_good = False, word_insert=False):
        bs = mp["BATCH_SIZE"]
        e = s + bs
        sw_idx = []
        sw_local = stopwords.words('english')
        sw_local.extend(['nt', 're', 'll', 'm', 'd', 've', 'nah', 'UNKNOWN'])
        for w in sw_local:
            if w in self.test_data.dec_in.tokenizer.word_to_index:
                sw_idx.append(self.test_data.dec_in.tokenizer.word_to_index[w])
        sw_idx = np.array(sw_idx)

        if enc_in is None:
            enc_in = self.test_data.enc_in.get_padded_seqs(s,e)
            if self.alternate_enc_in is not None:
                alternate_enc_in = self.alternate_enc_in.get_padded_seqs(s,e)

        if enc_in_mask is None:
            enc_in_mask = self.test_data.enc_in.get_masks(s,e, dtype=theano.config.floatX)
        #enc_in = self.word_dropout(enc_in, enc_in_mask, 
        #        token = self.test_data.enc_in.tokenizer.special_token_index[WORD_DROP_TOKEN], drop_frac=mp["WORD_DROPOUT_IN"])
        if word_insert:
            enc_in, enc_in_mask = self.word_insert(enc_in, enc_in_mask, token = self.test_data.enc_in.tokenizer.special_token_index[WORD_DROP_TOKEN])
        dec_in = np.zeros_like(self.test_data.dec_in.get_padded_seqs(s,e))
        for i in xrange(mp["SEQ_LEN"]):
            if no_feedback:
                dec_in_tmp = np.ones_like(dec_in) * self.test_data.dec_in.tokenizer.special_token_index[WORD_DROP_TOKEN]
            else:
                dec_in_tmp = dec_in

            res = self.th_generate(enc_in, dec_in_tmp, enc_in_mask)
            res = res.reshape((mp["BATCH_SIZE"], mp["SEQ_LEN"], self.out_vocab_size))
            if self.alternate_enc_in:
                res_alt = self.th_generate(alternate_enc_in, dec_in_tmp, enc_in_mask)
                res_alt = res_alt.reshape((mp["BATCH_SIZE"], mp["SEQ_LEN"], self.out_vocab_size))

                #print alternate_enc_in[:5]
                #print enc_in[:5]
                #sys.exit(0)
                #res[:, i] = res_alt[:, i]
                #print np.argmax(re
                #print res[:, i]
                #print "Argsort:", np.argsort(res[:, i])[:10]
                #print "Res:", res[:, i, np.argsort(-res[:, i])[:10]]
                #print "sum:", np.sum(res[:, i], axis=1, keepdims=True)
                #res[:, i] = res[:, i] * ((1.0-res_alt[:, i])**10.0)
                mul = -1
                if args.alternate_tradeoff < 0.0:
                    mul = 1
                chng = np.argmax(res[:, i], axis=1) != 0
                res[chng, i] = res[chng, i] + (1.0 + mul*res_alt[chng, i])*args.alternate_tradeoff
                res[chng, i] = res[chng, i] / np.sum(res[chng, i], axis=1, keepdims=True)

            if i+1 < mp["SEQ_LEN"]:
                if allow_unk == False:
                    unk_id = self.test_data.dec_out.tokenizer.special_token_index["UNKNOWN"]
                    res[:, i, unk_id] = 0.0
                if force_good:
                    res[:, :4, 0] = 0.0 # must be at least this long
                    if i > 0: 
                        # ignore duplicate non-stopwords
                        seen = dec_in[:, :(i+1)]
                        res_saved = np.array(res[:, i, :], dtype=theano.config.floatX)
                        print seen.shape
                        print res[:, i, :][np.arange(res.shape[0]), seen.T].shape
                        #res[:, i, seen] = 0
                        res[:, i, :][np.arange(res.shape[0]), seen.T] = 0
                        res[:, i, sw_idx] = res_saved[:, sw_idx]
                        # ignore duplicate adjacent words of any kind
                        res[:, i, dec_in[:, i]] = 0
                        res[:, i, 0] = res_saved[:, 0]
                if sample:
                    for j in xrange(dec_in.shape[0]):
                        temp = 0.5
                        res[j, i] = res[j, i] ** (1.0/temp)
                        res[j, i] /= np.sum(res[j, i])
                        dec_in[j, i+1] = np.random.choice(res.shape[2], p=res[j, i])
                else:
                    dec_in[:, i+1] = np.argmax(res[:, i], axis=1)

        gen_seq = TokenSequences(self.test_data.dec_out.tokenizer, reverse=False, start_pad=True, seq_len = mp["SEQ_LEN"])
        gen_seq.from_padded_seqs(dec_in, is_reversed=False, is_start_pad=True)
        if enc_txt is None:
            enc_txt = self.test_data.enc_in.get_text(s, e)
            dec_txt = self.test_data.dec_out.get_text(s, e)
            if self.test_data.filenames is not None:
                filenames = self.test_data.filenames[s:e]
        if filenames is None:
            filenames = ["" for t in xrange(bs)]
        
        in_all = []
        gt_all = []
        out_all = []
        for txt_in, txt_gt, txt_out in zip(enc_txt, dec_txt, gen_seq.get_text()):
            print "IN :", txt_in
            print "GT :", txt_gt
            print "OUT:", txt_out
            print ""
            in_all.append(" ".join(txt_in))
            gt_all.append(" ".join(txt_gt))
            out_all.append(" ".join(txt_out))

        return in_all, gt_all, out_all, filenames

"""def remove_style(data):
    data.enc_in.seqs_pad = None
    seqs = data.enc_in.seqs
    sw_tok = data.enc_in.tokenizer.text_to_seqs([sw])[0]
    sw_tok = filter(lambda x: x != data.enc_in.tokenizer.special_token_index["UNKNOWN"], sw_tok)
    sw_bs = np.zeros((data.enc_in.get_vocab_size(),), dtype=np.bool)
    sw_bs[sw_tok] = 1
    for i in xrange(len(seqs)):
        nseq = []
        for j in xrange(len(seqs[i])):
            if not sw_bs[seqs[i][j]]:
                nseq.append(seqs[i][j])
        seqs[i] = nseq
    data.enc_in.from_seqs(seqs)

    #print data.enc_in.get_text(0, 50)
    #sys.exit(0)

    return data"""

def train(args):
    #data = read_captions_images(mp["IN_VOCAB"], mp["OUT_VOCAB"], mp["SEQ_LEN"], val=True)
            #input_tokenizer=data_tmp.enc_in.tokenizer, output_tokenizer=data_tmp.dec_out.tokenizer)

    pos_list = ["ADJ", "PRT", "PUNCT", "ADV", "PRON", "CONJ"]

    if args.extended_pos:
        pos_list.extend(["ADP", "VERB", "X", "NUM", "DET", "PROPN"])
    if args.extended_pos_wverbs:
        pos_list.extend(["ADP", "X", "NUM", "DET", "PROPN"])
        

    style_remove_func = lambda x : remove_style_text(x, threads=mp["SPACY_THREADS"], batch_size=mp["SPACY_BATCH"], lemmatize=args.lemmatize, remove_order=args.remove_order, append_pos=args.use_pos, pos_list = pos_list, framenet=args.framenet)
    if args.random_removal:
        style_remove_func = lambda x : x

    prepend_type = None
    if args.dataset == "captions":
        data = read_captions_images(mp["IN_VOCAB"], mp["OUT_VOCAB"], mp["SEQ_LEN"], max_load=max_load, val=False, get_feats=False, read_importance=args.read_imp, full_dataset=True, style_remove_func=style_remove_func, rm_style_in=True)
    elif args.dataset == "romance":
        if args.read_imp:
            data = read_romance_new_importance(mp["IN_VOCAB"], mp["OUT_VOCAB"], mp["SEQ_LEN"], max_load=max_load)
        else:
            data = read_romance_new(mp["IN_VOCAB"], mp["OUT_VOCAB"], mp["SEQ_LEN"], max_load=max_load,
                    style_remove_func=style_remove_func)

        print data.enc_in.get_text(0,200)
        #sys.exit(0)
    elif args.dataset == "both":
        data = read_cap_and_rom(mp["IN_VOCAB"], mp["OUT_VOCAB"], mp["SEQ_LEN"], max_load=max_load,
                    style_remove_func=style_remove_func, shuffle=False)
        if args.cap_from_img:
            all_in_text = data.enc_in.get_text()
            img_txt = cPickle.load(open("./data/pre_gen_img_txt.pik", "rb"))
            tlen = len(all_in_text)
            cap_idx = 0
            for i in xrange(tlen):
                if all_in_text[i][0] == "CAPTIONDATASET":
                    all_in_text[i] = ["CAPTIONDATASET"] + img_txt[cap_idx]
                    cap_idx += 1
            data.enc_in.from_text(all_in_text)

        data.shuffle()
        #all_in_text = data.enc_in.get_text()
        #tlen = len(all_in_text)
        #print data.enc_in.get_text(tlen-100,tlen-1)
        #sys.exit(0)
        prepend_type = "ROMANCEDATASET"
    elif args.dataset == "poems":
        data = read_romance_new(mp["IN_VOCAB"], mp["OUT_VOCAB"], mp["SEQ_LEN"], max_load=max_load,
                style_remove_func=style_remove_func, rom_filename="./data/poems_10000.txt")
    elif args.dataset == "both_poems":
        data = read_cap_and_rom(mp["IN_VOCAB"], mp["OUT_VOCAB"], mp["SEQ_LEN"], max_load=max_load,
                    style_remove_func=style_remove_func, rom_filename="./data/poems_10000.txt")
        prepend_type = "ROMANCEDATASET"
    #elif args.dataset == "stylenet":
    #    data = 
    #print "WARNGING: not removing style"
    #remove_style2(data, threads=mp["SPACY_THREADS"], batch_size=mp["SPACY_BATCH"], remove_order=args.remove_order)

    model = EncDecModel()

    data_test = read_captions_images(mp["IN_VOCAB"], mp["OUT_VOCAB"], mp["SEQ_LEN"], max_load=max_load, val=True, get_feats=False, read_importance=args.read_imp, full_dataset=True,
            input_tokenizer = data.enc_in.tokenizer, output_tokenizer = data.dec_in.tokenizer, style_remove_func=style_remove_func, rm_style_in=True, prepend_type=prepend_type)
    #remove_style2(data_test, threads=mp["SPACY_THREADS"], batch_size=mp["SPACY_BATCH"], remove_order=args.remove_order)
    model.set_test_data(data_test)
    model.build(in_vocab_size = data.enc_in.get_vocab_size(), out_vocab_size=data.dec_out.get_vocab_size(), 
            word_count=data.dec_out.tokenizer.get_word_count())
    model.build_trainer()
    model.build_generator()

    save_object = {}
    save_object = data.save_tokenizer(save_object)
    model.train(data, save_name = args.outfile, save_object = save_object)

def test(args):
    save_model = cPickle.load(open(args.infile, "rb"))

    if args.dataset == "captions":
        data = read_captions_images(mp["IN_VOCAB"], mp["OUT_VOCAB"], mp["SEQ_LEN"], max_load=10000,
        input_tokenizer = save_model['input_tokenizer'], output_tokenizer = save_model['output_tokenizer'], val=True,
        get_feats=False, full_dataset=True)
        data.shuffle()
    elif args.dataset == "romance":
        data = read_romance_new(mp["IN_VOCAB"], mp["OUT_VOCAB"], mp["SEQ_LEN"], max_load=2000,
        input_tokenizer = save_model['input_tokenizer'], output_tokenizer = save_model['output_tokenizer'])
    remove_style2(data, threads=mp["SPACY_THREADS"], batch_size=mp["SPACY_BATCH"])
    data.expand(args.num_sample)
    #data = read_captions_images(mp["IN_VOCAB"], mp["OUT_VOCAB"], mp["SEQ_LEN"],
    #    input_tokenizer = save_model['input_tokenizer'], output_tokenizer = save_model['output_tokenizer'], val=True)
    #data = read_romance(mp["IN_VOCAB"], mp["OUT_VOCAB"], mp["SEQ_LEN"], 
    #    input_tokenizer = save_model['input_tokenizer'], output_tokenizer = save_model['output_tokenizer'])

    model = EncDecModel()
    model.set_test_data(data)
    model.build(in_vocab_size = data.dec_in.get_vocab_size(), out_vocab_size=data.dec_out.get_vocab_size(), 
            word_count=data.dec_out.tokenizer.get_word_count())
    model.build_generator()
    model.load_from_object(save_model)
    txt_in, txt_gt, txt_out, filename = model.generate(allow_unk=False, sample=args.sample)
    if args.json is not None:
        output = [{"in": i, "gt" : g, "out" : o, "fname" : f} 
                for i, g, o, f in zip(txt_in, txt_gt, txt_out, filename)]
        json.dump(output, open(args.json, "w"))

def test_img(args, batches = 1):

    fake_args = copy.deepcopy(args)
    fake_args.infile = fake_args.infile2
    fake_args.infile2 = None

    git = GenImgTags(fake_args)
    git.setup_test_from_input()

    save_model = cPickle.load(open(args.infile, "rb"))
    prepend_type=None
    if args.dataset == "captions":
        data = read_captions_images(mp["IN_VOCAB"], mp["OUT_VOCAB"], mp["SEQ_LEN"], max_load=0,
        input_tokenizer = save_model['input_tokenizer'], output_tokenizer = save_model['output_tokenizer'], val=True,
        get_feats=True, full_dataset=True)
        data.shuffle(seed=123)
    elif args.dataset == "romance":
        data = read_romance_new(mp["IN_VOCAB"], mp["OUT_VOCAB"], mp["SEQ_LEN"], max_load=8000,
        input_tokenizer = save_model['input_tokenizer'], output_tokenizer = save_model['output_tokenizer'])
    elif args.dataset == "both":
        prepend_type = "ROMANCEDATASET"
        if args.prepend_type == "romance":
            prepend_type = "ROMANCEDATASET"
            other_prepend_type = "CAPTIONDATASET"
        else:
            prepend_type = "CAPTIONDATASET" 
            other_prepend_tpe = "ROMANCEDATASET"
        data = read_captions_images(mp["IN_VOCAB"], mp["OUT_VOCAB"], mp["SEQ_LEN"], max_load=0,
        input_tokenizer = save_model['input_tokenizer'], output_tokenizer = save_model['output_tokenizer'], val=True,
        get_feats=True, full_dataset=True, prepend_type=None)
        data.shuffle(seed=123)
        
    all_text = []
    cur_s = 0
    num_ex = data.dec_in.get_num_seqs()
    max_batches = num_ex / mp["BATCH_SIZE"]
    if num_ex % mp["BATCH_SIZE"] != 0:
        max_batches+=1
        
    for i in xrange(min(batches, max_batches)):
        cur_s_tmp = cur_s
        if cur_s+mp["BATCH_SIZE"] > num_ex:
            cur_s_tmp = num_ex - mp["BATCH_SIZE"]
        gen_text = git.test_from_input(data, s=cur_s_tmp)
        if cur_s_tmp != cur_s:
            gen_text = gen_text[(mp["BATCH_SIZE"] - (num_ex - cur_s)):]
        cur_s += mp["BATCH_SIZE"]
        all_text.extend(gen_text)
    print gen_text
    #new_text = [" ".join(t) for t in gen_text]

    enc_in_alternate = None
    if args.gen_both:
        all_text_alternate = [[other_prepend_type] + sent for sent in all_text]
        enc_in_alternate = copy.deepcopy(data.enc_in)
        enc_in_alternate.from_text(all_text_alternate)
        #print all_text_alternate[:10]

    if prepend_type is not None:
        all_text_new = [[prepend_type] + sent for sent in all_text]
        all_text = all_text_new
        #print all_text_new[:10]
        #sys.exit(0)

    data.enc_in.from_text(all_text)
    if args.num_sample > 1:
        data.expand(args.num_sample)

    np.random.seed(1236)
    data = remove_unknown(data, encoder=True, remove_order=args.remove_order)

    if args.use_js: 
        from gen_img_tags_js import JointSpaceEmbedder
        args_new = copy.deepcopy(args)
        args_new.infile = args_new.infile_js
        js = JointSpaceEmbedder(args_new)
        js.init_getter()

    model = EncDecModel()
    model.set_test_data(data, enc_in_alternate)
    model.build(in_vocab_size = data.dec_in.get_vocab_size(), out_vocab_size=data.dec_out.get_vocab_size(), 
            word_count=data.dec_out.tokenizer.get_word_count())
    model.build_generator()
    model.load_from_object(save_model)
    output_all = []
    #seval = SentenceEvaluator()
    eval_store_gen = {}
    eval_store_gt = {}
    cur_s = 0
    num_ex *= args.num_sample
    batches *= args.num_sample
    max_batches *= args.num_sample
    all_sim = []
    for i in xrange(min(batches, max_batches)):
        cur_s_tmp = cur_s
        if cur_s+mp["BATCH_SIZE"] > num_ex:
            cur_s_tmp = num_ex - mp["BATCH_SIZE"]

        print "Generating from %d" % cur_s
        txt_in, txt_gt, txt_out, filename = model.generate(allow_unk=False, sample=args.sample, word_insert=False, force_good=False, s = cur_s_tmp)
        if cur_s_tmp != cur_s:
            txt_in = txt_in[(mp["BATCH_SIZE"] - (num_ex - cur_s)):]
            txt_gt = txt_gt[(mp["BATCH_SIZE"] - (num_ex - cur_s)):]
            txt_out = txt_out[(mp["BATCH_SIZE"] - (num_ex - cur_s)):]
            filename = filename[(mp["BATCH_SIZE"] - (num_ex - cur_s)):]

        #for gen, gt in zip(txt_out, txt_gt):
        #    seval.add_sentence_pair(" ".join(gen), " ".join(gt))
        for gen, gt, fname in zip(txt_out, txt_gt, filename):
            if fname in eval_store_gt:
                to_add = " ".join(gt)
                #if eval_store_gt[fname] == 5:
                #    print "Got here:", fname
                #    continue
                eval_store_gt[fname].append(to_add)
            else:
                eval_store_gt[fname] = [" ".join(gt).replace("UNKNOWN", "UNK").replace("NUMBER", "NUM")]
                eval_store_gen[fname] = [" ".join(gen)]
        if args.use_js:
            split_sent = remove_stopwords_sents([s.split() for s in txt_out]) 
            txt_out_tmp = [" ".join(s) for s in split_sent]
            cur_img_data = model.test_data.get_img_emb(cur_s, cur_s + len(txt_out_tmp))
            words_to_sim = [t.split() for j,t in enumerate(txt_out_tmp) ]
            #print zip(words_to_sim, txt_in)
            #sys.exit(0)
            sim = js.get_batch_score(words_to_sim, cur_img_data)
            all_sim.extend(sim)
            #print sim
        # save the output for later dumping to file
        output = [{"in": i, "gt" : g, "out" : o, "fname" : f} 
                for i, g, o, f in zip(txt_in, txt_gt, txt_out, filename)]
        output_all.extend(output)

        cur_s += mp["BATCH_SIZE"]
    
    all_scores = []
    if args.use_js and args.num_sample > 1:
        all_sim = np.array(all_sim)
        num_groups = len(output_all) / args.num_sample
        new_output = []
        for g in xrange(num_groups):
            s = g*args.num_sample
            e = (g+1)*args.num_sample
            if g == 0:
                print all_sim

            # sort order last key, then second last ....
            scores = [(-np.mean(v), -len(output_all[s+i]['out'].split()), -np.amin(v))\
                    for i, v in enumerate(all_sim[s:e])]
            all_scores.extend(scores)
            order = np.lexsort(zip(*scores))
            #print order
            new_output.append(output_all[s+order[0]])
            #jfor o in order:
            #    new_output.append(output_all[s+o])
            #    new_output[-1]['scores'] = scores[o]
        pprint(new_output)

        order = np.lexsort(zip(*all_scores))
        print "Best scoring:"
        for o in order[:10]:
            print all_scores[o]
            print output_all[o]

        print "\nWorst scoring:"
        for o in order[-10:]:
            print all_scores[o]
            print output_all[o]

        output_all = new_output
        
    print Bleu().compute_score(eval_store_gt, eval_store_gen)[0]
    #print seval

    # dump the first set to json
    if args.json is not None:
        json.dump(output_all, open(args.json, "w"))


def main():
    if args.mode == "train":
        train(args)
    elif args.mode == "test":
        test(args)
    elif args.mode == "test_img":
        test_img(args, args.num_test_batches)

if __name__ == "__main__":
    main()
