// Filename : sobel.cpp

#include "sobel.h"
#include "spacecomp_pre_cpp.h"
#include "space_chrono.h"

// Do NOT move macros, functions and includes beginning with "spacecomp".

sobel::sobel(SPACECOMP_CTOR_PARAMS(sobel, INDEX))
	: abstract_module(SPACECOMP_MBASE_ARGS(sobel, INDEX))
	, SPACECOMP_START_ILIST(sobel, INDEX)
	// Other initializers may go here
{
	SPACECOMP_THREAD(sobel, INDEX, thread);
	request_stack(2*8*1080*1920);
}

void sobel::thread(SPACECOMP_THREAD_PARAMS(sobel, INDEX, thread)) {
	spacecomp_thread_initialize();
	uint32_t input_addr_offset = 0;
	uint32_t output_addr_offset = 0;

	const uint32_t img_width = IMG_WIDTH;
	const uint32_t img_height = IMG_HEIGHT;
	const uint32_t slice_width = SLICE_WIDTH;

	// Sliding buffer holding the three lines needed to compute the current line of the output, each line is only read once from DDR
	uint8_t input_buffer[3][slice_width+8];

	// Sliding window holding the data needed to compute one pixel, each pixel is only read once from the input_buffer
	uint8_t convolution_window[3][3];
	uint8_t output_buffer[slice_width];

	// Set input buffer as three separate true dual port BRAMs, allows to read one element from each line every cycle
	#pragma HLS RESOURCE         variable=input_buffer core=RAM_T2P_BRAM 
	#pragma HLS ARRAY_PARTITION  variable=input_buffer complete dim=1

	// Partition the convolution window into individual registers, allows all elements to be available simultaneously
	#pragma HLS ARRAY_PARTITION variable=convolution_window complete dim=0

	while(1) {
		const bool initializing = spacecomp_thread_loop_start();

		// Wait for the controller notification
		ModuleRead(CONTROLLER0_ID, SPACE_BLOCKING);

		RegisterRead(REGISTER_FILE0_ID, PIXEL_DATA_IN_OFFSET_REG, input_addr_offset);
		RegisterRead(REGISTER_FILE0_ID, PIXEL_DATA_OUT_OFFSET_REG, output_addr_offset);

		// Set the address offsets to the first pixel of the slice
		output_addr_offset += INDEX * slice_width ;
		input_addr_offset += INDEX * slice_width ;

		// Set upper and lower border to 0 (sobel operator does not apply to image borders)
		for (unsigned int i = 0; i < slice_width; ++i){
			#pragma HLS UNROLL
			output_buffer[i] = 0;
		}

		DeviceWrite(ZYNQ_DDR0_ID, output_addr_offset, output_buffer, slice_width);
		DeviceWrite(ZYNQ_DDR0_ID, output_addr_offset + (img_height - 1)*img_width, output_buffer, slice_width);

		output_addr_offset += img_width;

		// Read the first three lines. One extra pixel is needed on each side to compute borders,
		// read 4 more bytes on each side to keep a 4-bytes aligned address and 4-bytes multiple read size.
		// First and last slices are handled with READ_OFFSET and READ_SIZE to avoid reading outside pixel data.
		for (unsigned int i = 0; i < 3; ++i) {
			DeviceRead(ZYNQ_DDR0_ID, input_addr_offset+READ_OFFSET-4, input_buffer[i], READ_SIZE);
			input_addr_offset += img_width;
		}

		// Perform Sobel line per line
		for (unsigned int line = 1; line < img_height-1; ++line) {

			// Instantiate the window for the first pixel of the line
			for (int i = 0; i < 3; i++) {
				#pragma HLS UNROLL

				#if IMG_BORDER_LEFT
				convolution_window[i][1] = input_buffer[i][0];
				convolution_window[i][2] = input_buffer[i][1];
                #else
				convolution_window[i][0] = input_buffer[i][3];
				convolution_window[i][1] = input_buffer[i][4];
				convolution_window[i][2] = input_buffer[i][5];
                #endif
			}

			// Compute the operator for each pixel of the line
			for (unsigned int pixel = 0; pixel < slice_width; ++pixel) {
				#pragma HLS PIPELINE II=1

				// Set output image borders to 0 based on slice position
				bool is_border_pixel = false;
				#if IMG_BORDER_LEFT
				is_border_pixel = (pixel == 0);
				#endif
				#if IMG_BORDER_RIGHT
				is_border_pixel = (pixel == slice_width - 1);
				#endif

				if (is_border_pixel) {
					output_buffer[pixel] = 0;
				} else {
					output_buffer[pixel] = sobel_operator(convolution_window, (uint32_t) line);
				}

				// Shift the window (3 parallel accesses to the 3 BRAMs, allows II = 1)
				if (pixel != slice_width-IMG_BORDER_RIGHT-1) {
					for (int i = 0; i < 3; i++) {
						#pragma HLS UNROLL
						
						convolution_window[i][0] = convolution_window[i][1];
						convolution_window[i][1] = convolution_window[i][2];
						convolution_window[i][2] = input_buffer[i][pixel-READ_OFFSET+6];
					}
				}
			}

			// Write the output line
			DeviceWrite(ZYNQ_DDR0_ID, output_addr_offset, output_buffer, slice_width);
			output_addr_offset += img_width;

			// Read the next line
			if (line != img_height - 2) {
				DeviceRead(ZYNQ_DDR0_ID, input_addr_offset+READ_OFFSET-4, input_buffer[(line+2) % 3], READ_SIZE);
				input_addr_offset += img_width;
			}
		}

		// Notify controller we are done
		ModuleWrite(CONTROLLER0_ID, SPACE_BLOCKING);
	}
}

uint8_t sobel::sobel_operator(uint8_t convolution_window[3][3], uint32_t line) {
	
	// Sobel operator kernels
	static const signed char x_op[3][3] =
	{ { -1,0,1 },
	{ -2,0,2 },
	{ -1,0,1 } };

	static const signed char y_op[3][3] =
	{ { 1,2,1 },
	{ 0,0,0 },
	{ -1,-2,-1 } };

	int x_weight = 0;
	int y_weight = 0;
	unsigned int edge_weight;

	for (unsigned int i = 0; i < 3; ++i) {
		#pragma HLS UNROLL

		for (unsigned int j = 0; j < 3; ++j) {
			#pragma HLS UNROLL
			
			// X direction gradient
			x_weight += convolution_window[(line+i-1) % 3][j] * x_op[i][j];

			// Y direction gradient
			y_weight += convolution_window[(line+i-1) % 3][j] * y_op[i][j];
		}
	}

	// Sobel gradient magnitude is the pythagorean sum of X and Y gradients
	// Approximate by sum of absolutes to avoid square root
	if (x_weight < 0)  x_weight = x_weight * -1;
	if (y_weight < 0)  y_weight = y_weight * -1;
	edge_weight = x_weight + y_weight;

	return (uint8_t)(edge_weight);
}

#include "spacecomp_post_cpp.h"