from collections import namedtuple
from enum import Enum, unique
import logging
import lzma
import pickle

from keras.models import Model
from keras.layers import *
from keras.utils.np_utils import to_categorical
from keras import regularizers

MlpLayer = namedtuple('MlpLayer', [
    'units',
    'dropout',
    'l2_regularization'
])

@unique
class MlpFeature(Enum):
    AVG_SENTENCE = 1
    FEE = 2
    AVG_SENTENCE_SYNSETS = 12
    AVG_SENTENCE_IMAGINED_PRETRAINED = 13

def mlp_layers_to_name(mlp_layers):
    fmt = lambda layer : 'u={0}&d={1}&l2={2}'.format(layer.units, layer.dropout, layer.l2_regularization)
    layers = ';'.join([fmt(layer) for layer in mlp_layers])
    return 'Dense:{0}'.format(layers)

def build_mlp_from_config(mlp_layers, embeddings_dims, number_of_labels, mlp_features):
    inputs = []
    features = []
    if isinstance(embeddings_dims, dict): # embeddings_dims are given as a dict
        embeddings_dim = embeddings_dims["embeddings_dim"]
        synset_embeddings_dim = embeddings_dims["synset_embeddings_dim"]
        if "imagined_embeddings_dim" in embeddings_dims.keys():
            imagined_embeddings_dim = embeddings_dims["imagined_embeddings_dim"]
    else: # embeddings_dims are given directly as one dim for a textual embedding instance
        embeddings_dim = embeddings_dims  

    if MlpFeature.AVG_SENTENCE in mlp_features:
        avg_sentence_input = Input(shape=(embeddings_dim,), dtype='float32', name=MlpFeature.AVG_SENTENCE.name)
        inputs.append(avg_sentence_input)
        features.append(avg_sentence_input)

    if MlpFeature.FEE in mlp_features:
        fee_input = Input(shape=(embeddings_dim,), dtype='float32', name=MlpFeature.FEE.name)
        inputs.append(fee_input)
        features.append(fee_input) 
        
    if MlpFeature.AVG_SENTENCE_SYNSETS in mlp_features:
        avg_sentence_input = Input(shape=(synset_embeddings_dim,), dtype='float32', name=MlpFeature.AVG_SENTENCE_SYNSETS.name)
        inputs.append(avg_sentence_input)
        features.append(avg_sentence_input)
        
    if MlpFeature.AVG_SENTENCE_IMAGINED_PRETRAINED in mlp_features:
        avg_sentence_input = Input(shape=(imagined_embeddings_dim,), dtype='float32', name=MlpFeature.AVG_SENTENCE_IMAGINED_PRETRAINED.name)
        inputs.append(avg_sentence_input)
        features.append(avg_sentence_input)

    if len(features) > 1:
        x = concatenate(features, name='merged_features') 
    else:
        x = features[0]
    
    # first_layer, hidden_layers = mlp_layers[0], mlp_layers[1:]

    # Add hidden layer
    for layer in mlp_layers:
        x = Dense(layer.units, activation='relu', kernel_regularizer=regularizers.l2(layer.l2_regularization))(x)
        x = Dropout(layer.dropout)(x)

    # Add output layer
    output = Dense(units=number_of_labels, name='output_layer')(x)

    output = Activation('softmax')(output)

    model = Model(inputs=inputs, outputs=[output])
    model.compile(optimizer='nadam',
                    loss='categorical_crossentropy',
                    metrics=['accuracy'])
    return model    

def load_dataset_mlp(resource_manager, corpus_name, dict_embeddings, lexicon, features):
    list_embeddingNames = []
    if "embeddings" in dict_embeddings.keys():
        embeddings = dict_embeddings["embeddings"]
        list_embeddingNames.append(embeddings.vsm_name)
    if "synset_embeddings" in dict_embeddings.keys():
        synset_embeddings = dict_embeddings["synset_embeddings"]
        list_embeddingNames.append(synset_embeddings.vsm_name)
        dataset_path = resource_manager.get_multimodal_dataset(corpus_name, list_embeddingNames)
    if "imagined_embeddings" in dict_embeddings.keys():
        imagined_embeddings = dict_embeddings["imagined_embeddings"]
        list_embeddingNames.append(imagined_embeddings.vsm_name)
        dataset_path = resource_manager.get_multimodal_dataset(corpus_name, list_embeddingNames)
    if "synset_embeddings" not in dict_embeddings.keys() and "imagined_embeddings" not in dict_embeddings.keys(): # unimodal
        dataset_path = resource_manager.get_dataset(corpus_name, embeddings.vsm_name)
    number_of_labels = lexicon.get_number_of_labels()
    
    with lzma.open(dataset_path, 'rb') as f:
        dataset = pickle.load(f)

    X = {}

    # Build average sentence embedding
    if MlpFeature.AVG_SENTENCE in features:
        X_avg_sentence = np.zeros([dataset.n, dataset.dim])

        for i, sentence_embedding in enumerate(dataset.token_embeddings):
            nonzero_tokens = [x for x in sentence_embedding if np.count_nonzero(x)]
            X_avg_sentence[i] = np.mean(nonzero_tokens, axis=0)

        X[MlpFeature.AVG_SENTENCE.name] = X_avg_sentence

    # Build fee embedding vector
    if MlpFeature.FEE in features:
        X_fee = np.zeros([dataset.n, dataset.dim])

        for i, fee_embedding in enumerate(dataset.fee_embeddings):
            X_fee[i] = fee_embedding

        X[MlpFeature.FEE.name] = X_fee                
    
    # Build average sentence embedding for synsets
    if MlpFeature.AVG_SENTENCE_SYNSETS in features:
        X_avg_sentence = np.zeros([dataset.n, dataset.synset_dim])

        for i, sentence_embedding in enumerate(dataset.synset_token_embeddings):
            X_avg_sentence[i] = np.mean(sentence_embedding, axis=0)

        X[MlpFeature.AVG_SENTENCE_SYNSETS.name] = X_avg_sentence
        
    # Build average sentence embedding for imagined pretrained
    if MlpFeature.AVG_SENTENCE_IMAGINED_PRETRAINED in features:
        X_avg_sentence = np.zeros([dataset.n, dataset.imagined_dim])

        for i, sentence_embedding in enumerate(dataset.imagined_token_embeddings):
            nonzero_tokens = [x for x in sentence_embedding if np.count_nonzero(x)]
            X_avg_sentence[i] = np.mean(nonzero_tokens, axis=0)

        X[MlpFeature.AVG_SENTENCE_IMAGINED_PRETRAINED.name] = X_avg_sentence
    
    # Build labels

    Y = np.zeros(dataset.n)
    for i, label in enumerate(dataset.labels):
        Y[i] = lexicon.get_id(label)
    Y_onehot = to_categorical(Y, number_of_labels)    

    logging.info('Shape of [%s] X tensor: %s', corpus_name, [X[name].shape for name in X.keys()])
    logging.info('Shape of [%s] y tensor: %s', corpus_name, Y_onehot.shape) 

    return dataset, X, Y_onehot