import os
import sys
import glob
import numpy as np
import pandas as pd

import keras
import tensorflow as tf
from keras import Model
from keras import backend as K
from keras.engine import Layer,InputSpec
from keras.layers.merge import concatenate
from InstanceNorm import InstanceNormalization
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,Multiply,GlobalAveragePooling3D,Dense


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



def BasicBlock(x,**kwargs):
    block_pre = kwargs.get('block_pre')
    block_number = kwargs.get('block_number')
    num_filters = kwargs.get('num_filters')
    kernel_size = kwargs.get('kernel_size')
    downsample = kwargs.get('downsample')
    leaky_ratio = kwargs.get('leaky_ratio',0.01)
    norm_func = kwargs.get('norm_func',BatchNormalization)
    expansion = kwargs.get('expansion',1)
    stride = kwargs.get('stride',1)
    activation_func = kwargs.get('activation_func')
    atrous_rate = kwargs.get('atrous_rate',1)
    padding = kwargs.get('padding','same')
    kernel_initializer = kwargs.get('kernel_initializer','he_normal')
    kernel_regularizer = kwargs.get('kernel_regularizer',None)
    

    residual = x
    
    block_prefix = block_pre+'_block%02d'%block_number
        
    kernel_sizes = [kernel_size,kernel_size]
    num_filter_list = [num_filters,num_filters]
    strides = [1,stride]
    
    out_result = x
    for idx in range(2):
        out_result = ConvUnit(out_result,norm_func=norm_func,activation_func=activation_func,num_filters=num_filter_list[idx],conv_stride=strides[idx],
                          kernel_size=kernel_sizes[idx],atrous_rate=atrous_rate,padding=padding,block_prefix=block_prefix,
                        kernel_initializer=kernel_initializer,kernel_regularizer=kernel_regularizer,
                        layer_idx=idx+1,conv_first=True)
        
    
    if len(downsample)>0:
        for idx in range(len(downsample)):
            x = downsample[idx](x)
        residual = x
    
    residual = keras.layers.add([residual, out_result])
    residual = LeakyReLU(alpha=leaky_ratio, name='%s_relu%d_2'%(block_pre,block_number))(residual)
    return residual


def Bottleneck(x,**kwargs):
    '''
    input_tensor.shape[-1] should be 4 times as output_tensor.shape[-1]
    '''
    block_pre = kwargs.get('block_pre')
    block_number = kwargs.get('block_number')
    num_filters = kwargs.get('num_filters')
    kernel_size = kwargs.get('kernel_size')
    downsample = kwargs.get('downsample')
    leaky_ratio = kwargs.get('leaky_ratio',0.01)
    norm_func = kwargs.get('norm_func',BatchNormalization)
    expansion = kwargs.get('expansion',4)
    stride = kwargs.get('stride',1)
    activation_func = kwargs.get('activation_func')
    atrous_rate = kwargs.get('atrous_rate',1)
    padding = kwargs.get('padding','same')
    kernel_initializer = kwargs.get('kernel_initializer','he_normal')
    kernel_regularizer = kwargs.get('kernel_regularizer',None)
    
    residual = x
    block_prefix = block_pre+'_block%02d'%block_number
        
    kernel_sizes = [1,kernel_size,kernel_size]
    num_filter_list = [num_filters,num_filters,num_filters*expansion]
    strides = [1,stride,1]
    
    out_result = x
    for idx in range(3):
        out_result = ConvUnit(out_result,norm_func=norm_func,activation_func=activation_func,num_filters=num_filter_list[idx],conv_stride=strides[idx],
                          kernel_size=kernel_sizes[idx],atrous_rate=atrous_rate,padding=padding,block_prefix=block_prefix,
                        kernel_initializer=kernel_initializer,kernel_regularizer=kernel_regularizer,
                        layer_idx=idx+1,conv_first=True)
    
        
    if len(downsample)>0:
        for idx in range(len(downsample)):
            x = downsample[idx](x)
        residual = x
    
    residual = keras.layers.add([residual, out_result])
    residual = LeakyReLU(alpha=leaky_ratio,name='%s_relu%d_4'%(block_pre,block_number))(residual)
    
    return residual


blocks_dict = {
    'BASIC': BasicBlock,
    'BOTTLENECK': Bottleneck
}

expansion_dict = {
    'BASIC': 1,
    'BOTTLENECK': 4
}


class HighResolutionModule():
    def __init__(self,**kwargs):
        self.num_inchannels = kwargs.get('num_inchannels')
        self.fuse_method = kwargs.get('fuse_method')
        self.num_branches = kwargs.get('num_branches')
        self.multi_scale_output = kwargs.get('multi_scale_output')
        self.num_blocks = kwargs.get('num_blocks')
        self.block = kwargs.get('block')
        self.stage_num = kwargs.get('stage_num')
        self.kernel_size = kwargs.get('kernel_size',3)
        self.leaky_ratio = kwargs.get('leaky_ratio',0.01)
        self.norm_func = kwargs.get('norm_func',BatchNormalization)
        self.num_channels = kwargs.get('num_channels')
        self.activation_func = LeakyReLU
        self.padding = 'same'
        self.kernel_initializer = 'he_normal'
        self.kernel_regularizer = None
        

        
        self.branches = []
    
    def _make_one_branch(self,input_tensor,branch_index,block,num_blocks,num_channels,stride=1):
        '''
        make a single branch()
        
        repeats block_func num_blocks[branch_index] times. Shape of output depends on num_channels[branch_index]
        
        :param input_tensor. tensor
        :param branch_index. int. index of current branch
        :param block. func. represents block func current tensor will go through
        :param num_blocks. list of int
        :param num_channels. list of int. represents num_channels of output tensor of current branch
        :param stride. int
        
        :output. tensor
        '''
        downsample = []
        prestr = 'HRmakestage_stage%02d'%self.stage_num
        if stride != 1 or self.num_inchannels[branch_index] != num_channels[branch_index]:
            '''
            Do downsample if unmatch exists in any dimension
            self.num_inchannles[branch_index] represents 
            num_channels[branch_index]
            '''            
            block_number = 1
            downsample.append(Conv3D(filters = num_channels[branch_index], kernel_size = 1, strides = stride, 
                               padding = 'same', name='%s_downsample_conv%d_1'%(prestr,branch_index), 
                               kernel_initializer='he_normal'))
            downsample.append(self.norm_func(name='%s_downsample_bn%d_1'%(prestr,block_number)))
                
        
        result = input_tensor
        
        result = block(result,block_pre = prestr+'_blockfunc_%d'%branch_index, block_number = 1,num_filters = num_channels[branch_index],
                      kernel_size = self.kernel_size,downsample=downsample,leaky_ratio = self.leaky_ratio, nurm_func = self.norm_func,
                      stride = stride)
        
        self.num_inchannels[branch_index] = num_channels[branch_index]
        
        for block_idx in range(1,num_blocks[branch_index]):
            result = block(result,block_pre = prestr+'_blockfunc_%d'%branch_index, block_number = block_idx+1,num_filters = num_channels[branch_index],
                      kernel_size = self.kernel_size,downsample=[],leaky_ratio = self.leaky_ratio, nurm_func = self.norm_func)
        
        return result
        
    def make_branches(self,input_tensors,num_branches,block,num_blocks,num_channels):
        '''
        Generate list of tensors. Each one represents output tensor of a single branch
        for each branch, call make_one_branch func
        shape of each tensor in the return list only have diff in channel dimension
        
        :param input_tensors. list of tensors

        :param block. func. represents block func current tensor will go through
        :param num_blocks. ilst of int
        :param num_channels. list of int. represents num_channels of output tensor of current branch
        :param stride. int
        
        :output. tensor
    
        '''
        self.branches = []
        for branch_index in range(num_branches):
            self.branches.append(self._make_one_branch(input_tensors[branch_index],branch_index,block,num_blocks,num_channels))
        
            
    
    
    def make_fuse_layers(self):
        '''
        Generator list of list of tensors
        shape of return tensor should have shape [self.branch_index_limit,self.num_branches]
        tensors in the same row should have same size(can be added directly)
        '''
        fuse_result = []
        ### todo : check this part
        if self.num_branches == 1:
            return fuse_result
        num_branches = self.num_branches
        num_inchannels = self.num_inchannels
        input_tensors = self.branches
        
        
        for output_branch_index in range(self.branch_index_limit):
            fuse_tensors = []
            for input_branch_index in range(num_branches):
                ##### for each input tensor
                block_pre = 'HRmakestage_fuse_stage%02d_%02d_%02d'%(self.stage_num,input_branch_index,output_branch_index)
                current_input_tensor = input_tensors[input_branch_index]
                
                if input_branch_index>output_branch_index:
 
                    '''
                    use 1*1 conv to have same output-channel
                    use upsample to have same space dimension size()
                    '''
                    current_tensor = ConvUnit(current_input_tensor,norm_func=self.norm_func,activation_func=self.activation_func,
                                              num_filters=num_inchannels[output_branch_index],
                                                kernel_size=1,atrous_rate=1,padding=self.padding,block_prefix=block_pre,
                                                kernel_initializer=self.kernel_initializer,kernel_regularizer=self.kernel_regularizer,
                                                layer_idx=1,conv_first=True)
                    current_tensor = UpSampling3D(size=2**(input_branch_index-output_branch_index),name='%s_upsample'%block_pre)(current_tensor)
                    fuse_tensors.append(current_tensor)
                elif output_branch_index == input_branch_index:
                    #no func 
                    current_tensor = current_input_tensor
                else:
                    '''
                    do downsample
                    '''
                    current_tensor = current_input_tensor
                    for branch_distance in range(output_branch_index-input_branch_index):
                        '''
                        only change channel of tensor at last branch
                        '''
                        if branch_distance == output_branch_index - input_branch_index - 1:
                            num_outchannels_conv3x3 = num_inchannels[output_branch_index]
                        else:
                            num_outchannels_conv3x3 = num_inchannels[input_branch_index]
                        
                        current_tensor = ConvUnit(current_tensor,norm_func=self.norm_func,activation_func=self.activation_func,
                                              num_filters=num_outchannels_conv3x3,conv_stride=2,
                                                kernel_size=self.kernel_size,atrous_rate=1,padding=self.padding,block_prefix=block_pre,
                                                kernel_initializer=self.kernel_initializer,kernel_regularizer=self.kernel_regularizer,
                                                layer_idx=branch_distance,conv_first=True)

                fuse_tensors.append(current_tensor)
            fuse_result.append(fuse_tensors)
        
        return fuse_result
    #
    def get_module(self,input_tensors):

        '''
        Generator self.branch_index_limit tensors with input_tensors
        '''
        '''
        self.make_branches
        
        input0 -- convBlock*N0 --- output0
        input1 -- convBlock*N1 --- output1
        input2 -- convBlock*N2 --- output2
        add conv blocks to each input in input_tensors in parallel
        
        '''
        self.make_branches(input_tensors,self.num_branches,self.block,self.num_blocks,self.num_channels)
        if self.num_branches == 1:
            return self.branches[0]
        ### self.branch_index_limit defines how many tensors should be feed to next stage
        self.branch_index_limit = self.num_branches
#         if not self.multi_scale_output:
#             self.branch_index_limit = 1
        fuse_result = self.make_fuse_layers()
            
        output_tensors = []
        for output_branch_index in range(self.branch_index_limit):
            if output_branch_index == 0:
                current_tensor = self.branches[0]
            else:
                current_tensor = fuse_result[output_branch_index][0]
            for input_branch_index in range(1,self.num_branches):
                current_tensor = keras.layers.add([current_tensor, fuse_result[output_branch_index][input_branch_index]])
            output_tensors.append(current_tensor)
        return output_tensors,self.num_inchannels
            

class HRSegNet():
    def __init__(self,**kwargs):
        '''
        init parameters and layers
        '''
        self.config = kwargs
        self.inchannels = kwargs.get('inchannels',64)
        self.norm_func = kwargs.get('norm_func')
        self.leaky_ratio = kwargs.get('leaky_ratio',0.15)
        self.kernel_size = kwargs.get('kernel_size',3)
        
        self.activation_func = kwargs.get('activation_func')
        self.padding = kwargs.get('padding','same')
        self.kernel_initializer = kwargs.get('kernel_initializer','he_normal')
        self.kernel_regularizer = kwargs.get('kernel_regularizer',None)
        
        self.num_classes = kwargs.get('num_classes')
        self.classification_layers = kwargs.get('classification_layers')
        self.dropout_rate = kwargs.get('dropout_rate',0)
        self.initial_strides = kwargs.get('initial_strides')
        
        
       
            
    def make_transition_layer(self,input_tensors,block,num_channels_pre_layer,num_channels_cur_layer,stage_num):
        '''
        generate num_branches_cur output_tensors 
        :param input_tensor: list of tensors with shape equals to num_branches in current stage,length = 
        :param num_channels_pre_layer: list of int.num_channels of last stage tensors
        :param num_channels_cur_layer: list of int.num_channels of current stage tensors
        :param stage: int. denotes No. of current stage
        '''
        num_branches_cur = len(num_channels_cur_layer)
        num_branches_pre = len(num_channels_pre_layer)
        
        transition_tensors = []
        prestr = 'HRSegModel_stage_%d_transition'%stage_num            
        for branch_idx_cur in range(num_branches_cur):
            if branch_idx_cur < num_branches_pre:
                ## if current branch does exist in pre stage. No downsample 
                if num_channels_cur_layer[branch_idx_cur] != num_channels_pre_layer[branch_idx_cur]:
                    current_tensor = ConvUnit(input_tensors[branch_idx_cur],norm_func=self.norm_func,activation_func=self.activation_func,
                                              num_filters=num_channels_cur_layer[branch_idx_cur],
                                                kernel_size=self.kernel_size,atrous_rate=1,padding=self.padding,block_prefix=prestr,
                                                kernel_initializer=self.kernel_initializer,kernel_regularizer=self.kernel_regularizer,
                                                layer_idx=branch_idx_cur,conv_first=True)
                    
                    transition_tensors.append(current_tensor)
                else:
                    #### keep current tensor
                    transition_tensors.append(input_tensors[branch_idx_cur])
            else:
                ### current branch does not exist in pre stage(downsample the smallest tensor in pre stage)
                current_tensor = input_tensors[-1]
                for branch_distance in range(branch_idx_cur+1-num_branches_pre):
                    ### Do not expand channel dimension until last layer of current branch
                    inchannels = num_channels_pre_layer[-1]
                    outchannles = num_channels_cur_layer[branch_idx_cur] if branch_distance == branch_idx_cur - num_branches_pre else inchannels
                    current_tensor = ConvUnit(current_tensor,norm_func=self.norm_func,activation_func=self.activation_func,
                                              num_filters=num_channels_cur_layer[branch_idx_cur],conv_stride=2,
                                                kernel_size=self.kernel_size,atrous_rate=1,padding=self.padding,block_prefix=prestr,
                                                kernel_initializer=self.kernel_initializer,kernel_regularizer=self.kernel_regularizer,
                                                layer_idx=branch_idx_cur,conv_first=True)
                transition_tensors.append(current_tensor)
        return transition_tensors
    
    def make_layer(self,x,num_channels,block,blocks,stride=1,expansion=4):
        '''
        defines function to make initial convStage
        '''
        prestr = 'HRSegModel_pre'
        downsample = []
        ## if any dimension does not match. Add conv(with stride)+norm layer
        if stride != 1 or self.inchannels != num_channels * expansion:
            '''
            downsample input tensor to match dimension size or channel size
            '''
            block_number,branch_index = 1,1
            ### Todo change downsampel to a single block(why not conv-bn-relu)
            downsample.append(Conv3D(filters = num_channels * expansion, kernel_size = 1, strides = stride, 
                               padding = 'same', name='%s_downsample_conv%d_3'%(prestr,branch_index), 
                               kernel_initializer='he_normal'))
            downsample.append(self.norm_func(name='%s_downsample_bn%d_3'%(prestr,block_number)))
            
        result = x
        
        downsample_func = [downsample] + [[] for _ in range(1,blocks)]
        
        for block_idx in range(blocks):
            result = block(result,block_pre = prestr+'_blockfunc', block_number = block_idx + 1,num_filters = num_channels,
                      kernel_size = self.kernel_size,downsample=downsample_func[block_idx],leaky_ratio = self.leaky_ratio, nurm_func = self.norm_func,
                      stride = stride)
        return result
    
    def make_stage(self,input_tensors,layer_config,num_inchannels,multi_scale_output=True,stage_num=1,merge_choice=False):
        '''
        defines make stage func
        :param 
        :param input_tensors. List of tensors
        :param layer_config. Dict. Represents config of current stage
        :param num_inchannels.
        :param multi_scale_output.Bool.
        
        output should be two variable
        output_tensor: single tensor: represents output tensor of current stage
        num_inchannels: int. represents
        '''
        num_modules = layer_config['NUM_MODULES']
        num_branches = layer_config['NUM_BRANCHES']
        num_blocks = layer_config['NUM_BLOCKS']
        num_channels = layer_config['NUM_CHANNELS']
        block = blocks_dict[layer_config['BLOCK']]
        fuse_method = layer_config['FUSE_METHOD']
        
        
        result = []
        for module_idx in range(num_modules):
            ##### last module of current not-multi-scale stage
            if not multi_scale_output and module_idx == num_modules - 1:
                reset_multi_scale_output = False
            else:
                reset_multi_scale_output = True
                
            current_parameter_dict = {
                'num_inchannels':num_inchannels,'fuse_method':fuse_method,
                'num_branches':num_branches,'multi_scale_output':multi_scale_output,
                'num_blocks':num_blocks,'block':block,
                'stage_num':stage_num,'kernel_size':self.kernel_size,
                'leaky_ratio':self.leaky_ratio,'norm_func':self.norm_func,
                'num_channels':num_channels,'activation_func':self.activation_func,
                'padding':self.padding,'kernel_initializer':self.kernel_initializer,
                'kernel_regularizer':self.kernel_regularizer        
                
            }
            hr_obj = HighResolutionModule(**current_parameter_dict)
                
            current_tensor_list,return_num_inchannels = hr_obj.get_module(input_tensors)
            result.append(current_tensor_list)
        
        return result,return_num_inchannels
        
    def generate_layer_config(self,stage_idx):
        layer_config = {'NUM_MODULES':self.config['num_modules'],
                        'NUM_BRANCHES':self.config['num_branches'][stage_idx],
                        'NUM_BLOCKS':[self.config['block_per_stage'][stage_idx] for _ in range(self.config['num_branches'][stage_idx])],
                        'NUM_CHANNELS':[self.inchannels*(2**layer_idx) for layer_idx in range(self.config['num_branches'][stage_idx])],
                        'BLOCK':self.config['block_func'],
                        'FUSE_METHOD':self.config['FUSE_METHOD']}
        
        return layer_config
    
    def get_model(self,input_shape,stages=4,downsample_stride=1,choice='tranpose',merge_choice=False):
        main_input = Input(shape=input_shape, name='main_input')
        
        prestr = 'HRSegModel_pre'
        num_filters = self.inchannels
        x = main_input
        for layer_idx in range(1,3):
            x = ConvUnit(x,norm_func=self.norm_func,activation_func=self.activation_func,
                                      num_filters=num_filters,conv_stride=self.initial_strides[layer_idx-1],
                                        kernel_size=self.kernel_size,atrous_rate=1,padding=self.padding,block_prefix=prestr,
                                        kernel_initializer=self.kernel_initializer,kernel_regularizer=self.kernel_regularizer,
                                        layer_idx=layer_idx,conv_first=True)

#         print ('input shape is ',x.shape)
        x = self.make_layer(x,num_channels=num_filters,block=BasicBlock,blocks=2,expansion=1)
        
        '''
        Each stage consists of 2 parts
        make_transition_layer:
        make stage: Add conv blocks in parallel & fuse tensors to get output 
        
        '''
        
        
        
        start_idx = 2
        for stage_idx in range(start_idx,stages+1):
            
            layer_config = self.generate_layer_config(stage_idx-2)
#             print ('stage_idx',stage_idx)
#             print layer_config
            if stage_idx == start_idx:
                y_list,pre_stage_channels = [[x for _ in range(layer_config['NUM_BRANCHES'])]],[self.inchannels]
            input_tensors = y_list[-1]
            
            multi_scale_output = False if stage_idx==stages else True
            num_channels,expansion_rate  = layer_config['NUM_CHANNELS'],expansion_dict[layer_config['BLOCK']]
            x_list = self.make_transition_layer(input_tensors,block = blocks_dict[layer_config['BLOCK']],
                                                num_channels_pre_layer=pre_stage_channels,
                                                num_channels_cur_layer=[num_channels[i] * expansion_rate for i in range(len(num_channels))],
                                                stage_num=stage_idx-1)
        
            y_list,pre_stage_channels = self.make_stage(x_list,layer_config,num_channels,multi_scale_output,stage_num = stage_idx-1)



        '''
        Concatenate fature map from differnt level together
        '''
        result_tensors = []
        prestr = 'HRSegNet_Stage%02d_FinalFeatureFuse'%stages
        for idx in range(len(y_list[0])):
            upsample_rate = 2**idx
            result_tensor = ConvUnit(y_list[0][idx],norm_func=self.norm_func,activation_func=self.activation_func,block_prefix=prestr,
                                      num_filters=self.inchannels,kernel_size=self.kernel_size,atrous_rate=1,padding=self.padding,
                                        kernel_initializer=self.kernel_initializer,kernel_regularizer=self.kernel_regularizer,
                                        layer_idx=idx+1,conv_first=True) 
            result_tensor = UpSampling3D(size=2**idx,name='%s_upsample%02d'%(prestr,idx+1))(result_tensor)
            result_tensors.append(result_tensor)   
        result_tensor = Add(name='%s_add'%prestr)(result_tensors)
        
        '''
        Classification Part
        '''
        final_feature = self.classification_layers(name='HRSegNet_classification')(result_tensor)
        if self.num_classes == 1:
            final_acti = 'sigmoid'
        else:
            final_acti = 'softmax'
        final_output = Dense(self.num_classes,activation=final_acti,name='HRSegNet_final_output')(final_feature)


        model = Model(main_input,final_output)

        return model,[main_input]+y_list[0],[]

