#ifndef NNET_CONV1D_H_
#define NNET_CONV1D_H_

#include "nnet_common.h"
#include "nnet_conv1d_latency.h"
#include "nnet_conv1d_resource.h"
#include "nnet_function_stubs.h"
#include <cstdlib>

namespace nnet {

struct conv1d_config {
    // Internal data type definitions
    typedef float bias_t;
    typedef float weight_t;
    typedef float accum_t;

    // Convolutional parameters
    static const unsigned pad_left = 0;
    static const unsigned pad_right = 0;
    static const unsigned in_width = 10;
    static const unsigned n_chan = 0;
    static const unsigned filt_width = 1;
    static const unsigned kernel_size = filt_width;
    static const unsigned n_filt = 1;
    static const unsigned stride_width = 1;
    static const unsigned dilation = 1;
    static const unsigned out_width = 10; //(N_IN + PAD_LEFT * PAD_RIGHT - (DILATION * (FILT_WIDTH - 1) + 1)) / STRIDE + 1

    static const unsigned reuse_factor = 1;
    static const bool store_weights_in_bram = false;
    static const unsigned n_zeros = 0; // not used yet
};

template <class data_T, class res_T, typename CONFIG_T>
void conv_1d_cl(data_T data[CONFIG_T::in_width * CONFIG_T::n_chan], res_T res[CONFIG_T::out_width * CONFIG_T::n_filt],
                typename CONFIG_T::weight_t weights[CONFIG_T::filt_width * CONFIG_T::n_chan * CONFIG_T::n_filt],
                typename CONFIG_T::bias_t biases[CONFIG_T::n_filt]) {
    // Inlining helps reduce latency, but may also cause timing issues in some cases, use carefully.
    // But without inlining Vitis HLS doesn't respect the parallelization factor config ¯\_(ツ)_/
    // Vitis2025.1 hangs in RTL simulation with this, though

    #pragma HLS INLINE recursive

    // #pragma HLS PIPELINE II = CONFIG_T::reuse_factor * CONFIG_T::n_partitions
    // ↑ This makes II=2 in for all n_partitions > 1, no matter what the actual II should be

    CONFIG_T::template conv_kernel<data_T, res_T, CONFIG_T>::conv(data, res, weights, biases);
}

template <class data_T, class res_T, typename CONFIG_T>
void pointwise_conv_1d_cl(data_T data[CONFIG_T::in_width * CONFIG_T::n_chan],
                          res_T res[CONFIG_T::out_width * CONFIG_T::n_filt],
                          typename CONFIG_T::weight_t weights[CONFIG_T::n_chan * CONFIG_T::n_filt],
                          typename CONFIG_T::bias_t biases[CONFIG_T::n_filt]) {
    assert(CONFIG_T::filt_width == 1);

    // Inlining helps reduce latency, but may also cause timing issues in some cases, use carefully.
    // But without inlining Vitis HLS doesn't respect the parallelization factor config ¯\_(ツ)_/¯

    #pragma HLS INLINE recursive

    // #pragma HLS PIPELINE II = CONFIG_T::reuse_factor * CONFIG_T::n_partitions
    // ↑ This makes II=2 in for all n_partitions > 1, no matter what the actual II should be

    CONFIG_T::template conv_kernel<data_T, res_T, CONFIG_T>::conv(data, res, weights, biases);
}

template <class data_T, class res_T, typename CONFIG_T>
class Conv1DLatency : public nnet::Conv1DKernel<data_T, res_T, CONFIG_T> {
  public:
    static void conv(data_T data[CONFIG_T::in_width * CONFIG_T::n_chan], res_T res[CONFIG_T::out_width * CONFIG_T::n_filt],
                     typename CONFIG_T::weight_t weights[CONFIG_T::filt_width * CONFIG_T::n_chan * CONFIG_T::n_filt],
                     typename CONFIG_T::bias_t biases[CONFIG_T::n_filt]) {
        // #pragma HLS INLINE recursive
        conv_1d_latency_cl<data_T, res_T, CONFIG_T>(data, res, weights, biases);
    }
};

template <class data_T, class res_T, typename CONFIG_T>
class Conv1DResource : public nnet::Conv1DKernel<data_T, res_T, CONFIG_T> {
  public:
    static void conv(data_T data[CONFIG_T::in_width * CONFIG_T::n_chan], res_T res[CONFIG_T::out_width * CONFIG_T::n_filt],
                     typename CONFIG_T::weight_t weights[CONFIG_T::filt_width * CONFIG_T::n_chan * CONFIG_T::n_filt],
                     typename CONFIG_T::bias_t biases[CONFIG_T::n_filt]) {
        // #pragma HLS INLINE recursive
        conv_1d_resource_cl<data_T, res_T, CONFIG_T>(data, res, weights, biases);
    }
};

template <class data_T, class res_T, typename CONFIG_T>
class BatchedDenseForConv1D : public nnet::Conv1DKernel<data_T, res_T, CONFIG_T> {
  public:
    static void conv(data_T data[CONFIG_T::in_width * CONFIG_T::n_chan], res_T res[CONFIG_T::out_width * CONFIG_T::n_filt],
                     typename CONFIG_T::weight_t weights[CONFIG_T::n_chan * CONFIG_T::n_filt],
                     typename CONFIG_T::bias_t biases[CONFIG_T::n_filt]) {

        #pragma HLS PIPELINE II = CONFIG_T::reuse_factor * CONFIG_T::n_partitions
        #pragma HLS INLINE RECURSIVE
        data_T data_tmp[CONFIG_T::n_partitions][CONFIG_T::in_width * CONFIG_T::n_chan / CONFIG_T::n_partitions];
        #pragma HLS ARRAY_PARTITION variable=data_tmp complete dim=0
        res_T res_tmp[CONFIG_T::n_partitions][CONFIG_T::out_width * CONFIG_T::n_filt / CONFIG_T::n_partitions];
        #pragma HLS ARRAY_PARTITION variable=res_tmp complete dim=0

        for (int jj = 0; jj < CONFIG_T::n_partitions; jj++) {
            #pragma HLS UNROLL
            for (int ii = 0; ii < CONFIG_T::in_width * CONFIG_T::n_chan / CONFIG_T::n_partitions; ii++) {
                #pragma HLS UNROLL
                data_tmp[jj][ii] = data[jj * CONFIG_T::in_width * CONFIG_T::n_chan / CONFIG_T::n_partitions + ii];
            }
        }

        // #pragma HLS ALLOCATION function instances=nnet::pointwise_conv_1d_latency_cl<data_T, res_T, CONFIG_T> limit=1
        // Vitis 2025.1 crashes with this, but 2023.2 is fine. Not tested on other versions.

        for (int jj = 0; jj < CONFIG_T::n_partitions; jj++) {
            nnet::pointwise_conv_1d_latency_cl<data_T, res_T, CONFIG_T>(data_tmp[jj], res_tmp[jj], weights, biases);
        }

        for (int jj = 0; jj < CONFIG_T::n_partitions; jj++) {
            #pragma HLS UNROLL
            for (int ii = 0; ii < CONFIG_T::out_width * CONFIG_T::n_filt / CONFIG_T::n_partitions; ii++) {
                #pragma HLS UNROLL
                res[jj * CONFIG_T::out_width * CONFIG_T::n_filt / CONFIG_T::n_partitions + ii] = res_tmp[jj][ii];
            }
        }
    }
};

} // namespace nnet

#endif
