import argparse

import json

import time

import datetime

import numpy as np

import code

import socket

import os

import theano

from theano import config

import theano.tensor as tensor

import cPickle as pickle

from imagernn.data_provider import getDataProvider, prepare_data

from imagernn.cnn_evaluatorTheano import CnnEvaluator

from imagernn.solver import Solver

from imagernn.imagernn_utils import decodeEvaluator, decodeGenerator, eval_split_theano

#from numbapro import cuda

from imagernn.utils import numpy_floatX, zipp, unzip, preProBuildWordVocab

from collections import defaultdict, OrderedDict

import signal

import sys

def dumpCheckpoint(filename, params, modelGen, modelEval, misc, it, val_ppl2):

print e

def main(params):

dumpCheckpoint(filename, params, modelGen, modelEval, misc, it, val_ppl2)

if __name__ == "__main__":

parser = argparse.ArgumentParser()

# global setup settings, and checkpoints

parser.add_argument('-d', '--dataset', dest='dataset', default='coco', help='dataset: flickr8k/flickr30k')

parser.add_argument('--fappend', dest='fappend', type=str, default='baseline', help='append this string to checkpoint filenames')

parser.add_argument('-o', '--checkpoint_output_directory', dest='checkpoint_output_directory', type=str, default='cv/', help='output directory to write checkpoints to')

parser.add_argument('--worker_status_output_directory', dest='worker_status_output_directory', type=str, default='status/', help='directory to write worker status JSON blobs to')

parser.add_argument('--write_checkpoint_ppl_threshold', dest='write_checkpoint_ppl_threshold', type=float, default=-1, help='ppl threshold above which we dont bother writing a checkpoint to save space')

parser.add_argument('--continue_training', dest='checkpoint_file_name', type=str, default='None', help='checkpoint file from which to resume training')

parser.add_argument('--use_pos_tag', dest='use_pos_tag', type=str, default='None', help='use_pos_tag')

# Some parameters about image features used

parser.add_argument('--feature_file', dest='feature_file', type=str, default='vgg_feats.mat', help='Which file should we use for read the CNN features')

parser.add_argument('--image_feat_size', dest='image_feat_size', type=int, default=4096, help='size of the input image features')

parser.add_argument('--data_file', dest='data_file', type=str, default='dataset.json', help='Which dataset file shpuld we use')

parser.add_argument('--mat_new_ver', dest='mat_new_ver', type=int, default=-1, help='If the .mat feature files are saved with new version (compressed) set this flag to 1')

parser.add_argument('--aux_inp_file', dest='aux_inp_file', type=str, default='None', help='Is there any auxillary inputs ? If yes indicate file here')

parser.add_argument('--swap_AuxFeat', dest='swap_aux', type=int, default=0, help='Feed image features through auxillary input!')

parser.add_argument('--advers_gen', dest='advers_gen', type=str, default=1, help='Should we use adverserial generator!')

parser.add_argument('--eval_model', dest='eval_model', type=str, default='cnn', help='which evaluator model to use type: cnn/lstm_eval')

# model parameters

parser.add_argument('--image_encoding_size', dest='image_encoding_size', type=int, default=512, help='size of the image encoding')

parser.add_argument('--word_encoding_size', dest='word_encoding_size', type=int, default=512, help='size of word encoding')

parser.add_argument('--hidden_size', dest='hidden_size', type=int, default=512, help='size of hidden layer in generator RNNs')

parser.add_argument('--hidden_depth', dest='hidden_depth', type=int, default=1, help='depth of hidden layer in generator RNNs')

parser.add_argument('--generator', dest='generator', type=str, default='lstm', help='generator to use')

parser.add_argument('-c', '--regc', dest='regc', type=float, default=1e-8, help='regularization strength')

parser.add_argument('--tanhC_version', dest='tanhC_version', type=int, default=0, help='use tanh version of LSTM?')

parser.add_argument('--beam_size', type=int, default=1, help='beam size in inference. 1 indicates greedy per-word max procedure. Good value is approx 20 or so, and more = better.')

parser.add_argument('--sent_encoding_size', dest='sent_encoding_size', type=int, default=300, help='size of sentence encoding layer on top of CNN')

parser.add_argument('--n_fmaps', dest='n_fmaps_psz', type=int, default=100, help='number of cnn feature maps per filter height')

parser.add_argument('--filter_hs', dest='filter_hs', metavar='N', type=int, nargs='+',default =[3,4,5], help='list fo filter heights to use in CNN')

parser.add_argument('--conv_non_linear', dest='conv_non_linear', type=str, default='relu', help='nonlinearity type: tanh/relu')

parser.add_argument('--max_sent_len', dest='maxlen', type=int, default=15, help='size of sentence encoding layer on top of CNN')

parser.add_argument('--sim_smooth_factor', dest='sim_smooth_factor', type=float, default=1.0, help='smoothing factor in softmax')

parser.add_argument('--sim_metric', dest='sim_metric', type=str, default='cosine', help='similarity metric to use to compare sent emb and image emb')

# Regarding word embedding sharing and such

parser.add_argument('--share_Wemb', dest='share_Wemb', type=int, default=0, help='If 1, share Wemb b/w eval and gen models')

parser.add_argument('--fix_Wemb', dest='fix_Wemb', type=int, default=0, help='If 1, gen model doesnt learn Wemb and keeps it fixed')

# optimization parameters

parser.add_argument('-m', '--max_epochs', dest='max_epochs', type=int, default=10, help='number of epochs to train for')

parser.add_argument('--solver', dest='solver', type=str, default='rmsprop', help='solver type: vanilla/adagrad/adadelta/rmsprop')

parser.add_argument('--decay_rate', dest='decay_rate', type=float, default=0.99, help='decay rate for adadelta/rmsprop')

parser.add_argument('--smooth_eps', dest='smooth_eps', type=float, default=1e-8, help='epsilon smoothing for rmsprop/adagrad/adadelta')

parser.add_argument('-lg', '--learning_rate_gen', dest='learning_rate_gen', type=float, default=4e-5, help='solver learning rate')

parser.add_argument('-ld', '--learning_rate_eval', dest='learning_rate_eval', type=float, default=1e-4, help='solver learning rate')

parser.add_argument('-b', '--batch_size', dest='batch_size', type=int, default=1, help='batch size')

parser.add_argument('-cb', '--eval_batch_size', dest='eval_batch_size', type=int, default=2, help='Batch size for eval descriminative network, 1 means it only gets a positive reference\

and a generated sample, n implies it also gets n-1 negative references ')

parser.add_argument('--rand_negs', dest='rand_negs', type=int, default=0, help='How many hard negetives obtianed by random permutations of positive to use to train eval n/w')

parser.add_argument('--sample_by_len', dest='sample_by_len', type=int, default=0, help='enable sampling by length of sentece to speed up training')

parser.add_argument('--grad_clip', dest='grad_clip', type=float, default=10.0, help='clip gradients (normalized by batch size)? elementwise. if positive, at what threshold?')

parser.add_argument('--drop_prob_encoder', dest='drop_prob_encoder', type=np.float32, default=0.5, help='what dropout to apply right after the encoder to an RNN/LSTM')

parser.add_argument('--drop_prob_decoder', dest='drop_prob_decoder', type=np.float32, default=0.5, help='what dropout to apply right before the decoder in an RNN/LSTM')

parser.add_argument('--drop_prob_eval', dest='drop_prob_eval', type=np.float32, default=0.5, help='what dropout to apply right before the decoder in an RNN/LSTM')

parser.add_argument('--drop_prob_aux', dest='drop_prob_aux', type=np.float32, default=0.5, help='what dropout to apply for the auxillary inputs to lstm')

parser.add_argument('--lr_decay', dest='lr_decay', type=float, default=1.0, help='decay factor for learning rate, applied every epoch')

parser.add_argument('--lr_decay_st_epoch', dest='lr_decay_st_epoch', type=float, default=100.0, help='from which epoch should the lr decay start')

# data preprocessing parameters

parser.add_argument('--word_count_threshold', dest='word_count_threshold', type=int, default=5, help='if a word occurs less than this number of times in training data, it is discarded')

# evaluation parameters

parser.add_argument('-p', '--eval_period', dest='eval_period', type=float, default=1.0, help='in units of epochs, how often do we evaluate on val set?')

parser.add_argument('--eval_max_images', dest='eval_max_images', type=int, default=-1, help='for efficiency we can use a smaller number of images to get validation error')

parser.add_argument('--softmax_smooth_factor', dest='softmax_smooth_factor', type=float, default=3.0, help='Is there any auxillary inputs ? If yes indicate file here')

parser.add_argument('--softmax_propogate', dest='softmax_propogate', type=int, default=1, help='Is there any auxillary inputs ? If yes indicate file here')

args = parser.parse_args()

params = vars(args) # convert to ordinary dict

if params['checkpoint_file_name'] != 'None':

checkpoint_init = pickle.load(open(params['checkpoint_file_name'], 'rb'))

model_init_from = checkpoint_init['model']

rg_init = checkpoint_init.get('rgrads',[])

#TODO: GET RID OF THEEEESSSEEEEEE !!!!!!!!!!!!!!!!

checkpoint_init['params'].pop('fappend')

for k in checkpoint_init['params']:

params[k] = checkpoint_init['params'][k]

params['batch_size'] = 1

params['eval_batch_size'] = 4

params['max_epochs'] = 1

params['decay_rate'] = 0.99

params['grad_clip'] = 10.0

params['write_checkpoint_ppl_threshold'] = -1

params['checkpoint_output_directory'] = 'cvCoco/advers'

if params['aux_inp_file'] != 'None':

params['en_aux_inp'] = 1

else:

params['en_aux_inp'] = 0

params['eval_batch_size'] += params['rand_negs']

if params['use_pos_tag'] != 'None':

sentTagMap = pickle.load(open(params['use_pos_tag'],'r'))

print 'parsed parameters:'

print json.dumps(params, indent = 2)

config.mode = 'FAST_RUN'

config.allow_gc = False

#config.exception_verbosity = 'high'

#main(params)