import sys
import os
sys.path.append(os.path.abspath("../.."))
from keras import layers, models
import tensorflow as tf
from BiddingKG.dl.table_head.models.my_average_pooling import MyAveragePooling1D
from BiddingKG.dl.table_head.models.self_attention import SeqSelfAttention


def get_model(input_shape, output_shape):
    # Input
    input_1 = layers.Input(shape=input_shape[1:], dtype="float32")
    input_2 = layers.Input(shape=input_shape[1:], dtype="float32")
    input_3 = layers.Input(shape=input_shape[1:], dtype="float32")
    input_4 = layers.Input(shape=input_shape[1:], dtype="float32")
    input_5 = layers.Input(shape=input_shape[1:], dtype="float32")
    input_6 = layers.Input(shape=input_shape[1:], dtype="float32")

    # Bi-LSTM
    bi_lstm_1 = layers.Bidirectional(layers.LSTM(16, return_sequences=True))(input_1)
    bi_lstm_2 = layers.Bidirectional(layers.LSTM(16, return_sequences=True))(input_2)
    bi_lstm_3 = layers.Bidirectional(layers.LSTM(16, return_sequences=True))(input_3)
    bi_lstm_4 = layers.Bidirectional(layers.LSTM(16, return_sequences=True))(input_4)
    bi_lstm_5 = layers.Bidirectional(layers.LSTM(16, return_sequences=True))(input_5)
    bi_lstm_6 = layers.Bidirectional(layers.LSTM(16, return_sequences=True))(input_6)

    # Self-Attention
    self_attention_1 = SeqSelfAttention(attention_activation='sigmoid')(bi_lstm_1)
    self_attention_2 = SeqSelfAttention(attention_activation='sigmoid')(bi_lstm_2)
    self_attention_3 = SeqSelfAttention(attention_activation='sigmoid')(bi_lstm_3)
    self_attention_4 = SeqSelfAttention(attention_activation='sigmoid')(bi_lstm_4)
    self_attention_5 = SeqSelfAttention(attention_activation='sigmoid')(bi_lstm_5)
    self_attention_6 = SeqSelfAttention(attention_activation='sigmoid')(bi_lstm_6)

    # Concat
    concat_1 = layers.concatenate([self_attention_1, self_attention_2, self_attention_3])
    concat_2 = layers.concatenate([self_attention_4, self_attention_5, self_attention_6])

    # Dense + Sigmoid
    dense_1 = layers.Dense(output_shape[0], activation="sigmoid")(concat_1)
    dense_2 = layers.Dense(output_shape[0], activation="sigmoid")(concat_2)

    # mask mean pooling
    pool_1 = MyAveragePooling1D(axis=1)(dense_1)
    pool_2 = MyAveragePooling1D(axis=1)(dense_2)

    # Concat
    concat = layers.concatenate([pool_1, pool_2])

    # Dense
    output = layers.Dense(10)(concat)
    output = layers.Dense(1, activation="sigmoid", name='output')(output)

    model = models.Model(inputs=[input_1, input_2, input_3, input_4, input_5, input_6],
                         outputs=output)

    model.summary()
    return model