"""
Originally found in 
    https://github.com/nikhilroxtomar/Semantic-Segmentation-Architecture
All credit goes to nikhilroxtomar.

This file is licensed under Apache License 2.0
"""
from keras.layers import (Activation, BatchNormalization, Concatenate, Conv2D,
                          Conv2DTranspose, Input, MaxPool2D, ZeroPadding2D)
from keras.models import Model
from tensorflow.keras import layers
from tensorflow.keras.applications import InceptionResNetV2

from ivus_segmentation.architectures.inception_resnetv2_unet import \
    build_inception_resnetv2_unet


def conv_block(input, num_filters):
    x = Conv2D(num_filters, 3, padding="same")(input)
    x = BatchNormalization()(x)
    x = Activation("relu")(x)

    x = Conv2D(num_filters, 3, padding="same")(x)
    x = BatchNormalization()(x)
    x = Activation("relu")(x)

    return x


def decoder_block(input, skip_features, num_filters):
    x = Conv2DTranspose(num_filters, (2, 2), strides=2, padding="same")(input)
    x = Concatenate()([x, skip_features])
    x = conv_block(x, num_filters)
    return x


def inception_resnetv2_unet(inputs):
    """Pre-trained InceptionResNetV2 Model"""
    encoder = InceptionResNetV2(
        include_top=False, weights="imagenet", input_tensor=inputs
    )

    """ Encoder """
    s1 = encoder.get_layer(inputs.name).output  # (512 x 512)

    s2 = encoder.get_layer("activation").output  # (255 x 255)
    s2 = ZeroPadding2D(((1, 0), (1, 0)))(s2)  # (256 x 256)

    s3 = encoder.get_layer("activation_3").output  # (126 x 126)
    s3 = ZeroPadding2D((1, 1))(s3)  # (128 x 128)

    s4 = encoder.get_layer("activation_74").output  # (61 x 61)
    s4 = ZeroPadding2D(((2, 1), (2, 1)))(s4)  # (64 x 64)

    """ Bridge """
    b1 = encoder.get_layer("activation_161").output  # (30 x 30)
    b1 = ZeroPadding2D((1, 1))(b1)  # (32 x 32)

    """ Decoder """
    d1 = decoder_block(b1, s4, 512)  # (64 x 64)
    d2 = decoder_block(d1, s3, 256)  # (128 x 128)
    d3 = decoder_block(d2, s2, 128)  # (256 x 256)
    d4 = decoder_block(d3, s1, 64)  # (512 x 512)

    # outputs = Conv2D(1, 1, padding="same", activation="sigmoid")(d4)

    # return inputs, outputs

    return d4


def inception_resnetv2_unet_frozen(inputs):
    """Pre-trained InceptionResNetV2 Model"""
    encoder = InceptionResNetV2(
        include_top=False, weights="imagenet", input_tensor=inputs
    )
    for layer in encoder.layers:
        layer.trainable = False

    """ Encoder """
    s1 = encoder.get_layer(inputs.name).output  # (512 x 512)

    s2 = encoder.get_layer("activation").output  # (255 x 255)
    s2 = ZeroPadding2D(((1, 0), (1, 0)))(s2)  # (256 x 256)

    s3 = encoder.get_layer("activation_3").output  # (126 x 126)
    s3 = ZeroPadding2D((1, 1))(s3)  # (128 x 128)

    s4 = encoder.get_layer("activation_74").output  # (61 x 61)
    s4 = ZeroPadding2D(((2, 1), (2, 1)))(s4)  # (64 x 64)

    """ Bridge """
    b1 = encoder.get_layer("activation_161").output  # (30 x 30)
    b1 = ZeroPadding2D((1, 1))(b1)  # (32 x 32)

    """ Decoder """
    d1 = decoder_block(b1, s4, 512)  # (64 x 64)
    d2 = decoder_block(d1, s3, 256)  # (128 x 128)
    d3 = decoder_block(d2, s2, 128)  # (256 x 256)
    d4 = decoder_block(d3, s1, 64)  # (512 x 512)

    # outputs = Conv2D(1, 1, padding="same", activation="sigmoid")(d4)

    # return inputs, outputs

    return d4


def unet(inputs):
    ### [First half of the network: downsampling inputs] ###
    x = layers.Conv2D(32, 3, strides=2, padding="same")(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)

    previous_block_activation = x  # Set aside residual

    # Blocks 1, 2, 3 are identical apart from the feature depth.
    for filters in [64, 128, 256]:
        x = layers.Activation("relu")(x)
        x = layers.SeparableConv2D(filters, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)

        x = layers.Activation("relu")(x)
        x = layers.SeparableConv2D(filters, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)

        x = layers.MaxPooling2D(3, strides=2, padding="same")(x)

        # Project residual
        residual = layers.Conv2D(filters, 1, strides=2, padding="same")(
            previous_block_activation
        )
        x = layers.add([x, residual])  # Add back residual
        previous_block_activation = x  # Set aside next residual

    ### [Second half of the network: upsampling inputs] ###

    for filters in [256, 128, 64, 32]:
        x = layers.Activation("relu")(x)
        x = layers.Conv2DTranspose(filters, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)

        x = layers.Activation("relu")(x)
        x = layers.Conv2DTranspose(filters, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)

        x = layers.UpSampling2D(2)(x)

        # Project residual
        residual = layers.UpSampling2D(2)(previous_block_activation)
        residual = layers.Conv2D(filters, 1, padding="same")(residual)
        x = layers.add([x, residual])  # Add back residual
        previous_block_activation = x  # Set aside next residual

    # Add a per-pixel classification layer
    return x
