from keras.utils import conv_utils
from keras import backend as K
from keras.engine import InputSpec
from keras.layers import Conv2D, ZeroPadding2D, DepthwiseConv2D, SeparableConv2D
import tensorflow as tf

'''
conv, mask = PConv2D(filters, kernel_size, strides=2, padding='same')([img_in, mask_in])
     model =  Conv2D(filters, kernel_size, strides=1, padding="same")(model)
'''

class aPConv2D(Conv2D):
    def __init__(self, *args, n_channels=3, mono=False, **kwargs):
        super().__init__(*args, **kwargs)
        self.input_spec = InputSpec(ndim=4)

    def build(self, input_shape):        
        """Adapted from original _Conv() layer of Keras        
        param input_shape: list of dimensions for [img, mask]
        """
        
        if self.data_format == 'channels_first':
            channel_axis = 1
        else:
            channel_axis = -1
            
        if input_shape[channel_axis] is None:
            raise ValueError('The channel dimension of the inputs should be defined. Found `None`.')
            
        self.input_dim = input_shape[channel_axis]
        
        # Image kernel
        kernel_shape = self.kernel_size + (self.input_dim, self.filters)
        # print(kernel_shape)
        self.kernel = self.add_weight(shape=kernel_shape,
                                      initializer=self.kernel_initializer,
                                      name='img_kernel',
                                      regularizer=self.kernel_regularizer,
                                      constraint=self.kernel_constraint)
        self.kernel_1x1xd = K.ones(shape=(1,1, self.input_dim,1, 1))
        self.pconv_padding = (
            (int((self.kernel_size[0]-1)/2), int((self.kernel_size[0]-1)/2)), 
            (int((self.kernel_size[0]-1)/2), int((self.kernel_size[0]-1)/2)), 
        )
 
        # Window size - used for normalization
        self.window_size = self.kernel_size[0] * self.kernel_size[1]
        
        if self.use_bias:
            self.bias = self.add_weight(shape=(self.filters,),
                                        initializer=self.bias_initializer,
                                        name='bias',
                                        regularizer=self.bias_regularizer,
                                        constraint=self.bias_constraint)
        else:
            self.bias = None
        self.built = True

    def call(self, inputs, mask=None):
        mask1p0 = K.ones(shape=(1,inputs.shape[1], inputs.shape[2],inputs.shape[3]))
        images = K.spatial_2d_padding(inputs, self.pconv_padding, self.data_format) # 24x24
        mask_1p0 = K.spatial_2d_padding(mask1p0, self.pconv_padding, self.data_format)  # 24x24
        mask_1p1 = K.ones_like(mask_1p0)
        mask1p0_output = K.depthwise_conv2d(
            mask_1p0, self.kernel, 
            strides=self.strides,
            padding='valid',
            data_format=self.data_format,
            dilation_rate=self.dilation_rate
        )
        print('mask1p0_output', mask1p0_output.shape)
        asd = mask1p0_output.eval(session=K.get_session())
        print(asd[0,:,:,0])
        # Apply convolutions to image
        img_output = K.depthwise_conv2d(
            images, self.kernel, 
            strides=self.strides,
            padding='valid',
            data_format=self.data_format,
            dilation_rate=self.dilation_rate
        )   

        mask1p1_output = K.depthwise_conv2d(
            # mask_1p1, self.kernel_mask1p1, 
            mask_1p1, self.kernel, 
            strides=self.strides,
            padding='valid',
            data_format=self.data_format,
            dilation_rate=self.dilation_rate
        )
        mask1p1_output = mask1p1_output.eval(session=K.get_session())

        mask_ratio = mask1p1_output / (mask1p0_output + 1e-8)

        # Normalize iamge output
        img_output = img_output * mask_ratio

        img_output = K.expand_dims(img_output, axis=-1)
        print('img_output', img_output.shape)
        stride3d = self.strides+(self.input_dim,)
        # print(stride3d)
        img_output = K.conv3d(img_output, 
            self.kernel_1x1xd, 
            strides=stride3d,
            padding='valid',
            data_format=self.data_format#,
        )   
        print('img_output2', img_output.shape)
        img_output = K.squeeze(img_output, axis=-1)
        # Apply bias only to the image (if chosen to do so)
        if self.use_bias:
            img_output = K.bias_add(
                img_output,
                self.bias,
                data_format=self.data_format)
        print('img_output bias', img_output.shape)
        # Apply activations on the image
        if self.activation is not None:
            img_output = self.activation(img_output)
           
        # return [img_output, mask_output]
        
        print('img_output final', img_output.shape)
        return img_output
    
    def compute_output_shape(self, input_shape):
        if self.data_format == 'channels_last':
            space = input_shape[1:-1]
            new_space = []
            for i in range(len(space)):
                new_dim = conv_utils.conv_output_length(
                    space[i],
                    self.kernel_size[i],
                    padding='same',
                    stride=self.strides[i],
                    dilation=self.dilation_rate[i])
                new_space.append(new_dim)
            new_shape = (input_shape[0],) + tuple(new_space) + (self.filters,)
            return [new_shape, new_shape]
        if self.data_format == 'channels_first':
            space = input_shape[2:]
            new_space = []
            for i in range(len(space)):
                new_dim = conv_utils.conv_output_length(
                    space[i],
                    self.kernel_size[i],
                    padding='same',
                    stride=self.strides[i],
                    dilation=self.dilation_rate[i])
                new_space.append(new_dim)
            new_shape = (input_shape, self.filters) + tuple(new_space)
            return [new_shape, new_shape]
