import keras
import tensorflow as tf
from keras.models import Sequential
from keras.layers import Dense, Dropout, Flatten , Activation, Masking, LSTM, GRU, Multiply
from keras.layers import Conv2D, MaxPooling2D
from keras import backend as K
from keras.callbacks import Callback
from keras.layers import Lambda, Input, Dense, Concatenate ,Conv2DTranspose, TimeDistributed
from keras.layers import LeakyReLU,BatchNormalization,AveragePooling2D,Reshape
from keras.layers import UpSampling2D,ZeroPadding2D
from keras.losses import mse, binary_crossentropy
from keras.models import Model
from keras.layers import Lambda,TimeDistributed, Bidirectional
from keras import layers


def Net_2LSTM():

    input1 = Input((400,49))
    input2 = Input((400,23))

    #print("input1:",input1.shape)

    l = Bidirectional(LSTM(32, dropout=0.2, recurrent_dropout=0.2, input_shape=(400,49)))(input1)
    d1 = Dense(32,activation='relu')(l)

    #print('d1:',d1.shape)

    l2 = Bidirectional(LSTM(16, dropout=0.2, recurrent_dropout=0.2, input_shape=(400,23)))(input2)
    d2 = Dense(32,activation="relu")(l2)


    d3 = Concatenate()([d1,d2])
    d4 = Dense(64,activation="relu")(d3)

    #branch1: PHQ_Binary
    b1 = Dense(16,activation="relu")(d4)
    o1 = Dense(2,activation="sigmoid",name="classify_1")(b1)

    #print('o1:',o1.shape)

    #branch2: PHQ_Score
    b2 = Dense(16,activation="relu")(d4)
    o2 = Dense(1,activation="relu", name="regress_1")(b2)

    #branch3: PCL-C (PTSD)
    b3 = Dense(16, activation="relu")(d4)
    o3 = Dense(2,activation="sigmoid",name="classify_2")(b3)

    #branch4: PTSD Severity
    b4 = Dense(16, activation="relu")(d4)
    o4 = Dense(1,activation="relu", name="regress_2")(b4)

    model = Model([input1, input2],[o1,o2,o3,o4])

    return model

def Net_2LSTM_1regress():

    input1 = Input((3628,49))
    input2 = Input((3628,23))

    #print("input1:",input1.shape)

    m1 = Masking(mask_value=0)(input1)
    l = Bidirectional(LSTM(32, dropout=0.2, recurrent_dropout=0.2))(m1)
    d1 = Dense(32,activation='relu')(l)

    #print('d1:',d1.shape)

    m2 = Masking(mask_value=0)(input2)
    l2 = Bidirectional(LSTM(16, dropout=0.2, recurrent_dropout=0.2))(m2)
    d2 = Dense(32,activation="relu")(l2)


    d3 = Concatenate()([d1,d2])
    d4 = Dense(64,activation="relu")(d3)

    #branch2: PHQ_Score
    b2 = Dense(16,activation="relu")(d4)
    o2 = Dense(1,activation="relu", name="regress_1")(b2)

    model = Model([input1, input2],o2)

    return model

def Net_2LSTM_1regress_lastLinear():

    input1 = Input((3628,49))
    input2 = Input((3628,23))

    #print("input1:",input1.shape)

    m1 = Masking(mask_value=0)(input1)
    l = Bidirectional(LSTM(32, dropout=0.2, recurrent_dropout=0.2))(m1)
    d1 = Dense(32,activation='relu')(l)

    #print('d1:',d1.shape)

    m2 = Masking(mask_value=0)(input2)
    l2 = Bidirectional(LSTM(16, dropout=0.2, recurrent_dropout=0.2))(m2)
    d2 = Dense(32,activation="relu")(l2)


    d3 = Concatenate()([d1,d2])
    d4 = Dense(64,activation="relu")(d3)

    #branch2: PHQ_Score
    b2 = Dense(16,activation="relu")(d4)
    o2 = Dense(1,name="regress_1")(b2)

    model = Model([input1, input2],o2)

    return model


def Net_3feature_1regress_lastLinear():

    input1 = Input((3628,49))
    input2 = Input((1574,23))
    input3 = Input((128,256))

    #print("input1:",input1.shape)

    #FAUs

    m1 = Masking(mask_value=0)(input1)

    # ATTENTION PART STARTS HERE
    attention_probs_1 = Dense(49, activation='softmax', name='attention_vec_1')(m1)
    attention_mul_1 = Multiply(name='attention_mul_1')([m1, attention_probs_1])
    # ATTENTION PART FINISHES HERE

    #print("att_1:",attention_mul_1)

    l = Bidirectional(LSTM(32, dropout=0.2, recurrent_dropout=0.2))(attention_mul_1)
    d1 = Dense(32,activation='relu')(l)

    #print('d1:',d1.shape)

    #eGeMAPS

    m2 = Masking(mask_value=0)(input2)

    # ATTENTION PART STARTS HERE
    attention_probs_2 = Dense(23, activation='softmax', name='attention_vec_2')(m2)
    attention_mul_2 = Multiply(name='attention_mul_2')([m2, attention_probs_2])
    # ATTENTION PART FINISHES HERE

    l2 = Bidirectional(LSTM(16, dropout=0.2, recurrent_dropout=0.2))(attention_mul_2)
    d2 = Dense(32,activation="relu")(l2)

    #print('d2:',d2.shape)

    #text
    m3 = Masking(mask_value=0)(input3)

    # ATTENTION PART STARTS HERE
    attention_probs_3 = Dense(256, activation='softmax', name='attention_vec_3')(m3)
    attention_mul_3 = Multiply(name='attention_mul_3')([m3, attention_probs_3])
    # ATTENTION PART FINISHES HERE

    #print("att_3:",attention_mul_3.shape)

    l3 = Bidirectional(LSTM(64, dropout=0.2, recurrent_dropout=0.2))(attention_mul_3)
    d3 = Dense(32, activation='relu')(l3)

    #Concatenate
    d3 = Concatenate()([d1,d2,d3])
    d4 = Dense(64,activation="relu")(d3)

    o2 = Dense(1,name="regress_1")(d4)

    model = Model([input1, input2, input3],o2)

    return model

def Net_3feature_1binary():

    input1 = Input((3628,49))
    input2 = Input((1574,23))
    input3 = Input((128,256))

    #print("input1:",input1.shape)

    #FAUs

    m1 = Masking(mask_value=0)(input1)

    # ATTENTION PART STARTS HERE
    attention_probs_1 = Dense(49, activation='softmax', name='attention_vec_1')(m1)
    attention_mul_1 = Multiply(name='attention_mul_1')([m1, attention_probs_1])
    # ATTENTION PART FINISHES HERE

    #print("att_1:",attention_mul_1)

    l = Bidirectional(LSTM(32, dropout=0.2, recurrent_dropout=0.2))(attention_mul_1)
    d1 = Dense(32,activation='relu')(l)

    #print('d1:',d1.shape)

    #eGeMAPS

    m2 = Masking(mask_value=0)(input2)

    # ATTENTION PART STARTS HERE
    attention_probs_2 = Dense(23, activation='softmax', name='attention_vec_2')(m2)
    attention_mul_2 = Multiply(name='attention_mul_2')([m2, attention_probs_2])
    # ATTENTION PART FINISHES HERE

    l2 = Bidirectional(LSTM(16, dropout=0.2, recurrent_dropout=0.2))(attention_mul_2)
    d2 = Dense(32,activation="relu")(l2)

    #print('d2:',d2.shape)

    #text
    m3 = Masking(mask_value=0)(input3)

    # ATTENTION PART STARTS HERE
    attention_probs_3 = Dense(256, activation='softmax', name='attention_vec_3')(m3)
    attention_mul_3 = Multiply(name='attention_mul_3')([m3, attention_probs_3])
    # ATTENTION PART FINISHES HERE

    #print("att_3:",attention_mul_3.shape)

    l3 = Bidirectional(LSTM(64, dropout=0.2, recurrent_dropout=0.2))(attention_mul_3)
    d3 = Dense(32, activation='relu')(l3)

    #Concatenate
    d3 = Concatenate()([d1,d2,d3])
    d4 = Dense(64,activation="relu")(d3)

    o2 = Dense(2, activation="sigmoid", name="classify_1")(d4)

    model = Model([input1, input2, input3],o2)

    return model

def Net_3feature_atten_1binary():

    input1 = Input((3628,49))
    input2 = Input((1574,23))
    input3 = Input((128,256))

    #print("input1:",input1.shape)

    #FAUs

    m1 = Masking(mask_value=0)(input1)

    # ATTENTION PART STARTS HERE
    attention_probs_1 = Dense(49, activation='softmax', name='attention_vec_1')(m1)
    attention_mul_1 = Multiply(name='attention_mul_1')([m1, attention_probs_1])
    # ATTENTION PART FINISHES HERE

    #print("att_1:",attention_mul_1)

    l = Bidirectional(LSTM(32, dropout=0.2, recurrent_dropout=0.2))(attention_mul_1)
    d1 = Dense(32,activation='relu')(l)

    #print('d1:',d1.shape)

    #eGeMAPS

    m2 = Masking(mask_value=0)(input2)

    # ATTENTION PART STARTS HERE
    attention_probs_2 = Dense(23, activation='softmax', name='attention_vec_2')(m2)
    attention_mul_2 = Multiply(name='attention_mul_2')([m2, attention_probs_2])
    # ATTENTION PART FINISHES HERE

    l2 = Bidirectional(LSTM(16, dropout=0.2, recurrent_dropout=0.2))(attention_mul_2)
    d2 = Dense(32,activation="relu")(l2)

    #print('d2:',d2.shape)

    #text
    m3 = Masking(mask_value=0)(input3)

    #print("att_3:",attention_mul_3.shape)

    # ATTENTION PART STARTS HERE
    attention_probs_3 = Dense(23, activation='softmax', name='attention_vec_2')(m3)
    attention_mul_3 = Multiply(name='attention_mul_2')([m3, attention_probs_3])
    # ATTENTION PART FINISHES HERE

    l3 = Bidirectional(LSTM(64, dropout=0.2, recurrent_dropout=0.2))(attention_mul_3)
    d3 = Dense(32, activation='relu')(l3)

    #Concatenate
    d3 = Concatenate()([d1,d2,d3])
    d4 = Dense(64,activation="relu")(d3)

    o2 = Dense(2, activation="sigmoid", name="classify_1")(d4)

    model = Model([input1, input2, input3],o2)

    return model

