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, prepare_seq_features

from imagernn.solver import Solver

from imagernn.imagernn_utils import decodeGenerator, eval_split_theano, eval_prep_refs

#from numbapro import cuda

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

from imagernn.recurrent_feat_encoder import RecurrentFeatEncoder

from imagernn.attention_network import AttentionNetwork

from collections import defaultdict

import operator

def main(params):

print e

if __name__ == "__main__":

parser = argparse.ArgumentParser()

# global setup settings, and checkpoints

parser.add_argument('--use_theano', dest='use_theano', default=1, help='Should we use thano and gpu!?. Actually dont try with value 0 :-|')

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('--restore_grads', dest='restore_grads', type=int, default=1, help='restore grads from checkpoint or not')

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('--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=1, help='Feed image features through auxillary input!')

# parameters for loading multiple features per video using labels.txt

parser.add_argument('--labelsFile', dest='labels', type=str, default='labels.txt', help='labels.txt file for this dataset')

parser.add_argument('--featfromlbl', dest='featfromlbl', type=str, default='ALL ALL', help='should we use lables.txt, if yes which feature?'

'use + sign to seperately specify for img and aux')

parser.add_argument('--poolmethod', dest='poolmethod', type=str, default='max', help='What pooling to use if multiple features are found')

parser.add_argument('--uselabel', dest='uselabel', type=int, default=3, help='which features should use labels.txt, img/aux or both, 0 - None, 1 - img, 2 - aux, 3 - both')

parser.add_argument('--use_partial_sent',dest='use_partial_sent', nargs='+',type=int, default=[], help='Indexes of sentences to use')

parser.add_argument('--use_train_subset',dest='use_train_subset', nargs='+',type=int, default=[], help='Indexes of training subsets to use')

# Parameters to enable class based factorization of lstm output

parser.add_argument('--class_out_factoring', dest='class_out_factoring', type=int, default=0, help='Enable Class based output factorization in generator')

parser.add_argument('--nClasses', dest='nClasses', type=int, default=200, help='Number of classes to use')

parser.add_argument('--class_inp_file', dest='class_inp_file', type=str, default=None, help='If clustering is already done, provide the inp file')

parser.add_argument('--clust_tool_dir', dest='clust_tool_dir', type=str, default=None, help='Needed if clustering needs to be triggered now.'

' Support currently is only for Percy Liang\'s Brown clustering implementation')

parser.add_argument('--cls_zmean', dest='cls_zmean', type=int, default=0, help='cls_zmean to use')

parser.add_argument('--cls_diff_layer', dest='cls_diff_layer', type=int, default=0, help='cls_diff_layer')

# 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('--en_residual_conn', dest='en_residual_conn', type=int, default=0, 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?')

# 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 types supported: rmsprop')

parser.add_argument('--decay_rate', dest='decay_rate', type=float, default=0.999, 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('-l', '--learning_rate', dest='learning_rate', type=float, default=1e-3, help='solver learning rate')

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

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

parser.add_argument('--use_dropout', dest='use_dropout', type=int, default=1, help='enable or disable dropout')

parser.add_argument('--drop_prob_encoder', dest='drop_prob_encoder', type=float, 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=float, 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=float, default=0.5, help='what dropout to apply for the auxillary inputs to lstm')

# refere paper: "Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks" http://arxiv.org/abs/1506.03099

parser.add_argument('--sched_sampling_mode', dest='sched_sampling_mode', type=str, default=None, help='should we implement scheduled sampling during training')

parser.add_argument('--sched_sampling_const', dest='sched_sampling_const', type=float, default=1.0, help='scheduling constant, exact nature depends on the mode')

parser.add_argument('--sslin_slope', dest='sslin_slope', type=np.float, default=1.0, help='slope of decay in linear scheduling')

parser.add_argument('--sslin_min', dest='sslin_min', type=np.float, default=1.0, help='min truth constant in linear scheduling')

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('--maxlen', dest='maxlen', type=int, default=None, help='enable sampling by length of sentece to speed up training')

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""")

parser.add_argument('--use_unk_token', dest='use_unk_token', type=int, default=0, help=' replace unknows with UNK token')

parser.add_argument('--reverse_sentence', dest='reverse_sentence', type=int, default=0, help='Should we reverse the sentences when feeding it to the RNN?')

parser.add_argument('--use_video_feat', dest='use_video_feat', type=int, default=0, help='Use video features for training')

# 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_batch_size', dest='eval_batch_size', type=int, default=100, help="""for faster validation performance evaluation, what batch

size to use on val img/sentences?""")

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('--eval_metric', dest='eval_metric', type=str, default='perplex', help="""Specify the evaluation metric to use on validation. Possible

values are perplex, meteor, cider""")

parser.add_argument('--metrics_to_track', dest='met_to_track',nargs='+', type=str, default=[], help="""Specify the evaluation metric to use on validation. Possible

values are perplex, meteor, cider""")

parser.add_argument('--eval_beamsize', dest='eval_beamsize', type=int, default=1, help='what beamsize to use in validation when using AEMs')

# parameters controlling use of external data (For eg. lsmdc uses COCO Nearest neighbhors)

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

parser.add_argument('--ed_sample_prob', dest='ed_sample_prob', type=np.float, default=0.0, help='probability with which to sample external data')

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

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

# parameters to use a feature encoding recurrent network

parser.add_argument('--feat_encoder', dest='feat_encoder', type=str, default=None, help='Which encoder should we use')

parser.add_argument('--use_encoder_for', dest='use_encoder_for', type=int, default=0, help='Is it for image feat or aux input')

parser.add_argument('--use_shared_mem_enc', dest='use_shared_mem_enc', type=int, default=1, help='Is it for image feat or aux input')

parser.add_argument('--featenc_hidden_size', dest='featenc_hidden_size', type=int, default=512, help='Should img or aux features be read from disk')

# Implement option to encode GT sentences and use it as features to the generator

parser.add_argument('--encode_gt_sentences', dest='encode_gt_sentences', type=int, default=0, help='Should img or aux features be read from disk')

parser.add_argument('--n_encgt_sent', dest='n_encgt_sent', type=int, default=5, help='Should img or aux features be read from disk')

# parameters to use an attention mechanism.

parser.add_argument('--use_attention_network', dest='use_attn', type=str, default=None, help='Which encoder should we use')

parser.add_argument('--attn_hidden_config', dest='attn_hidden_config',nargs='+', type=int, default=[100], help='Is it for image feat or aux input')

# Some params to enable partial feature reading from disk!!

parser.add_argument('--disk_feature', dest='disk_feature', type=int, default=0, help='Should img or aux features be read from disk')

# Implementing gumbel softmax

parser.add_argument('--use_gumbel_mse', dest='use_gumbel_mse', type=int, default=0, help='Should img or aux features be read from disk')

parser.add_argument('--gumbel_temp_init', dest='gumbel_temp_init', type=np.float, default=0, help='Should img or aux features be read from disk')

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',[])

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

params['en_aux_inp'] = 1

else:

params['en_aux_inp'] = 0

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)