'''
Created on 2019年3月26日

@author: User
'''

import os
import sys
sys.path.append(os.path.abspath("../.."))

from BiddingKG.dl.common.Utils import *
from keras import layers
from keras import models
from keras import losses,optimizers
import tensorflow as tf
from BiddingKG.dl.common.models import transformer_model

def getBiLSTMModel():
    input = layers.Input(shape=(100,))
    
    lstm = layers.Bidirectional(layers.LSTM(128,return_sequences=False))(input)
    
    matrix = layers.Dense(48,activation="relu")(lstm)
    
    out = layers.Dense(2,activation="softmax")(matrix)
    
    model = models.Model(inputs=input,outputs=out)
    
    model.compile(optimizer=optimizers.Adam(lr=0.01), loss=losses.categorical_crossentropy,metrics=[precision,recall,f1_score])
    model.summary()
    return model

def getTextCNNModel():
    input = layers.Input(shape=(100,60))
    
    list_pool = []
    
    list_size = [10,15,20]

    for size in list_size:    
        c = layers.Conv1D(filters=4,kernel_size=size,strides=1,activation="relu")(input)
        p = layers.AvgPool1D(pool_size=int(c.shape[1]))(c)
        list_pool.append(p)
    
    concat = layers.merge(list_pool,mode="concat")
    
    flatten = layers.Flatten()(concat)
    
    matrix = layers.Dense(12,activation="relu")(flatten)
    
    
    out = layers.Dense(2,activation="softmax")(matrix)
    
    model = models.Model(inputs=input,outputs=out)
    model.compile(optimizer=optimizers.Adadelta(),loss=losses.categorical_crossentropy,metrics=[precision,recall,f1_score])
    model.summary()
    return model

def get_context_form_model(vocab_len,char_dim,lstm_dim,context_dim,res_dim,residual_stacks):
    input = tf.placeholder(shape=[None,9,30],dtype=tf.int32)
    label = tf.placeholder(shape=[None,2],dtype=tf.int32)

    # input_center = tf.slice(input,[0,4,0],[-1,1,-1])
    with tf.variable_scope("embedding"):
        embedding = tf.get_variable("char_embedding",shape=[vocab_len,char_dim])
        embedding_input = tf.nn.embedding_lookup(embedding,input)


    with tf.variable_scope("center"):
        center_embedding = tf.slice(embedding_input,[0,4,0,0],[-1,1,-1,-1])
        with tf.variable_scope("bert"):
            bert_input = tf.reshape(center_embedding,[-1,30,char_dim])
            bert_output = transformer_model(input_tensor=bert_input,name="bert")

            forward_cell = tf.contrib.rnn.BasicLSTMCell(lstm_dim, state_is_tuple=True)
            backward_cell = tf.contrib.rnn.BasicLSTMCell(lstm_dim,state_is_tuple=True)
            outputs, final_out = tf.nn.bidirectional_dynamic_rnn(forward_cell,backward_cell,bert_output, dtype=tf.float32)
            bi_output = tf.concat(final_out,-1) #[-1,lstm_dim*2]
            bi_output = tf.reshape(tf.slice(bi_output,[1,0,0],[1,-1,-1]),[-1,lstm_dim*2])

            center_output = tf.nn.relu(tf.matmul(bi_output,tf.get_variable("center_v",shape=[lstm_dim*2,res_dim])))
            center_output = tf.nn.dropout(center_output,1)

    with tf.variable_scope("context"):
        with tf.variable_scope("bi_rnn"):
            resize_input = tf.reshape(embedding_input,[-1,30,char_dim])
            forward_cell = tf.contrib.rnn.BasicLSTMCell(lstm_dim, state_is_tuple=True)
            backward_cell = tf.contrib.rnn.BasicLSTMCell(lstm_dim,state_is_tuple=True)
            outputs, final_out = tf.nn.bidirectional_dynamic_rnn(forward_cell,backward_cell,resize_input, dtype=tf.float32)
            bi_output = tf.concat(final_out,-1) #[-1,lstm_dim*2]
            bi_output = tf.slice(bi_output,[1,0,0],[1,-1,-1])

            context_input = tf.reshape(bi_output,[-1,lstm_dim*2])
            context_v = tf.get_variable("context_v",shape=[lstm_dim*2,context_dim])

            context_emb = tf.nn.relu(tf.matmul(context_input,context_v))

            context_output = tf.reshape(context_emb,[-1,9*context_dim])

            context_output = tf.nn.relu(tf.matmul(context_output,tf.get_variable("context_output_v",shape=[9*context_dim,res_dim])))

            context_output = tf.nn.dropout(context_output,1)


    with tf.variable_scope("residual"):
        input_x = tf.concat([context_output,center_output],axis=-1)
        for stack in range(residual_stacks):
            stack_w = tf.get_variable("stack_w_%d"%stack,shape=[res_dim*2,res_dim*2])
            stack_b = tf.get_variable("stack_b_%d"%stack,shape=[1,res_dim*2])
            stack_out = tf.matmul(input_x,stack_w)+stack_b
            input_x = tf.nn.relu(stack_out+input_x)



    with tf.variable_scope("softmax"):
        softmax_input = tf.identity(input_x)
        softmax_w = tf.get_variable("softmax_w",shape=[res_dim*2,2])
        softmax_output = tf.nn.softmax(tf.matmul(softmax_input,softmax_w))



    # with tf.variable_scope("softmax"):
    #     softmax_input = tf.reshape(bi_output,[-1,lstm_dim*2])
    #     softmax_w = tf.get_variable("softmax_w",shape=[lstm_dim*2,2])
    #     softmax_output = tf.nn.softmax(tf.matmul(softmax_input,softmax_w))

    _p = precision(tf.cast(label,tf.float32),softmax_output)
    _r = recall(tf.cast(label,tf.float32),softmax_output)
    loss = -tf.reduce_mean(tf.cast(label,tf.float32)*tf.log(softmax_output))

    optimizer = tf.train.AdamOptimizer(learning_rate=0.01)
    global_step = tf.Variable(0, trainable=False)
    grads_vars = optimizer.compute_gradients(loss)
    capped_grads_vars = [[tf.clip_by_value(g, -5, 5), v] for g, v in grads_vars]
    train = optimizer.apply_gradients(capped_grads_vars, global_step)

    return {"input":input,"output":softmax_output,"label":label,"train":train,"embedding":embedding,"loss":loss,"precision":_p,"recall":_r}






if __name__=="__main__":
    #getBiLSTMModel()
    getTextCNNModel()