#edit-mode: -*- python -*-
#coding:gbk

'''
Configure args:
generating (bool, default: False)
    1. True:
        (1). fname_test_pflist: file name of test img_sen_arr list
        (2). fname_dct: file name of the generating dictionary
    2. False:
        (1). fname_train_pflist: file name of train img_sen_arr list
        (2). fname_test_pflist: file name of test img_sen_arr list
        (3). fname_dct: file name of the dictionary
'''

import math
import numpy as np

drop_rate=0.5
batch_size = 128
default_decay_rate(1e-4*batch_size)

lr_fast = 0.1
lr_common = 5e-3
lr_ori = lr_fast

word_embedding_dim=512
hidden_dim = 512
#original multimodal_dim = 1024
multimodal_dim = 256

img_feat_dim = 9*4096 # NEED to be modified
img_feat_dim_reduce = 9*256 # NEED to be modified
num_channels = 4096 # NEEd to be modified
saliency_filter_size = 3 # NEED to be modified
sal_map_dim = 3*3 # NEED to be modified
num_channels_exp = 4096+180 #NEED to be modified
num_filters_reduce = 256 # NEED to be modified

color_feat_dim = 180*9
default_initial_std(1/math.sqrt(hidden_dim))

generating = get_config_arg('generating', bool, False)

img_feat_norm_pD = 0.0001

if generating:
    fname_test_pflist = get_config_arg('fname_test_pflist', str, '')
    fname_dct_ques_pD = get_config_arg('fname_dct_ques', str, '')
    fname_dct_ans_pD = get_config_arg('fname_dct_ans', str, '')
    ques_dct_size = len(np.load(fname_dct_ques_pD))
    ans_dct_size = len(np.load(fname_dct_ans_pD))
    Inputs("img_feat", "question", "answer", "color")
    Outputs("predicted_id")
    #Outputs("sal_norm")

    TestData(PyData(files = fname_test_pflist,
                 load_data_module="pyDataPro_img_sen_ans_color",
                 load_data_object="ImgColorQuesAnsDataProvider",
                 load_data_args="img_size:%d;feat_norm_factor:%f;ques_dct_size:%d;ans_dct_size:%d;mode:%s"%(img_feat_dim, img_feat_norm_pD, ques_dct_size, ans_dct_size, 'test')
          )
    )

    Layer(
        name = "img_feat",
        type = "data",
        size = img_feat_dim,
    )

    Layer(
       name = "question",
       type = "data",
       size = ques_dct_size,
    )

    Layer(
        name = "answer",
        type = "data",
        size = ans_dct_size,
    )

    Layer(
        name = "color",
        type = "data",
        size = color_feat_dim,
    )

    Layer(
        name = "word_embedding",
        type = "mixed",
		active_type = "linear",
        size = word_embedding_dim,
        bias = False,
        inputs = TableProjection(
                   "question",
                   parameter_name="wordvecs",
                   learning_rate=lr_ori,
                   momentum=0.9,
                   decay_rate=5e-4*batch_size,
                   initial_mean=0.0,
                   initial_std=0.01),
    )

else:
    # Input slots:
    # "img_feat" represents the feature of images
    # "word": current word index
    # "next_word": next word index from orignial dictionary + new obj dictionary
    Inputs("img_feat", "question", "answer", "color")

    # Output slots:
    # multiclass cross entropy
    Outputs("cost")

    # pyDataProvidor args (pD):
    flist_tr_pD = get_config_arg('fname_train_pflist', str, '')
    flist_te_pD = get_config_arg('fname_test_pflist', str, '')
    fname_dct_ques_pD = get_config_arg('fname_dct_ques', str, '')
    fname_dct_ans_pD = get_config_arg('fname_dct_ans', str, '')
    ques_dct_size = len(np.load(fname_dct_ques_pD))
    ans_dct_size = len(np.load(fname_dct_ans_pD))

    TrainData(PyData(files = flist_tr_pD,
                 load_data_module="pyDataPro_img_sen_ans_color",
                 load_data_object="ImgColorQuesAnsDataProvider",
                 load_data_args="img_size:%d;feat_norm_factor:%f;ques_dct_size:%d;ans_dct_size:%d;mode:%s"%(img_feat_dim, img_feat_norm_pD, ques_dct_size, ans_dct_size, 'train'),
                 async_load_data=False
          )
    )

    TestData(PyData(files = flist_te_pD,
                 load_data_module="pyDataPro_img_sen_ans_color",
                 load_data_object="ImgColorQuesAnsDataProvider",
                 load_data_args="img_size:%d;feat_norm_factor:%f;ques_dct_size:%d;ans_dct_size:%d;mode:%s"%(img_feat_dim, img_feat_norm_pD, ques_dct_size, ans_dct_size, 'test'),
                 async_load_data=False
          )
    )

    Layer(
        name = "img_feat",
        type = "data",
        size = img_feat_dim,
    )

    Layer(
        name = "question",
        type = "data",
        size = ques_dct_size,
    )

    Layer(
       name = "answer",
       type = "data",
       size = ans_dct_size,
    )

    Layer(
        name = "color",
        type = "data",
        size = color_feat_dim,
    )

    Layer(
        name = "word_embedding",
        type = "mixed",
		active_type = "linear",
        size = word_embedding_dim,
        bias = False,
        inputs = TableProjection(
                   "question",
                   parameter_name="wordvecs",
                   learning_rate=lr_ori,
                   momentum=0.9,
                   decay_rate=5e-4*batch_size,
                   initial_mean=0.0,
                   initial_std=0.5),
    )

Layer(
    name = "hidden1",
    type = "mixed",
    size = hidden_dim,
    active_type = "stanh",
    bias = True,
    inputs = [
        FullMatrixProjection("word_embedding", learning_rate=lr_ori),
    ]
)

'''
'''
Layer(
    name = "rnn1",
    type = "lstmemory",
    active_type = "relu",
    active_state_type = "linear",
    active_gate_type = "sigmoid",
    bias = Bias(initial_std=0),
    inputs = Input("hidden1", initial_std=0, learning_rate=lr_ori),
)

Layer(
    name = "word_rnn_mix",
    type = "mixed",
    size = hidden_dim,
    active_type = "stanh",
    bias = True,
    inputs = [
        FullMatrixProjection("hidden1", learning_rate=lr_ori),
        FullMatrixProjection("rnn1", learning_rate=lr_ori),
    ],
    drop_rate = drop_rate,
)

Layer(
    name = "mix_avg",
    type = "average",
    inputs = [Input("word_rnn_mix")],
)

Layer(
    name = "fc1",
    type = "mixed",
    device = -1,
    size = 9*num_channels_exp,
    bias = False,
    active_type = 'softmax',
    inputs = [
        FullMatrixProjection("mix_avg", initial_std = 0, learning_rate=lr_ori),
    ]
)

Layer(
    name = "concat_color_img",
    type = "concat",
    inputs = [
        Input("img_feat",initial_std = 0, learning_rate = lr_ori),
        Input("color", initial_std = 0, learning_rate = lr_ori),
    ]
)

Layer(
    name = "sal_norm",
    type = "mixed",
    size = sal_map_dim,
    bias = False,
    active_type='softmax',
    inputs = ConvOperator(
               input_layer_names=["concat_color_img", "fc1"],
               num_filters=1,
               conv_conf=Conv(filter_size=saliency_filter_size,
                              channels = num_channels_exp,
                              padding = 1,
                              stride = 1,
                              groups = 1)),
)


Layer(
    name = "sal_exp",
    type = "featmap_expand",
        device = -1,
    num_filters = num_channels_exp,
    inputs = [
        Input("sal_norm", initial_std=0, learning_rate=lr_ori),
        ],
)

Layer(
    name = "product",
    type = "mixed",
    device = -1,
	active_type = "linear",
    size = sal_map_dim*num_channels_exp,
    bias = False,
    inputs = DotMulOperator(
        input_layer_names=["sal_exp", "concat_color_img"],
        scale = 1.0),
)


Layer(
    name = "reduce",
    type = "cudnn_conv",
    active_type = "stanh",
    bias = Bias(learning_rate=lr_ori*2,
                momentum=0.9,
                initial_mean=0,
                initial_std=0.0),
    inputs = Input("product",
                   learning_rate=lr_ori,
                   momentum=0.9,
                   decay_rate=5e-4*batch_size,
                   initial_mean=0.0,
                   initial_std=0.1,
                   conv = Conv(filter_size = 3,
                               channels = num_channels,
                               padding = 1,
                               stride = 1,
                               groups = 1)),
    num_filters = num_filters_reduce,
    shared_biases = True,
)

Layer(
    name = "hidden2",
    type = "mixed",
    size = word_embedding_dim,
    active_type = "stanh",
    inputs = [
        FullMatrixProjection("reduce",
                            initial_std=0.01,
                            learning_rate=lr_ori),
        FullMatrixProjection("mix_avg",
                            learning_rate=lr_ori,
                            momentum=0.05,
                            decay_rate=5e-4*batch_size,
                            initial_mean=0.0,
                            initial_std=1),
        FullMatrixProjection("img_feat",
                            initial_std=0.01,
                            learning_rate=lr_ori),
    ],
    drop_rate = drop_rate,
)

Layer(
    name = "hidden3",
    type = "mixed",
    size = hidden_dim,
    active_type = "stanh",
    bias = True,
    inputs = [
        FullMatrixProjection("hidden2", learning_rate=lr_ori),
    ]
)


if generating:
    Layer(
        name = "output",
        type = "mixed",
        size = ans_dct_size,
        active_type = "softmax",
        bias = Bias(
                 learning_rate=lr_ori,
                 momentum=0.9,
                 initial_mean=0,
                 initial_std=0),
        inputs = TransposedFullMatrixProjection(
                   "hidden3",
                   parameter_name="outputvecs",
                   learning_rate=lr_ori,
                   momentum=0.9,
                   decay_rate=5e-4*batch_size,
                   initial_mean=0.0,
                   initial_std=0.01),

    )

    Layer(
       name = "predicted_id",
       type = "maxid",
       inputs = "output",
    )
else:
    Layer(
        name = "output",
        type = "mixed",
        size = ans_dct_size,
        active_type = "softmax",
        bias = Bias(
                 learning_rate=lr_ori,
                 momentum=0.9,
                 initial_mean=0,
                 initial_std=0),
        inputs = TransposedFullMatrixProjection(
                   "hidden3",
                   parameter_name="outputvecs",
                   learning_rate=lr_ori,
                   momentum=0.9,
                   decay_rate=5e-4*batch_size,
                   initial_mean=0.0,
                   initial_std=0.01),

    )

    Layer(
       name = "cost",
       type = "multi-class-cross-entropy",
       inputs = ["output", "answer"],
    )

    Layer(
       name = "error_layer",
       type = "classification_error",
       inputs = ["output", "answer"]
    )

    Evaluator(
       name = "error_classfication",
       type = "sum",
       inputs = "error_layer",
    )

Settings(
    algorithm = 'sgd',
    learning_method = 'adadelta',
    ada_epsilon = 0.01,
    ada_rou = 0.99,
    batch_size = batch_size,
    average_window = 0.5,
    max_average_window = int(437543 / batch_size),
    learning_rate = lr_common,
    learning_rate_decay_a = 5e-7,
    learning_rate_decay_b = 0.5,
)
