from keras.layers import BatchNormalization, Dense, Activation
from keras.layers import Reshape, Concatenate, Input, Add, Lambda, Multiply
from keras.models import Model, Sequential
import tensorflow as tf
from tqdm import trange
from keras.layers import dot, concatenate
import keras.backend as k


class MDNN2:
    def __init__(self, channels, nfeatures, sizes, activation='tanh'):
        self.nfeatures = nfeatures
        self.atommodels = {}
        for c in channels:
            model = Sequential()
            model.add(Dense(units=sizes[0], input_dim=nfeatures, name='{}_dense_0'.format(c)))
            model.add(Activation(activation, name='{}_activation_0'.format(c)))
            for i, s in enumerate(sizes[1:-1]):
                model.add(Dense(units=s, name='{}_dense_{}'.format(c, i+1)))
                model.add(Activation(activation, name='{}_activation_{}'.format(c, i+1)))
            model.add(Dense(units=1, name='{}_dense_{}'.format(c, len(sizes)-1)))
            model.add(Activation('linear', name='{}_activation_{}'.format(c, len(sizes)-1)))
            self.atommodels[c] = model


    def create_model_features(self, atomelements):

        outputs = []
        in_layer = Input(shape=(len(atomelements), self.nfeatures))

        unpack_layer = Lambda(lambda x: tf.unstack(x, axis=1))
        unpack_out = unpack_layer(in_layer)

        for i, atomin in enumerate(unpack_out):
            if atomelements[i] not in self.atommodels:
                print('Element {} not found in layers'.format(atomelements[i]))
                raise

            output = self.atommodels[atomelements[i]](atomin)
            outputs.append(output)

        final_out = concatenate(outputs, 'concat')

        return Model(in_layer, final_out)