import os
import sys
import glob
import numpy as np

import tensorflow as tf
from keras import Model
from keras.regularizers import l2
from keras import backend as K
from keras.engine import Layer,InputSpec
from keras.layers.merge import concatenate
from keras.callbacks import TensorBoard,Callback
from keras.layers.advanced_activations import LeakyReLU
from keras.preprocessing.image import ImageDataGenerator
from keras.layers.normalization import BatchNormalization
from keras.callbacks import ModelCheckpoint, ReduceLROnPlateau
from keras import initializers, regularizers, constraints,optimizers
from keras.callbacks import ModelCheckpoint, LearningRateScheduler,TensorBoard
from keras.layers import Add,Input,Conv3D,Convolution3D,Dropout,UpSampling3D,Concatenate,MaxPooling3D,\
GlobalAveragePooling3D,Dense,GlobalMaxPooling3D,Lambda,Activation,Reshape,Permute, PReLU, Deconvolution3D

sys.path.append('../baseLayers/')
from CommonLayers import NormActi,ConvUnit,ASPP


def NoduleSegConvBlock_proxima(x,**kwargs):
    block_idx = kwargs.get('block_idx',1)
    strides = kwargs.get('strides')
    atrous_rates = kwargs.get('atrous_rates')
    num_filters = kwargs.get('num_filters')
    norm_func = kwargs.get('norm_func')
    activation_func = kwargs.get('activation_func')
    kernel_size = kwargs.get('kernel_size',3)
    padding = kwargs.get('padding','same')
    conv_first = kwargs.get('conv_first',True)
    kernel_initializer = kwargs.get('kernel_initializer','he_normal')
    kernel_regularizer = kwargs.get('kernel_regularizer',l2(1e-4))
    ASPP_choice = kwargs.get('ASPP_choice',False)
    
    
    repeat_times = 3
    block_prefix = 'NoduleSegConvBlock%02d'%block_idx

    tensors = []
    for layer_idx in range(repeat_times):
        if layer_idx == 0:
            conv_stride = int(strides)
        else:
            conv_stride = 1
        if ASPP_choice:
            x = ASPP(x,atrous_rate=atrous_rates,norm_func=norm_func,activation_func=activation_func,num_filters=int(num_filters/4),
                    kenrel_size=kernel_size,padding=padding,block_idx=block_idx,kernel_initializer=kernel_initializer,
                    kernel_regularizer=kernel_regularizer,layer_idx=layer_idx+1)

        x = ConvUnit(x,norm_func=norm_func,activation_func=activation_func,num_filters=num_filters,
                    kernel_size=kernel_size,conv_stride=conv_stride,atrous_rate=atrous_rates,padding=padding,
                     block_prefix=block_prefix,kernel_initializer=kernel_initializer,kernel_regularizer=kernel_regularizer,
                    layer_idx=layer_idx+1,conv_first=conv_first)
        
        tensors.append(x)
    return Add(name='NoduleSegConvBlock%02d_add'%block_idx)([tensors[0],tensors[-1]])


def NoduleSegEncoder_proxima(input_shape,**kwargs):
    num_blocks = kwargs.get('num_blocks',5)
    
    strides = kwargs.get('strides')
    atrous_rates = kwargs.get('atrous_rates')
    base_filters = kwargs.get('base_filters',32)
    norm_func = kwargs.get('norm_func')
    activation_func = kwargs.get('activation_func')
    kernel_size = kwargs.get('kernel_size',3)
    padding = kwargs.get('padding','same')
    dropout_rate = kwargs.get('dropout_rate',0.5)
    
    classification_layers = kwargs.get('classification_layers')
    num_classes = kwargs.get('num_classes',1)
    kernel_initializer = kwargs.get('kernel_initializer','he_normal')
    kernel_regularizer = kwargs.get('kernel_regularizer',l2(1e-4))
    ASPP_choice = kwargs.get('ASPP_choice',False)
    ASPP_block_list = kwargs.get('ASPP_block_list',[num_blocks-1])
    
    if type(dropout_rate)== float or type(dropout_rate)== int:
        dropout_rates = [dropout_rate for _ in range(num_blocks)]
    else:
        dropout_rates = dropout_rate
    input_tensor = Input(shape=input_shape)
    tensors = [input_tensor]
    acc_stride = 1
    enlarge_ratio = 1
    real_stride_ratios = []
    last_enlarge = 1
    print ('num_blocks is ',num_blocks)
    for block_idx in range(num_blocks):
        current_ASPP_choice = ASPP_choice if block_idx in ASPP_block_list else False
        if strides[block_idx]>1:
            enlarge_ratio *= strides[block_idx]
            current_stride = int(enlarge_ratio/last_enlarge/atrous_rates[block_idx])
        else:
            current_stride = strides[block_idx]
        last_enlarge = current_stride
        real_stride_ratios.append(current_stride)
        current_tensor = NoduleSegConvBlock_proxima(tensors[-1],activation_func=activation_func,
                                       norm_func=norm_func,strides=current_stride,
                                       atrous_rates=atrous_rates[block_idx],kernel_initializer=kernel_initializer,
                                      kernel_regularizer=kernel_regularizer,kernel_size=kernel_size,
                                      padding=padding,block_idx=block_idx+1,
                                       num_filters = base_filters*(2**block_idx),
                                       ASPP_choice=current_ASPP_choice)
        current_tensor = Dropout(rate=dropout_rates[block_idx],name='NoduleSegEncoderDropout%02d'%(block_idx+1))(current_tensor)
        tensors.append(current_tensor)

    final_feature = classification_layers(name='NoduleSegEncoder_classification')(tensors[-1])
    if num_classes == 1:
        final_acti = 'sigmoid'
    else:
        final_acti = 'softmax'
    final_output = Dense(num_classes,activation=final_acti,name='NoduleSegEncoder_final_output')(final_feature)
    
    
    model = Model(input_tensor,final_output)
    
    return model,tensors,real_stride_ratios