##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

## Created by: Yang Wang

## School of Automation, Huazhong University of Science & Technology (HUST)

## wangyang_sky@hust.edu.cn

## Copyright (c) 2018

##

## This source code is licensed under the MIT-style license found in the

## LICENSE file in the root directory of this source tree

##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

#coding:utf-8

import tensorflow as tf

def my_conv2d(inputs, filters, kernel_size, strides=(1, 1), padding='same', activation=None, name=None, reuse=None):

    return tf.layers.conv2d(inputs=inputs, filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, activation=activation,

                kernel_initializer=tf.truncated_normal_initializer(stddev=0.1), bias_initializer=tf.constant_initializer(0.1), name=name, reuse=reuse)

def build_cnn(input_image=None, image_size=32, n_colors=3, activation_function=tf.nn.relu, reuse=None, name='VGG_NET_CNN'):

    # VGG_NET 32       # [samples, W, H, colors]

    with tf.variable_scope(name, reuse=reuse): 

        input_image = tf.reshape(input_image, shape=[-1, image_size, image_size, n_colors], name='Reshape_inputs')

        # layer_1   # 4个3*3*32

        h_conv1_1 = my_conv2d(input_image, filters=32, kernel_size=(3,3), activation=activation_function, name='conv1_1')

        h_conv1_2 = my_conv2d(h_conv1_1, filters=32, kernel_size=(3,3), activation=activation_function, name='conv1_2')

        h_conv1_3 = my_conv2d(h_conv1_2, filters=32, kernel_size=(3,3), activation=activation_function, name='conv1_3')

        h_conv1_4 = my_conv2d(h_conv1_3, filters=32, kernel_size=(3,3), activation=activation_function, name='conv1_4')

        h_pool1 = tf.layers.max_pooling2d(h_conv1_4, pool_size=(2,2), strides=(2,2), padding='same', name='max_pooling_1')    # shape is (None, 16, 16, 32)

        # layer_2

        h_conv2_1 = my_conv2d(h_pool1, filters=64, kernel_size=(3,3), activation=activation_function, name='conv2_1')

        h_conv2_2 = my_conv2d(h_conv2_1, filters=64, kernel_size=(3,3), activation=activation_function, name='conv2_2')

        h_pool2 = tf.layers.max_pooling2d(h_conv2_2, pool_size=(2,2), strides=(2,2), padding='same', name='max_pooling_2')    # shape is (None, 8, 8, 64)

        # layer_3

        h_conv3_1 = my_conv2d(h_pool2, filters=128, kernel_size=(3,3), activation=activation_function, name='conv3_1')

        h_pool3 = tf.layers.max_pooling2d(h_conv3_1, pool_size=(2,2), strides=(2,2), padding='same', name='max_pooling_3')    # shape is (None, 4, 4, 128)

    return h_pool3

def build_convpool_max(input_image, nb_classes, image_size=32, n_colors=3, 

        n_timewin=7, dropout_rate=0.5, name='CNN_Max', train=True, reuse=False):

    """

    Builds the complete network with maxpooling layer in time.

    :param input_image: list of EEG images (one image per time window)

    :param nb_classes: number of classes

    :param image_size: size of the input image (assumes a square input)

    :param n_colors: number of color channels in the image

    :param n_timewin: number of time windows in the snippet

    :return: a pointer to the output of last layer

    """

    with tf.name_scope(name):

        with tf.name_scope('Parallel_CNNs'):

            convnets = []

            # Build 7 parallel CNNs with shared weights

            for i in range(n_timewin):

                if i==0:

                    convnet = build_cnn(input_image[i],image_size=image_size,n_colors=n_colors, reuse=reuse)

                else:

                    convnet = build_cnn(input_image[i],image_size=image_size,n_colors=n_colors, reuse=True)

                convnets.append(convnet)    # list contains [None, 4, 4, 128]

            convnets = tf.stack(convnets)   # [n_timewin, nSamples, 4, 4, 128]

            convnets = tf.transpose(convnets, [1,0,2,3,4]) # [nSamples, n_timewin, 4, 4, 128]

        with tf.variable_scope('Max_pooling_over_flames'):

            # convpooling using Max pooling over frames

            convnets = tf.reshape(convnets, shape=[ -1, n_timewin, 4*4*128, 1])

            convpool = tf.nn.max_pool(convnets, # [nSamples, 1，4*4*128, 1]

                ksize=[1, n_timewin, 1, 1], strides=[1, 1, 1, 1], padding='VALID', name='convpool_max')

        convpool_flat = tf.reshape(convpool, [-1, 4*4*128])

        h_fc1_drop1 = tf.layers.dropout(convpool_flat, rate=dropout_rate, training=train, name='dropout_1')

        # input shape [batch, 4*4*128] output shape [batch, 512]

        h_fc1 = tf.layers.dense(h_fc1_drop1, 512, activation=tf.nn.relu, name='fc_relu_512')

        # dropout 

        h_fc1_drop2 = tf.layers.dropout(h_fc1, rate=dropout_rate, training=train, name='dropout_2')

        # inputshape [batch, 512] output shape [batch, nb_classes]    # the loss function contains the softmax activation

        prediction = tf.layers.dense(h_fc1_drop2, nb_classes, name='fc_softmax')

    return prediction

def build_convpool_conv1d(input_image, nb_classes, image_size=32, n_colors=3, 

        n_timewin=7, dropout_rate=0.5, name='CNN_Conv1d', train=True, reuse=False):

    """

    Builds the complete network with 1D-conv layer to integrate time from sequences of EEG images.

    :param input_image: list of EEG images (one image per time window)

    :param nb_classes: number of classes

    :param image_size: size of the input image (assumes a square input)S

    :param n_colors: number of color channels in the image

    :param n_timewin: number of time windows in the snippet

    :return: a pointer to the output of last layer

    """

    with tf.name_scope(name):

        with tf.name_scope('Parallel_CNNs'):

            convnets = []

            # Build 7 parallel CNNs with shared weights

            for i in range(n_timewin):

                if i==0:

                    convnet = build_cnn(input_image[i],image_size=image_size,n_colors=n_colors, reuse=reuse)

                else:

                    convnet = build_cnn(input_image[i],image_size=image_size,n_colors=n_colors, reuse=True)

                convnets.append(convnet)

            convnets = tf.stack(convnets)

            convnets = tf.transpose(convnets, [1,0,2,3,4])

        with tf.variable_scope('Conv1d_over_flames'):

            convnets = tf.reshape(convnets, shape=[ -1, n_timewin, 4*4*128, 1])

            convpool = my_conv2d(convnets, filters=64, kernel_size=(3, 4*4*128), strides=(1, 1), padding='valid', activation=tf.nn.relu, name='convpool_conv1d')

        with tf.variable_scope('Output_layers'):

            convpool_flat = tf.reshape(convpool, [-1, (n_timewin-2)*64])

            h_fc1_drop1 = tf.layers.dropout(convpool_flat, rate=dropout_rate, training=train, name='dropout_1')

            h_fc1 = tf.layers.dense(h_fc1_drop1, 256, activation=tf.nn.relu, name='fc_relu_256')

            h_fc1_drop2 = tf.layers.dropout(h_fc1, rate=dropout_rate, training=train, name='dropout_2')

            prediction = tf.layers.dense(h_fc1_drop2, nb_classes, name='fc_softmax')

    return prediction

def build_convpool_lstm(input_image, nb_classes, grad_clip=110, image_size=32, n_colors=3, 

        n_timewin=7, dropout_rate=0.5, num_units=128, batch_size=32, name='CNN_LSTM', train=True, reuse=False):

    """

    Builds the complete network with LSTM layer to integrate time from sequences of EEG images.

    :param input_image: list of EEG images (one image per time window)

    :param nb_classes: number of classes

    :param grad_clip:  the gradient messages are clipped to the given value during

                        the backward pass.

    :param image_size: size of the input image (assumes a square input)

    :param n_colors: number of color channels in the image

    :param n_timewin: number of time windows in the snippet

    :param num_units: number of units in the LSTMCell

    :return: a pointer to the output of last layer

    """

    with tf.name_scope(name):

        with tf.name_scope('Parallel_CNNs'):

            convnets = []

            # Build 7 parallel CNNs with shared weights

            for i in range(n_timewin):

                if i==0:

                    convnet = build_cnn(input_image[i],image_size=image_size,n_colors=n_colors, reuse=reuse)

                else:

                    convnet = build_cnn(input_image[i],image_size=image_size,n_colors=n_colors, reuse=True)

                convnets.append(convnet)

            convnets = tf.stack(convnets)

            convnets = tf.transpose(convnets, [1,0,2,3,4]) # 调换轴 shape: (nSamples, n_timewin, 4, 4, 128)

        with tf.variable_scope('LSTM_layer'):

            # (nSamples, n_timewin, 4, 4, 128) ==>  (nSamples, n_timewin, 4*4*128)

            convnets = tf.reshape(convnets, shape=[-1, n_timewin, 4*4*128], name='Reshape_for_lstm')

            #lstm cell inputs:[batchs, time_steps, 4*4*128]

            with tf.variable_scope('LSTM_Cell'):

                lstm_cell = tf.contrib.rnn.BasicLSTMCell(num_units=num_units, forget_bias=1.0, state_is_tuple=True)

                outputs, final_state = tf.nn.dynamic_rnn(lstm_cell, convnets, dtype=tf.float32, time_major=False)

                # outputs.shape is (batch_size, time_steps, num_units)

                outputs = tf.transpose(outputs, [1,0,2])        # (time_steps, batch_size, num_units)

                outputs = outputs[-1]

        with tf.variable_scope('Output_layers'):

            h_fc1_drop1 = tf.layers.dropout(outputs, rate=dropout_rate, training=train, name='dropout_1')

            h_fc1 = tf.layers.dense(h_fc1_drop1, 256, activation=tf.nn.relu, name='fc_relu_256')

            h_fc1_drop2 = tf.layers.dropout(h_fc1, rate=dropout_rate, training=train, name='dropout_2')

            prediction = tf.layers.dense(h_fc1_drop2, nb_classes, name='fc_softmax')

    return prediction

def build_convpool_mix(input_image, nb_classes, grad_clip=110, image_size=32, n_colors=3, 

        n_timewin=7, dropout_rate=0.5, num_units=128, batch_size=32, name='CNN_Mix', train=True, reuse=False):

    """

    Builds the complete network with LSTM and 1D-conv layers combined

    :param input_image: list of EEG images (one image per time window)

    :param nb_classes: number of classes

    :param grad_clip:  the gradient messages are clipped to the given value during

                        the backward pass.

    :param imsize: size of the input image (assumes a square input)

    :param n_colors: number of color channels in the image

    :param n_timewin: number of time windows in the snippet

    :return: a pointer to the output of last layer

    """

    with tf.name_scope(name):

        with tf.name_scope('Parallel_CNNs'):

            convnets = []

            # Build 7 parallel CNNs with shared weights

            for i in range(n_timewin):

                if i==0:

                    convnet = build_cnn(input_image[i],image_size=image_size,n_colors=n_colors, reuse=reuse)

                else:

                    convnet = build_cnn(input_image[i],image_size=image_size,n_colors=n_colors, reuse=True)

                convnets.append(convnet)

            convnets = tf.stack(convnets)

            convnets = tf.transpose(convnets, [1,0,2,3,4])

        with tf.variable_scope('Conv1d_over_flames'):

            convpool = tf.reshape(convnets, shape=[ -1, n_timewin, 4*4*128, 1])

            convpool = my_conv2d(convpool, filters=64, kernel_size=(3, 4*4*128), strides=(1, 1), padding='valid', activation=tf.nn.relu, name='convpool_conv1d')

            conv1d_out = tf.reshape(convpool, [-1, (n_timewin-2)*64])

        with tf.variable_scope('LSTM_layer'):

            # (nSamples, n_timewin, 4, 4, 128) ==>  (nSamples, n_timewin, 4*4*128)

            convnets = tf.reshape(convnets, shape=[-1, n_timewin, 4*4*128], name='Reshape_for_lstm')

            #lstm cell inputs:[batchs, time_steps, 4*4*128]

            with tf.variable_scope('LSTM_Cell'):

                lstm_cell = tf.contrib.rnn.BasicLSTMCell(num_units=num_units, forget_bias=1.0, state_is_tuple=True)

                outputs, final_state = tf.nn.dynamic_rnn(lstm_cell, convnets, dtype=tf.float32, time_major=False)

                # outputs.shape is (batch_size, time_steps, num_units)

                outputs = tf.transpose(outputs, [1,0,2])

                lstm_out = outputs[-1]

        with tf.variable_scope('Output_layers'):

            dense_in = tf.concat((conv1d_out, lstm_out), axis=1, name='concat_conv1d_lstm')    # shape [batch, (n_timewin-2)*64+num_units]

            h_fc1_drop1 = tf.layers.dropout(dense_in, rate=dropout_rate, training=train, name='dropout_1')

            h_fc1 = tf.layers.dense(h_fc1_drop1, 512, activation=tf.nn.relu, name='fc_relu_512')

            h_fc1_drop2 = tf.layers.dropout(h_fc1, rate=dropout_rate, training=train, name='dropout_2')

            prediction = tf.layers.dense(h_fc1_drop2, nb_classes, name='fc_softmax')

    return prediction