import math
from hls4ml.converters.onnx_to_hls import onnx_handler, get_onnx_attribute, get_onnx_input_name, compute_pads_1d, compute_pads_2d
from hls4ml.converters.utils import compute_padding_1d, compute_padding_2d

@onnx_handler('Conv')
def parse_conv_layer(reader, node, inputs_map, input_shapes, graph, config):

    layer = {}
    layer['name'] = node.name
    layer['data_format'] = 'channels_first' #ONNX's default is channel first
    layer['inputs'] = get_onnx_input_name(node, graph)
    reader.add_input(layer['name'], node.input)

    strides = get_onnx_attribute(node, 'strides')
    kernel_shape = get_onnx_attribute(node, 'kernel_shape')

    if len(input_shapes[0]) == 3: # Conv1D
        layer['class_name'] = 'Conv1D'

        layer['in_width']= input_shapes[0][2]
        layer['n_chan']=input_shapes[0][1]
        layer['filt_width']= kernel_shape[0]
        layer['n_filt']= reader.get_weights_data(layer['name'], 'kernel').shape[2]
        layer['stride_width'] = strides[0]
        pads = compute_pads_1d(node, layer)

        layer['pad_left'] = pads[0]
        layer['pad_right'] = pads[1]
        
        if all(x == 0 for x in pads): # No padding, i.e., 'VALID' padding
            layer['padding'] = 'valid'
        else:
            layer['padding'] = 'same'
        
        (layer['out_width'],_,_) = compute_padding_1d(layer['padding'],
                                                      layer['in_width'],
                                                      layer['stride_width'],
                                                      layer['filt_width'])

        output_shape = [input_shapes[0][0], layer['n_filt'], layer['out_width']]
        
    elif len(input_shapes[0]) == 4: # Conv2D
        
        layer['class_name'] = 'Conv2D'

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

        layer['filt_height']=kernel_shape[0]
        layer['filt_width']=kernel_shape[1]
        
        layer['n_filt']=next((x.type.tensor_type.shape.dim[1].dim_value for x in graph.value_info if x.name == node.output[0]), None)
        layer['stride_height'] = strides[0]
        layer['stride_width'] = strides[1]
        pads = compute_pads_2d(node, layer)

        layer['pad_top'] = pads[0]
        layer['pad_bottom'] = pads[2]
        layer['pad_left'] = pads[1]
        layer['pad_right'] = pads[3]

        if all(x == 0 for x in pads): # No padding, i.e., 'VALID' padding in Keras/Tensorflow
            layer['padding'] = 'valid'
        else: #Only 'valid' and 'same' padding are available in Keras
            layer['padding'] = 'same'
            
        (layer['out_height'], layer['out_width'],_,_,_,_) = compute_padding_2d(layer['padding'],
                                                                               layer['in_height'],
                                                                               layer['in_width'],
                                                                               layer['stride_height'],
                                                                               layer['stride_width'],
                                                                               layer['filt_height'],
                                                                               layer['filt_width'])

        output_shape = [input_shapes[0][0], layer['n_filt'], layer['out_height'], layer['out_width']]

    return layer, output_shape
    