from keras.layers import Input
from keras.models import Model
from qkeras import *
import hls4ml

def get_model():
    # Define a dummy model with only one QSeparableConv1D layer
    input_layer = Input(shape=(32, 3))
    layer = QSeparableConv1D(
                filters=16,
                kernel_size=3,
                depthwise_quantizer=quantized_bits(16, 6, alpha=1),
                pointwise_quantizer=quantized_bits(16, 6, alpha=1),
                bias_quantizer=quantized_bits(16, 6, alpha=1)
            )(input_layer)
    model = Model(inputs=input_layer, outputs=layer)
    return model

model = get_model()

model.summary()

config = hls4ml.utils.config_from_keras_model(model, granularity="name")

# Configure the project to be SeparableConv1D compatible
config['Model']['Precision'] = 'ap_fixed<16,6>'
config['Model']['ReuseFactor'] = 1
config['Model']['Strategy'] = 'Latency'

# Use the Vivado backend (2020.1)
cfg = hls4ml.converters.create_config(backend='Vitis')
cfg['IOType'] = 'io_stream'
cfg['HLSConfig'] = config
cfg['KerasModel'] = model
cfg['OutputDir'] = 'hls4ml_prj'
cfg['Part'] = 'xcku115-flvb2104-2-i'

hls_model = hls4ml.converters.keras_to_hls(cfg)

hls_model.compile()

hls_model.build(reset=True, csim=False, synth=True)