import keras
from keras import layers
## VGG encodeur, cette partie ne change pas
vgg = VGG16(weights='imagenet', include_top=False, input_shape=(height,width,3))

#### Premier block d'upsampling avec skip-connections #####
vgg_output = vgg.get_layer('block5_conv3').output
# un upsampling comme dans l'exercice précédent
up1 = UpSampling2D(size=(2, 2))(vgg_output)
# Récupération des couches de l'encodeur dont la taille correspond à up1.
block4_conv1 = vgg.get_layer('block4_conv1').output
block4_conv2 = vgg.get_layer('block4_conv2').output
block4_conv3 = vgg.get_layer('block4_conv3').output

# Concaténation des différentes couches
merge1 = layers.concatenate([up1, block4_conv1, block4_conv2, block4_conv3])

# Une couche de convolution dont le rôle est d'apprendre à fusionner ces couches
conv1 = Conv2D(filters=128, kernel_size=(5, 5), padding='same', activation='relu')(merge1)

##### Suite du décodeur ####

"""
for i in range(4, 0, -1) :
    for j in range(1, 4) :
        try :
            vgl = vgg.get_layer('block'+str(i+1)+'_conv'+str(j)).output
            break
        except :
            pass
    # un upsampling comme dans l'exercice précédent
    up = UpSampling2D(size=(2, 2))(vgl)
    # Récupération des couches de l'encodeur dont la taille correspond à up.
    #blockConv = [vgg.get_layer('block'+str(i)+'_conv'+str(j)).output for j in range(1, 4)]
    blaockConv = [up]
    for j in range(1, 4) :
        try :
            tmp = vgg.get_layer('block'+str(i)+'_conv'+str(j)).output
            blockConv.append(tmp)
        except :
            pass

# Concaténation des différentes couches
merge = keras.layers.concatenate(blockConv)

# Une couche de convolution dont le rôle est d'apprendre à fusionner ces couches
convs.append(Conv2D(filters=128, kernel_size=(5, 5), padding='same', activation='relu')(merge))

"""

#### Deuxième block d'upsampling avec skip-connections #####
#vgg_output = conv1#vgg.get_layer('block4_conv3').output
# un upsampling comme dans l'exercice précédent
up2 = UpSampling2D(size=(2, 2))(conv1)
# Récupération des couches de l'encodeur dont la taille correspond à up1.
block3_conv1 = vgg.get_layer('block3_conv1').output
block3_conv2 = vgg.get_layer('block3_conv2').output
block3_conv3 = vgg.get_layer('block3_conv3').output

# Concaténation des différentes couches
merge2 = layers.concatenate([up2, block3_conv1, block3_conv2, block3_conv3])

# Une couche de convolution dont le rôle est d'apprendre à fusionner ces couches
conv2 = Conv2D(filters=128, kernel_size=(5, 5), padding='same', activation='relu')(merge2)


#### Troisième block d'upsampling avec skip-connections #####
#vgg_output = conv2#vgg.get_layer('block3_conv2').output
# un upsampling comme dans l'exercice précédent
up3 = UpSampling2D(size=(2, 2))(conv2)
# Récupération des couches de l'encodeur dont la taille correspond à up1.
block2_conv1 = vgg.get_layer('block2_conv1').output
block2_conv2 = vgg.get_layer('block2_conv2').output

# Concaténation des différentes couches
merge3 = layers.concatenate([up3, block2_conv1, block2_conv2])

# Une couche de convolution dont le rôle est d'apprendre à fusionner ces couches
conv3 = Conv2D(filters=128, kernel_size=(5, 5), padding='same', activation='relu')(merge3)


#### Quatrième block d'upsampling avec skip-connections #####
#vgg_output = conv2#vgg.get_layer('block1_conv2').output
# un upsampling comme dans l'exercice précédent
up4 = UpSampling2D(size=(2, 2))(conv3)
# Récupération des couches de l'encodeur dont la taille correspond à up1.
block1_conv1 = vgg.get_layer('block1_conv1').output
block1_conv2 = vgg.get_layer('block1_conv2').output

# Concaténation des différentes couches
merge4 = layers.concatenate([up4, block1_conv1, block1_conv2])

# Une couche de convolution dont le rôle est d'apprendre à fusionner ces couches
conv4 = Conv2D(filters=128, kernel_size=(5, 5), padding='same', activation='relu')(merge4)
#x = conv4
## Couches finales de classification

fcn_output = SoftmaxMap()(conv4)
model = models.Model(vgg.input,
                         outputs = fcn_output,
                         name="encodeur-décodeur")
model.compile(optimizer=tf.keras.optimizers.Adam(lr=0.00001), loss=crossentropy_for_img_segmentation)
model.summary()