import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from keras.layers import Dropout


def residual_block(y, nb_channels, _strides=(1, 1, 1), _kernel_size=(1, 1, 1), _project_shortcut=True):
    shortcut = y

    y = layers.Conv3D(nb_channels, kernel_size=_kernel_size, strides=_strides, padding='same')(y)

    if _project_shortcut or _strides != (1, 1, 1):
        shortcut = layers.Conv3D(nb_channels, kernel_size=(1, 1, 1), strides=_strides, padding='same')(shortcut)

    y = layers.add([shortcut, y])
    y = layers.ReLU()(y)
    return y
    

def network(grad_directions, output_size, dropout):

    inputs = keras.Input(shape=(3,3,3,grad_directions,),  name='diffusion_data')
    x = residual_block(inputs, 256)
    x = layers.Conv3D(512, kernel_size=(2,2,2))(x)
    x = residual_block(x, 512)
    x = layers.LayerNormalization()(x)
    x = layers.Flatten()(x)
    i = inputs[:,1,1,1]

    x2 = layers.Reshape((1,1,1,100))(i)
    x2 = residual_block(x2,704)
    x2 = layers.LayerNormalization()(x2)
    x2 = layers.Flatten()(x2)

    f = layers.Concatenate(axis=-1)([x,x2])
    f = layers.Reshape((40,40,3))(f)
    encoder = keras.Model(inputs = inputs, outputs = f, name="encoder")

    pi = keras.Input(shape=(40,40,3,),  name='encoded_data')
    x = layers.Conv2D(32, (3,3), activation = 'relu')(pi)
    x = layers.LayerNormalization()(x)
    x = layers.Conv2D(64, (3,3), activation = 'relu')(x)
    x = layers.Conv2D(64, (3,3), activation = 'relu')(x)
    x = layers.LayerNormalization()(x)
    x = layers.Dropout(dropout)(x)
    x = layers.MaxPool2D(pool_size=(2,2), strides = None)(x)
    x = layers.Conv2D(128, (3,3), activation = 'relu')(x)
    x = layers.Conv2D(128, (3,3), activation = 'relu')(x)
    x = layers.MaxPool2D(pool_size=(2,2), strides = None)(x)
    x = layers.LayerNormalization()(x)
    x = layers.Conv2D(256, (3,3), activation = 'relu')(x)
    x = layers.Conv2D(256, (3,3), padding = 'same', activation = 'relu')(x)
    x = layers.Conv2D(512, (3,3), activation = 'relu')(x)
    x = layers.Flatten()(x)
    x = layers.Dropout(dropout)(x)
    out = layers.Dense(output_size, activation = 'linear')(x)
    predictor = keras.Model(inputs = pi, outputs = out, name="predictor")

    features = encoder(inputs)
    preds = predictor(features)
    model = keras.Model(inputs = inputs, outputs = preds, name="Model")

    return model
    

tensorflow.keras.layers.LayerNormalization tensorflow.keras.layers.Concatenate tensorflow.keras.layers.Dropout tensorflow.keras.layers.ReLU tensorflow.keras.Input tensorflow.keras.layers.Dense tensorflow.keras.layers.Conv2D tensorflow.keras.layers.Conv3D tensorflow.keras.layers.MaxPool2D tensorflow.keras.Model tensorflow.keras.layers.add tensorflow.keras.layers.Reshape tensorflow.keras.layers.Flatten
[(15, 'tensorflow.keras.layers.add', 'layers.add', (['[shortcut, y]'], {}), False, 'from tensorflow.keras import layers\n'), (22, 'tensorflow.keras.Input', 'keras.Input', ([], {'shape': '(3, 3, 3, grad_directions)', 'name': '"""diffusion_data"""'}), False, 'from tensorflow import keras\n'), (37, 'tensorflow.keras.Model', 'keras.Model', ([], {'inputs': 'inputs', 'outputs': 'f', 'name': '"""encoder"""'}), False, 'from tensorflow import keras\n'), (39, 'tensorflow.keras.Input', 'keras.Input', ([], {'shape': '(40, 40, 3)', 'name': '"""encoded_data"""'}), False, 'from tensorflow import keras\n'), (57, 'tensorflow.keras.Model', 'keras.Model', ([], {'inputs': 'pi', 'outputs': 'out', 'name': '"""predictor"""'}), False, 'from tensorflow import keras\n'), (61, 'tensorflow.keras.Model', 'keras.Model', ([], {'inputs': 'inputs', 'outputs': 'preds', 'name': '"""Model"""'}), False, 'from tensorflow import keras\n'), (10, 'tensorflow.keras.layers.Conv3D', 'layers.Conv3D', (['nb_channels'], {'kernel_size': '_kernel_size', 'strides': '_strides', 'padding': '"""same"""'}), False, 'from tensorflow.keras import layers\n'), (16, 'tensorflow.keras.layers.ReLU', 'layers.ReLU', ([], {}), False, 'from tensorflow.keras import layers\n'), (24, 'tensorflow.keras.layers.Conv3D', 'layers.Conv3D', (['(512)'], {'kernel_size': '(2, 2, 2)'}), False, 'from tensorflow.keras import layers\n'), (26, 'tensorflow.keras.layers.LayerNormalization', 'layers.LayerNormalization', ([], {}), False, 'from tensorflow.keras import layers\n'), (27, 'tensorflow.keras.layers.Flatten', 'layers.Flatten', ([], {}), False, 'from tensorflow.keras import layers\n'), (30, 'tensorflow.keras.layers.Reshape', 'layers.Reshape', (['(1, 1, 1, 100)'], {}), False, 'from tensorflow.keras import layers\n'), (32, 'tensorflow.keras.layers.LayerNormalization', 'layers.LayerNormalization', ([], {}), False, 'from tensorflow.keras import layers\n'), (33, 'tensorflow.keras.layers.Flatten', 'layers.Flatten', ([], {}), False, 'from tensorflow.keras import layers\n'), (35, 'tensorflow.keras.layers.Concatenate', 'layers.Concatenate', ([], {'axis': '(-1)'}), False, 'from tensorflow.keras import layers\n'), (36, 'tensorflow.keras.layers.Reshape', 'layers.Reshape', (['(40, 40, 3)'], {}), False, 'from tensorflow.keras import layers\n'), (40, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(32)', '(3, 3)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (41, 'tensorflow.keras.layers.LayerNormalization', 'layers.LayerNormalization', ([], {}), False, 'from tensorflow.keras import layers\n'), (42, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(64)', '(3, 3)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (43, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(64)', '(3, 3)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (44, 'tensorflow.keras.layers.LayerNormalization', 'layers.LayerNormalization', ([], {}), False, 'from tensorflow.keras import layers\n'), (45, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (46, 'tensorflow.keras.layers.MaxPool2D', 'layers.MaxPool2D', ([], {'pool_size': '(2, 2)', 'strides': 'None'}), False, 'from tensorflow.keras import layers\n'), (47, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(128)', '(3, 3)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (48, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(128)', '(3, 3)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (49, 'tensorflow.keras.layers.MaxPool2D', 'layers.MaxPool2D', ([], {'pool_size': '(2, 2)', 'strides': 'None'}), False, 'from tensorflow.keras import layers\n'), (50, 'tensorflow.keras.layers.LayerNormalization', 'layers.LayerNormalization', ([], {}), False, 'from tensorflow.keras import layers\n'), (51, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(256)', '(3, 3)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (52, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(256)', '(3, 3)'], {'padding': '"""same"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (53, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(512)', '(3, 3)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (54, 'tensorflow.keras.layers.Flatten', 'layers.Flatten', ([], {}), False, 'from tensorflow.keras import layers\n'), (55, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (56, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['output_size'], {'activation': '"""linear"""'}), False, 'from tensorflow.keras import layers\n'), (13, 'tensorflow.keras.layers.Conv3D', 'layers.Conv3D', (['nb_channels'], {'kernel_size': '(1, 1, 1)', 'strides': '_strides', 'padding': '"""same"""'}), False, 'from tensorflow.keras import layers\n')]