from collections import namedtuple

import keras
from keras.activations import elu
from keras.engine import Input, Model as KerasModel
from keras.layers import Conv1D, Dense, Dropout, Embedding, GlobalAveragePooling1D, GlobalMaxPooling1D, TimeDistributed

import dataset


def selu(x):
    """Scaled Exponential Linear Unit. (Klambauer et al., 2017)
    # Arguments
        x: A tensor or variable to compute the activation function for.
    # References
        - [Self-Normalizing Neural Networks](https://arxiv.org/abs/1706.02515)
    # copied from keras.io
    """
    alpha = 1.6732632423543772848170429916717
    scale = 1.0507009873554804934193349852946
    return scale * elu(x, alpha)


Model = namedtuple("Model", ["in_domains", "in_flows", "out_client", "out_server"])


def get_embedding(embedding_size, input_length, filter_size, kernel_size, hidden_dims, drop_out=0.5) -> KerasModel:
    x = y = Input(shape=(input_length,))
    y = Embedding(input_dim=dataset.get_vocab_size(), output_dim=embedding_size)(y)
    y = Conv1D(filter_size,
               kernel_size,
               activation=selu)(y)
    y = GlobalAveragePooling1D()(y)
    y = Dense(hidden_dims, activation=selu)(y)
    return KerasModel(x, y)


def get_model(cnnDropout, flow_features, domain_features, window_size, domain_length, cnn_dims, kernel_size,
              dense_dim, cnn, model_output="both") -> Model:
    ipt_domains = Input(shape=(window_size, domain_length), name="ipt_domains")
    encoded = TimeDistributed(cnn, name="domain_cnn")(ipt_domains)
    ipt_flows = Input(shape=(window_size, flow_features), name="ipt_flows")
    merged = keras.layers.concatenate([encoded, ipt_flows], -1)
    # CNN processing a small slides of flow windows
    y = Conv1D(cnn_dims,
               kernel_size,
               activation=selu,
               input_shape=(window_size, domain_features + flow_features))(merged)
    # remove temporal dimension by global max pooling
    y = GlobalMaxPooling1D()(y)
    y = Dropout(cnnDropout)(y)
    y = Dense(dense_dim, activation=selu)(y)
    out_client = Dense(1, activation='sigmoid', name="client")(y)
    out_server = Dense(1, activation='sigmoid', name="server")(y)
    
    return Model(ipt_domains, ipt_flows, out_client, out_server)


def get_new_model(dropout, flow_features, domain_features, window_size, domain_length, cnn_dims, kernel_size,
                  dense_dim, cnn, model_output="both") -> Model:
    ipt_domains = Input(shape=(window_size, domain_length), name="ipt_domains")
    ipt_flows = Input(shape=(window_size, flow_features), name="ipt_flows")
    encoded = TimeDistributed(cnn, name="domain_cnn")(ipt_domains)
    merged = keras.layers.concatenate([encoded, ipt_flows], -1)
    y = Dense(dense_dim, activation=selu)(merged)
    out_server = Dense(1, activation="sigmoid", name="server")(y)
    merged = keras.layers.concatenate([merged, y], -1)
    # CNN processing a small slides of flow windows
    y = Conv1D(filters=cnn_dims,
               kernel_size=kernel_size,
               activation=selu,
               padding="same",
               input_shape=(window_size, domain_features + flow_features))(merged)
    # remove temporal dimension by global max pooling
    y = GlobalMaxPooling1D()(y)
    y = Dropout(dropout)(y)
    y = Dense(dense_dim,
              activation=selu,
              name="dense_client")(y)
    out_client = Dense(1, activation='sigmoid', name="client")(y)
    
    return Model(ipt_domains, ipt_flows, out_client, out_server)