import collections.abc

from hls4ml.converters.keras_to_hls import keras_handler, parse_default_keras_layer


@keras_handler('ZeroPadding1D')
def parse_zeropadding1d_layer(keras_layer, input_names, input_shapes, data_reader):
    assert keras_layer['class_name'] == 'ZeroPadding1D'

    layer = parse_default_keras_layer(keras_layer, input_names)

    padding = keras_layer['config']['padding']
    if isinstance(padding, int):
        layer['pad_left'] = padding
        layer['pad_right'] = padding
    elif isinstance(padding, collections.abc.Sequence):
        layer['pad_left'] = padding[0]
        layer['pad_right'] = padding[1]

    if layer['data_format'] == 'channels_first':
        output_shape = [
            input_shapes[0][0],  # Batch
            input_shapes[0][1],  # Channels
            layer['pad_left'] + input_shapes[0][2] + layer['pad_right'],  # Width
        ]
        layer['out_width'] = output_shape[2]
        layer['n_chan'] = output_shape[1]

        layer['in_width'] = input_shapes[0][2]
    else:
        output_shape = [
            input_shapes[0][0],  # Batch
            layer['pad_left'] + input_shapes[0][1] + layer['pad_right'],  # Width
            input_shapes[0][2],  # Channels
        ]
        layer['out_width'] = output_shape[1]
        layer['n_chan'] = output_shape[2]

        layer['in_width'] = input_shapes[0][1]

    return layer, output_shape


@keras_handler('ZeroPadding2D')
def parse_zeropadding2d_layer(keras_layer, input_names, input_shapes, data_reader):
    assert keras_layer['class_name'] == 'ZeroPadding2D'

    layer = parse_default_keras_layer(keras_layer, input_names)

    padding = keras_layer['config']['padding']
    if isinstance(padding, int):
        layer['pad_top'] = padding
        layer['pad_bottom'] = padding
        layer['pad_left'] = padding
        layer['pad_right'] = padding
    elif isinstance(padding, collections.abc.Sequence):
        height_pad, width_pad = padding
        if isinstance(height_pad, collections.abc.Sequence):
            layer['pad_top'] = height_pad[0]
            layer['pad_bottom'] = height_pad[1]
        else:
            layer['pad_top'] = height_pad
            layer['pad_bottom'] = height_pad
        if isinstance(width_pad, collections.abc.Sequence):
            layer['pad_left'] = width_pad[0]
            layer['pad_right'] = width_pad[1]
        else:
            layer['pad_left'] = width_pad
            layer['pad_bottom'] = width_pad

    if layer['data_format'] == 'channels_first':
        output_shape = [
            input_shapes[0][0],  # Batch
            input_shapes[0][1],  # Channels
            layer['pad_top'] + input_shapes[0][2] + layer['pad_bottom'],  # Height
            layer['pad_left'] + input_shapes[0][3] + layer['pad_right'],  # Width
        ]
        layer['out_height'] = output_shape[2]
        layer['out_width'] = output_shape[3]
        layer['n_chan'] = output_shape[1]

        layer['in_height'] = input_shapes[0][2]
        layer['in_width'] = input_shapes[0][3]
    else:
        output_shape = [
            input_shapes[0][0],  # Batch
            layer['pad_top'] + input_shapes[0][1] + layer['pad_bottom'],  # Height
            layer['pad_left'] + input_shapes[0][2] + layer['pad_right'],  # Width
            input_shapes[0][3],  # Channels
        ]
        layer['out_height'] = output_shape[1]
        layer['out_width'] = output_shape[2]
        layer['n_chan'] = output_shape[3]

        layer['in_height'] = input_shapes[0][1]
        layer['in_width'] = input_shapes[0][2]

    return layer, output_shape
