# Modified from https://github.com/flyyufelix/DenseNet-Keras/blob/master/densenet121.py

from keras.models import Model
from keras.layers import Input, merge, ZeroPadding2D
from keras.layers.core import Dense, Dropout, Activation
from keras.layers.convolutional import Convolution2D
from keras.layers.pooling import AveragePooling2D, GlobalAveragePooling2D, MaxPooling2D
from keras.layers.normalization import BatchNormalization
import keras.backend as K

import tensorflow as tf

from utils.custom_layers import Scale

def ParallelDenseNet121(num_inputs=4, input_size=224, nchannels=3, nb_dense_block=4, growth_rate=32, nb_filter=64, reduction=0.0, dropout_rate=0.0, weight_decay=1e-4, classes=1000, num_gpu=1):
    '''Instantiate the DenseNet 121 architecture,
        # Arguments
            nb_dense_block: number of dense blocks to add to end
            growth_rate: number of filters to add per dense block
            nb_filter: initial number of filters
            reduction: reduction factor of transition blocks.
            dropout_rate: dropout rate
            weight_decay: weight decay factor
            classes: optional number of classes to classify images
            weights_path: path to pre-trained weights
        # Returns
            A Keras model instance.
    '''
    eps = 1.1e-5

    # compute compression factor
    compression = 1.0 - reduction

    # Handle Dimension Ordering for different backends
    global concat_axis
    if K.image_dim_ordering() == 'tf':
        concat_axis = 3
        img_inputs = [Input(shape=(input_size, input_size, nchannels), name=f'input_{i}') for i in range(num_inputs)]
    else:
        concat_axis = 1
        img_inputs = [Input(shape=(nchannels, input_size, input_size), name='data') for i in range(num_inputs)]

    # From architecture for ImageNet (Table 1 in the paper)
    nb_layers = [6,12,24,16] # For DenseNet-121

    # Initial convolution
    init_conv_layers = [
        ZeroPadding2D((3, 3)),
        Convolution2D(nb_filter, 7, 7, subsample=(2, 2), bias=False),
        BatchNormalization(epsilon=eps, axis=concat_axis),
        Scale(axis=concat_axis),
        Activation('relu'),
        ZeroPadding2D((1, 1)),
        MaxPooling2D((3, 3), strides=(2, 2), name='pool1')
    ]
    x = img_inputs.copy()
    x = allocate_layers(x=x, layers=init_conv_layers, num_gpu=num_gpu)
#     for j in range(len(x)):
#         for layer in init_conv_layers:
#             x[j] = layer(x[j])

    # Add dense blocks
    for block_idx in range(nb_dense_block - 1):
        stage = block_idx+2
        x, nb_filter = dense_block(x, stage, nb_layers[block_idx], nb_filter, growth_rate, dropout_rate=dropout_rate, weight_decay=weight_decay, num_gpu=num_gpu)

        # Add transition_block
        x = transition_block(x, stage, nb_filter, compression=compression, dropout_rate=dropout_rate, weight_decay=weight_decay, num_gpu=num_gpu)
        nb_filter = int(nb_filter * compression)

    final_stage = stage + 1
    x, nb_filter = dense_block(x, final_stage, nb_layers[-1], nb_filter, growth_rate, dropout_rate=dropout_rate, weight_decay=weight_decay, num_gpu=num_gpu)

    top_activation = 'softmax' if classes > 1 else 'sigmoid'
    
    finl_conv_layers = [
        BatchNormalization(epsilon=eps, axis=concat_axis),
        Scale(axis=concat_axis),
        Activation('relu'),
        GlobalAveragePooling2D(name='pool'+str(final_stage)),
        Dense(classes),
        Activation(top_activation, name='prob')
    ]
    x = allocate_layers(x=x, layers=finl_conv_layers, num_gpu=num_gpu)
#     for j in range(len(x)):
#         for layer in finl_conv_layers:
#             x[j] = layer(x[j])

    x = merge.maximum(x)
    model = Model(img_inputs, x, name='parallel_dense')
    
    return model


def conv_block(x, stage, branch, nb_filter, dropout_rate=None, weight_decay=1e-4, num_gpu=1):
    '''Apply BatchNorm, Relu, bottleneck 1x1 Conv2D, 3x3 Conv2D, and option dropout
        # Arguments
            x: input tensor 
            stage: index for dense block
            branch: layer index within each dense block
            nb_filter: number of filters
            dropout_rate: dropout rate
            weight_decay: weight decay factor
    '''
    eps = 1.1e-5
    conv_name_base = 'conv' + str(stage) + '_' + str(branch)
    relu_name_base = 'relu' + str(stage) + '_' + str(branch)

    # 1x1 Convolution (Bottleneck layer)
    inter_channel = nb_filter * 4  
    layers = []
    layers.extend([
        BatchNormalization(epsilon=eps, axis=concat_axis),
        Scale(axis=concat_axis),
        Activation('relu'),
        Convolution2D(inter_channel, 1, 1, bias=False)
    ])
    
    if dropout_rate:
        layers.append(Dropout(dropout_rate))

    # 3x3 Convolution
    layers.extend([
        BatchNormalization(epsilon=eps, axis=concat_axis),
        Scale(axis=concat_axis),
        Activation('relu'),
        ZeroPadding2D((1, 1)),
        Convolution2D(nb_filter, 3, 3, bias=False)
    ])

    if dropout_rate:
        layers.append(Dropout(dropout_rate))
        
    x = allocate_layers(x=x, layers=layers, num_gpu=num_gpu)
#     for j in range(len(x)):
#         for layer in layers:
#             x[j] = layer(x[j])
            
    return x


def transition_block(x, stage, nb_filter, compression=1.0, dropout_rate=None, weight_decay=1E-4, num_gpu=1):
    ''' Apply BatchNorm, 1x1 Convolution, averagePooling, optional compression, dropout 
        # Arguments
            x: input tensor
            stage: index for dense block
            nb_filter: number of filters
            compression: calculated as 1 - reduction. Reduces the number of feature maps in the transition block.
            dropout_rate: dropout rate
            weight_decay: weight decay factor
    '''

    eps = 1.1e-5
    conv_name_base = 'conv' + str(stage) + '_blk'
    relu_name_base = 'relu' + str(stage) + '_blk'
    pool_name_base = 'pool' + str(stage) 

    layers = []
    layers.extend([
        BatchNormalization(epsilon=eps, axis=concat_axis),
        Scale(axis=concat_axis),
        Activation('relu'),
        Convolution2D(int(nb_filter * compression), 1, 1, bias=False)
    ])

    if dropout_rate:
        layers.append(Dropout(dropout_rate))

    layers.append(AveragePooling2D((2, 2), strides=(2, 2), name=pool_name_base))

    x = allocate_layers(x=x, layers=layers, num_gpu=num_gpu)
#     for j in range(len(x)):
#         for layer in layers:
#             x[j] = layer(x[j])
            
    return x


def dense_block(x, stage, nb_layers, nb_filter, growth_rate, dropout_rate=None, weight_decay=1e-4, grow_nb_filters=True, num_gpu=1):
    ''' Build a dense_block where the output of each conv_block is fed to subsequent ones
        # Arguments
            x: input tensor
            stage: index for dense block
            nb_layers: the number of layers of conv_block to append to the model.
            nb_filter: number of filters
            growth_rate: growth rate
            dropout_rate: dropout rate
            weight_decay: weight decay factor
            grow_nb_filters: flag to decide to allow number of filters to grow
    '''

    eps = 1.1e-5
    concat_feat = x

    for i in range(nb_layers):
        branch = i+1
        x = conv_block(concat_feat, stage, branch, growth_rate, dropout_rate, weight_decay, num_gpu=num_gpu)
        concat_feat = allocate_concat(x, concat_feat, num_gpu=num_gpu)
#         for j in range(len(x)):
#             concat_feat[j] = merge.concatenate([concat_feat[j], x[j]], axis=concat_axis)

        if grow_nb_filters:
            nb_filter += growth_rate

    return concat_feat, nb_filter


def allocate_layers(x, layers, num_gpu):
    nx = len(x)
    group_size = nx // num_gpu + (nx % num_gpu > 0)
    for ng in range(num_gpu):
        with tf.device(f'/GPU:{ng}'):
            for j in range(ng * group_size, min((ng + 1) * group_size, nx)):
                for layer in layers:
                    x[j] = layer(x[j])
    return x

def allocate_concat(x, concat_feat, num_gpu):
    nx = len(x)
    group_size = nx // num_gpu + (nx % num_gpu > 0)
    for ng in range(num_gpu):
        with tf.device(f'/GPU:{ng}'):
            for j in range(ng * group_size, min((ng + 1) * group_size, nx)):
                concat_feat[j] = merge.concatenate([concat_feat[j], x[j]], axis=concat_axis)
    return concat_feat
    