#include "hls_stream.h"

#define MAX_WIDTH  1024
#define MAX_HEIGHT 1024
#define MAX_KERNEL 11

typedef unsigned char pixel_t;
typedef signed char coeff_t;
typedef int acc_t;
typedef hls::stream<pixel_t> pixel_stream;
typedef hls::stream<coeff_t> coeff_stream;
typedef hls::stream<acc_t> acc_stream;

void convolution(pixel_stream& in_stream, coeff_stream& kernel_stream, acc_stream& out_stream, int height, int width, int ksize) {
#pragma HLS INTERFACE axis port=in_stream
#pragma HLS INTERFACE axis port=kernel_stream
#pragma HLS INTERFACE axis port=out_stream
#pragma HLS INTERFACE s_axilite port=height
#pragma HLS INTERFACE s_axilite port=width
#pragma HLS INTERFACE s_axilite port=ksize
#pragma HLS INTERFACE ap_ctrl_none port=return

    pixel_t line_buffer[MAX_WIDTH + MAX_KERNEL - 1][MAX_KERNEL - 1];
    pixel_t window_buffer[MAX_KERNEL][MAX_KERNEL];
    coeff_t kernel_buffer[MAX_KERNEL][MAX_KERNEL];

    // Initialize line buffer
    for (int i = 0; i < ksize - 1; i++) {
        for (int j = 0; j < width + ksize - 1; j++) {
            line_buffer[j][i] = 0;
        }
    }

    for (int i = 0; i < height; i++) {
        for (int j = 0; j < width; j++) {
            // Shift line buffer
            for (int k = 0; k < ksize - 2; k++) {
                line_buffer[j][k] = line_buffer[j][k+1];
            }
            line_buffer[j][ksize-2] = in_stream.read();

            // Shift window buffer
            for (int k = 0; k < ksize - 1; k++) {
                for (int l = 0; l < ksize - 2; l++) {
                    window_buffer[k][l] = window_buffer[k][l+1];
                }
                window_buffer[k][ksize-2] = (j > 0) ? line_buffer[j-1][k] : 0;
                window_buffer[k][ksize-1] = line_buffer[j][k];
            }

            // Shift kernel buffer
            for (int k = 0; k < ksize - 1; k++) {
                for (int l = 0; l < ksize - 1; l++) {
                    kernel_buffer[k][l] = kernel_buffer[k][l+1];
                }
                kernel_buffer[k][ksize-2] = kernel_stream.read();
            }

            // Perform convolution
            acc_t acc = 0;
            for (int k = 0; k < ksize; k++) {
                for (int l = 0; l < ksize; l++) {
                    acc += window_buffer[k][l] * kernel_buffer[k][l];
                }
            }
            out_stream.write(acc);
        }

        // Shift line buffer after processing a row
        for (int j = 0; j < width + ksize - 2; j++) {
            line_buffer[j][ksize-2] = (i == height - 1) ? 0 : line_buffer[j][ksize-1];
        }
    }
}