#include "calib.h"
#include "stdlib.h"
#include <cstdint>
#include <cmath>

void calib(
	hls::stream<CalibData> 		&inCalibData,
	hls::stream<InterpolateCalibSubpix> 	&outSubpixCalib
)
{
#if !(VITIS_HLS)
#pragma HLS DATA_PACK variable=inCalibData
#pragma HLS DATA_PACK variable=outSubpixCalib
#else
#pragma HLS aggregate variable=inCalibData compact=bit
#pragma HLS aggregate variable=outSubpixCalib compact=bit
#endif
#pragma HLS INTERFACE axis register both port=inCalibData
#pragma HLS INTERFACE axis register both port=outSubpixCalib

	static RgnSubPixCalib calibData_buf[ROW_COUNT] = {}; //
#pragma HLS BIND_STORAGE variable=calibData_buf type=ram_2p impl=bram
#pragma HLS RESET variable=calibData_buf off

	//HLS reset to "0"
	//CAM_MAX_Y = 4095
	static SRowData interpolate_buf[ROW_COUNT] = {}; //size = 4096 * 72 = 288kbit， RAM36 * 8
#pragma HLS BIND_STORAGE variable=interpolate_buf type=ram_2p impl=bram
#pragma HLS ARRAY_PARTITION variable=interpolate_buf cyclic factor=3 dim=1
#pragma HLS RESET variable=interpolate_buf off

	// 每行状态：
	// [41:2] = 40位内部帧标记
	// [1:0]  = 点数：0、1、2
	static ap_uint<42> rowState[BANK_COUNT][BANK_DEPTH] = {};
	#pragma HLS ARRAY_PARTITION variable=rowState complete dim=1
	#pragma HLS BIND_STORAGE variable=rowState type=ram_2p impl=bram
	#pragma HLS RESET variable=rowState off

	// 与 rowState 一起保留，运行时复位不清零。
	static ap_uint<40> frameGeneration = 0;
	#pragma HLS RESET variable=frameGeneration off

	//row data initialization
	SRowData nullRowData;
	nullRowData.ptNum = 0;
	nullRowData.pt1.x = 0.0f;
	nullRowData.pt1.y = 0;;  //support up to 4k camera
	nullRowData.pt1.rid = 0;
	nullRowData.pt1.value = 0;
	nullRowData.pt1.flag = 0;	//bit0:OverLap | bit1:Rflct
	nullRowData.pt1.matchFlag = 0;

	nullRowData.pt2.x = 0.0f;
	nullRowData.pt2.y = 0;;  //support up to 4k camera
	nullRowData.pt2.rid = 0;
	nullRowData.pt2.value = 0;
	nullRowData.pt2.flag = 0;	//bit0:OverLap | bit1:Rflct
	nullRowData.pt2.matchFlag = 0;

#if 0
	for(uint16_t i= 0; i < CAM_MAX_Y; i ++)
	{
#pragma HLS PIPELINE II=1
		interpolate_buf[i] = nullRowData;
	}
#endif

	CalibData 	InData;
	InData = inCalibData.read();
	if(0b10 != InData.Sync)  //wait for syncFirst
		return;

	uint32_t frmNo, timeStamp, encInfo, frmWH;
	uint16_t frmX, frmY, Lines;

	Byte8 dataBuff;
	//1st: frameNo
	dataBuff.dData = InData.calibK[0];
	frmNo = dataBuff.nData[0];
	//2nd: timeStamp
	InData = inCalibData.read();
	dataBuff.dData = InData.calibK[0];
	timeStamp = dataBuff.nData[0];
	//3rd:encInfo
	InData = inCalibData.read();
	dataBuff.dData = InData.calibK[0];
	encInfo = dataBuff.nData[0];
	//4th: frmW, frmH
	InData = inCalibData.read();
	dataBuff.dData = InData.calibK[0];
	frmWH = dataBuff.nData[0];
	//5th: frmX, WinNum
	InData = inCalibData.read();
	dataBuff.dData = InData.calibK[0];
	frmX = dataBuff.nData[0] & 0xffff;
	Lines = dataBuff.nData[0] >> 16;
	//6th: frmY
	InData = inCalibData.read();
	dataBuff.dData = InData.calibK[0];
	frmY = dataBuff.nData[0] & 0xffff;

	const ap_uint<40> max_generation = 0xffffffffffULL;

	if (frameGeneration == max_generation)
	{
GENERATION_WRAP_CLEAR:
	    for (unsigned addr = 0; addr < BANK_DEPTH; ++addr)
	    {
#pragma HLS PIPELINE II=1

CLEAR_BANKS:
	        for (unsigned b = 0; b < BANK_COUNT; ++b)
	        {
#pragma HLS UNROLL
	            rowState[b][addr] = 0;
	        }
	    }

	    // 全部清零完成后，再发布 generation 1。
	    frameGeneration = 1;
	}
	else
	{
	    frameGeneration = frameGeneration + 1;
	}

	const ap_uint<40> currFrameTag = frameGeneration;

	Byte8 dataConver_x, dataConver_y;
	//double kx, ky;
	float kx_1, kx_2, ky_1, ky_2;

	//line_pt0
	RgnSubPixCalib calibPnt;
	float rawx = 0.0f, x_square = 0.0f, rawy = 0.0f, y_square = 0.0f;
	float calib_x = 0.0f, calib_y = 0.0f;
	//line_pt1
	RgnSubPixCalib calibPnt_pt1;
	float rawx_pt1 = 0.0f, x_square_pt1 = 0.0f, rawy_pt1 = 0.0f, y_square_pt1 = 0.0f;
	float calib_x_pt1 = 0.0f, calib_y_pt1 = 0.0f;

	float s2_result_x=0.0f, s3_result_x=0.0f, s4_result_x=0.0f;
	float s2_result_y=0.0f, s3_result_y=0.0f, s4_result_y=0.0f;
	float s2_result_x_pt1=0.0f, s3_result_x_pt1=0.0f, s4_result_x_pt1=0.0f;
	float s2_result_y_pt1=0.0f, s3_result_y_pt1=0.0f, s4_result_y_pt1=0.0f;

	bool vldFlag = 0;
	bool vldFlag_pt1 = 0;
	uint16_t in_dataSize = 0;
	for(int n = 0; n < Lines; n ++)
	{
#pragma HLS LOOP_TRIPCOUNT min=3072 max=3072
		for(int j = 0; j <= 5; j ++) //6 clocks: 0-pt0, 1-pt1, 2，3，4，5 - K
		{
#pragma HLS PIPELINE II=1
			InData = inCalibData.read();
			if( 0 == j)//read subpix
			{
				calibPnt.x = (float)calib_x;
				calibPnt.y = (float)calib_y;
				if( (n > 0) && (vldFlag == 1))
				{
					calibData_buf[in_dataSize] = calibPnt;
					in_dataSize++;
					//outSubpixCalib.write(calibPnt);
				}

				dataBuff.dData = InData.calibK[0];
				rawx = (float)dataBuff.fData[0];
				ap_uint<32> apData32 = (ap_uint<32>)dataBuff.nData[1];
				ap_uint<CAM4K_Y_WIDTH> dy_12;
				dy_12(CAM4K_Y_WIDTH-1,0)= apData32.range(CAM4K_Y_WIDTH-1, 0);
				rawy = (float)dy_12;
				if(dy_12 == 580)
					int kkk = 1;
				calibPnt.rid = apData32.range(10+CAM4K_Y_WIDTH, CAM4K_Y_WIDTH);
				calibPnt.flag = apData32.range(14+CAM4K_Y_WIDTH, 11+CAM4K_Y_WIDTH);
				dataBuff.dData = InData.calibK[1];
				apData32 = (ap_uint<32>)dataBuff.nData[0];
				calibPnt.value = apData32.range(7,0);
				calibPnt.Sync = InData.Sync;
				calibPnt.rsv = 0;
				if(rawx < 0)
					vldFlag = 0;
				else
					vldFlag = 1;

				x_square = rawx * rawx;
				y_square = rawy * rawy;
			}
			else if( 1 == j)//read subpix_pt1
			{
				calibPnt_pt1.x = (float)calib_x_pt1;
				calibPnt_pt1.y = (float)calib_y_pt1;
				if( (n > 0) && (vldFlag_pt1 == 1))
				{
					calibData_buf[in_dataSize] = calibPnt_pt1;
					in_dataSize++;
					//outSubpixCalib.write(calibPnt_pt1);
				}

				dataBuff.dData = InData.calibK[0];
				rawx_pt1 = (float)dataBuff.fData[0];
				ap_uint<32> apData32 = (ap_uint<32>)dataBuff.nData[1];
				ap_uint<CAM4K_Y_WIDTH> dy_12;
				dy_12(CAM4K_Y_WIDTH-1, 0)= apData32.range(CAM4K_Y_WIDTH-1, 0);
				rawy_pt1 = (float)dy_12;
				calibPnt_pt1.rid = apData32.range(10+CAM4K_Y_WIDTH, CAM4K_Y_WIDTH);
				calibPnt_pt1.flag = apData32.range(14+CAM4K_Y_WIDTH,11+CAM4K_Y_WIDTH);
				dataBuff.dData = InData.calibK[1];
				apData32 = (ap_uint<32>)dataBuff.nData[0];
				calibPnt_pt1.value = apData32.range(7,0);
				calibPnt_pt1.Sync = InData.Sync;
				calibPnt_pt1.rsv = 0;
				if(rawx_pt1 < 0)
					vldFlag_pt1 = 0;
				else
					vldFlag_pt1 = 1;

				x_square_pt1 = rawx_pt1 * rawx_pt1;
				y_square_pt1 = rawy_pt1 * rawy_pt1;
			}
			else if(2 == j)//k7*x^2, k6 * x
			{
				//kx = InData.calibK[0];
				//ky = InData.calibK[1];
				dataConver_x.dData = InData.calibK[0];
				kx_2 = dataConver_x.fData[0];
				ky_2 = dataConver_x.fData[1];
				dataConver_y.dData = InData.calibK[1];
				kx_1 = dataConver_y.fData[0];
				ky_1 = dataConver_y.fData[1];

				s2_result_x = kx_2 * x_square + kx_1 * rawx;
				s2_result_y = ky_2 * y_square + ky_1 * rawy;
				s2_result_x_pt1 = kx_2 * x_square_pt1 + kx_1 * rawx_pt1;
				s2_result_y_pt1 = ky_2 * y_square_pt1 + ky_1 * rawy_pt1;
			}
			else if(3 == j)//k5 * x^2, k4 * x
			{
				dataConver_x.dData = InData.calibK[0];
				kx_2 = dataConver_x.fData[0];
				ky_2 = dataConver_x.fData[1];
				dataConver_y.dData = InData.calibK[1];
				kx_1 = dataConver_y.fData[0];
				ky_1 = dataConver_y.fData[1];

				s3_result_x = (s2_result_x + kx_2) * x_square + kx_1 * rawx;
				s3_result_y = (s2_result_y + ky_2) * y_square + ky_1 * rawy;
				s3_result_x_pt1 = (s2_result_x_pt1 + kx_2) * x_square_pt1 + kx_1 * rawx_pt1;
				s3_result_y_pt1 = (s2_result_y_pt1 + ky_2) * y_square_pt1 + ky_1 * rawy_pt1;
			}
			else if(4 == j) //k3 * x^2, k2 * x
			{
				dataConver_x.dData = InData.calibK[0];
				kx_2 = dataConver_x.fData[0];
				ky_2 = dataConver_x.fData[1];
				dataConver_y.dData = InData.calibK[1];
				kx_1 = dataConver_y.fData[0];
				ky_1 = dataConver_y.fData[1];

				s4_result_x = (s3_result_x + kx_2) * x_square + kx_1 * rawx;
				s4_result_y = (s3_result_y + ky_2) * y_square + ky_1 * rawy;
				s4_result_x_pt1 = (s3_result_x_pt1 + kx_2) * x_square_pt1 + kx_1 * rawx_pt1;
				s4_result_y_pt1 = (s3_result_y_pt1 + ky_2) * y_square_pt1 + ky_1 * rawy_pt1;
			}
			else if(5 == j)//k1 * x, k0
			{
				dataConver_x.dData = InData.calibK[0];
				kx_2 = dataConver_x.fData[0];
				ky_2 = dataConver_x.fData[1];
				dataConver_y.dData = InData.calibK[1];
				kx_1 = dataConver_y.fData[0];
				ky_1 = dataConver_y.fData[1];

				calib_x = (s4_result_x + kx_2) * rawx + kx_1;
				calib_y = (s4_result_y + ky_2) * rawy + ky_1;
				calib_x_pt1 = (s4_result_x_pt1 + kx_2) * rawx_pt1 + kx_1;
				calib_y_pt1 = (s4_result_y_pt1 + ky_2) * rawy_pt1 + ky_1;
			}
		}
	}
	if( (Lines > 0)&& (vldFlag == 1))
	{
		calibPnt.x = (float)calib_x;
		calibPnt.y = (float)calib_y;
		calibData_buf[in_dataSize] = calibPnt;
		in_dataSize ++;
	}

	if( (Lines > 0)&& (vldFlag_pt1 == 1))
	{
		calibPnt_pt1.x = (float)calib_x_pt1;
		calibPnt_pt1.y = (float)calib_y_pt1;
		calibData_buf[in_dataSize] = calibPnt_pt1;
		in_dataSize ++;
	}

	//interpolate
	//maximum 2 points per line
	RgnSubPixCalib pre_data_1 = {0.0f, 0.0f, 0, 0, 0, 0, 0};
	RgnSubPixCalib pre_data_2 = {0.0f, 0.0f, 0, 0, 0, 0, 0};
	uint16_t minY = 0xffff;
	uint16_t maxY = 0;
	uint16_t inter_y_1 = 0;
	uint16_t inter_y_2 = 0;
	uint16_t total_point_num = 0;
	RgnSubPixCalib proc_calib_data = {};
	if(in_dataSize > 0)
		proc_calib_data = calibData_buf[0];  //delpay one data
	for(uint16_t i = 1; i <= in_dataSize; i ++)
	{
#pragma HLS PIPELINE II=2
#pragma HLS LOOP_TRIPCOUNT min=3072 max=3072

		RgnSubPixCalib nxt_calib_data;
		if(i < in_dataSize)
			nxt_calib_data = calibData_buf[i];
		else
		{
			nxt_calib_data.x = 0;
			nxt_calib_data.y = 0;
			nxt_calib_data.rid = 0;
			nxt_calib_data.value = 0;
			nxt_calib_data.flag = 0;
			nxt_calib_data.Sync = 0;
			nxt_calib_data.rsv = 0;
		}

		//if(nxt_calib_data.y == 580)
		//	int kkk = 1;

        const bool proc_valid = (proc_calib_data.rid != 0);
        const bool hit1 = proc_valid && (pre_data_1.rid == proc_calib_data.rid);
        const bool hit2 = proc_valid && (pre_data_2.rid == proc_calib_data.rid);
        const bool keep1 = (pre_data_1.rid != 0) && ((pre_data_1.rid == proc_calib_data.rid) || (pre_data_1.rid == nxt_calib_data.rid));
        const bool keep2 = (pre_data_2.rid != 0) && ((pre_data_2.rid == proc_calib_data.rid) || (pre_data_2.rid == nxt_calib_data.rid));

        if (hit1 || hit2)  //interpolate
		{
            const bool use_slot1 = hit1;
            // 保存旧状态，后续计算使用这些局部快照
            const RgnSubPixCalib pre_data = use_slot1 ? pre_data_1 : pre_data_2;
            const uint16_t start_y = use_slot1 ? inter_y_1 : inter_y_2;
            uint16_t count = 0;
            // 检查范围和正向线段，避免除零及越界转换。
            const bool segment_valid = (proc_calib_data.y > pre_data.y) && (proc_calib_data.y >= 0.0f) && (proc_calib_data.y < (float)ROW_COUNT);

            if (segment_valid && start_y < ROW_COUNT)
            {
                // y已保证非负，转无符号整数等价于向下取整
                const uint16_t end_y = (uint16_t)proc_calib_data.y;
                if (start_y <= end_y)
                {
                    const uint16_t remaining = end_y - start_y + 1;
                    count = (remaining > 3) ? 3 : remaining;
                }
            }

            const uint16_t next_y = start_y + count;

            // 下一轮状态不依赖坐标计算或BRAM写回
            if (use_slot1)
            {
                pre_data_1 = proc_calib_data;
                inter_y_1 = next_y;
            }
            else
            {
                pre_data_2 = proc_calib_data;
                inter_y_2 = next_y;
            }

            if (count > 0)
            {
                float dx = proc_calib_data.x - pre_data.x;
                float dy = proc_calib_data.y - pre_data.y;

#pragma HLS BIND_OP variable=dx op=fsub impl=fabric latency=3
#pragma HLS BIND_OP variable=dy op=fsub impl=fabric latency=3

                const float k = dx / dy;

                // 每个bank至多处理一个插值点。
                // start_y、start_y+1、start_y+2分别落到不同bank。
                const unsigned int start_bank = (unsigned int)start_y % BANK_COUNT;

                ap_uint<1> added[BANK_COUNT];
#pragma HLS ARRAY_PARTITION variable=added complete dim=1
                INIT_ADDED:
                for (unsigned int b = 0; b < BANK_COUNT; ++b)
                {
#pragma HLS UNROLL
                    added[b] = 0;
                }

                INTERPOLATE_BANKS:
                for (uint16_t b = 0; b < BANK_COUNT; ++b)
                {
#pragma HLS UNROLL
                    // 计算属于当前bank的插值点偏移j。
                    // 展开后b是常量，每个分支固定访问rowState[b]。
                    const uint16_t j = (b >= start_bank) ? (b - start_bank) : (b + BANK_COUNT - start_bank);

                    if (j < count)
                    {
                        const uint16_t iy =  start_y + j;

                        const uint16_t bank_addr = iy / BANK_COUNT;

                        // -------- 行槽位分配 --------
                        const ap_uint<42> old_state = rowState[b][bank_addr];
                        const ap_uint<40> old_tag = old_state.range(41, 2);
                        const ap_uint<2> stored_num = old_state.range(1, 0);
                        ap_uint<2> num = 0;
                        if (old_tag == currFrameTag)
                            num = stored_num;

                        const bool accepted = (num < 2);
                        const bool write_pt2 = (num == 1);
                        ap_uint<2> next_num = num;
                        if (accepted)
                            next_num = num + 1;

                        ap_uint<42> new_state = 0;
                        new_state.range(41, 2) = currFrameTag;
                        new_state.range(1, 0) = next_num;

                        // 此状态更新不使用浮点插值结果
                        rowState[b][bank_addr] = new_state;
                        added[b] = accepted ? 1 : 0;

                        // -------- 坐标计算与点数据写入 --------
                        if (accepted)
                        {
                            BuffCalibSubpix inter_point;
                            inter_point.y = (uint16_t)iy;
                            inter_point.x = pre_data.x + k * ((float)iy - pre_data.y);
                            inter_point.rid = pre_data.rid;
                            inter_point.value = pre_data.value;
                            inter_point.matchFlag = 0;
                            inter_point.flag = pre_data.flag;

                            if (write_pt2)
                                interpolate_buf[iy].pt2 = inter_point;
                            else
                                interpolate_buf[iy].pt1 = inter_point;
                        }
                    }
                }

                // 先扩展为普通整数，避免ap_uint<1>加法位宽问题
                const uint16_t added_count = (uint16_t)added[0] + (uint16_t)added[1] + (uint16_t)added[2];
                total_point_num = total_point_num + added_count;

                // 当前处理的行连续，范围无需在展开循环中反复更新。
                // 与原代码相同，即使某行已满，也将访问范围纳入min/max。
                const uint16_t last_y =  start_y + count - 1;

                if (start_y < minY)
                    minY = start_y;
                if (last_y > maxY)
                    maxY = last_y;
            }
        }
        else if (proc_valid)
        {
            // 新region：只接收范围内的点
            const bool y_valid = (proc_calib_data.y >= 0.0f) && (proc_calib_data.y < (float)ROW_COUNT);

            if (y_valid)
            {
                const uint16_t first_y = (uint16_t)ceilf(proc_calib_data.y);

                if (!keep1)
                {
                    pre_data_1 = proc_calib_data;
                    inter_y_1 = first_y;
                }
                else if (!keep2)
                {
                    pre_data_2 = proc_calib_data;
                    inter_y_2 = first_y;
                }
            }
        }
		proc_calib_data = nxt_calib_data;
	}

	//output
	InterpolateCalibSubpix outData = {0.0f, 0, 0, 0, 0, 0};
	UintFloat tmp;
	//1st: frameNo
	outData.Sync = 0b10;
	tmp.unData = frmNo;
	outData.x = tmp.fdata;
	outSubpixCalib.write(outData);
	//2nd: timeStamp
	outData.Sync = 0b00;
	tmp.unData = timeStamp;
	outData.x = tmp.fdata;
	outSubpixCalib.write(outData);
	//3rd:encInfo
	tmp.unData = encInfo;
	outData.x = tmp.fdata;
	outSubpixCalib.write(outData);
	//4th: frmW, frmH
	tmp.unData = frmWH;
	outData.x = tmp.fdata;
	outSubpixCalib.write(outData);
	//5th: frmX, WinNum
	tmp.unData = (total_point_num << 16) | frmX;
	outData.x = tmp.fdata;
	outSubpixCalib.write(outData);
	//6th: frmY
	tmp.unData = frmY;
	outData.x = tmp.fdata;
	if(total_point_num == 0)
		outData.Sync = 0b11;
	else
		outData.Sync = 0b00;
	outSubpixCalib.write(outData);

	if(total_point_num > 0)
	{
		uint16_t out_counter = 0;
		uint16_t start = minY * 2;
		uint16_t end = maxY * 2 + 2;
		for(uint16_t loop = start; loop < end; loop ++)
		{
#pragma HLS PIPELINE II=1
#pragma HLS LOOP_TRIPCOUNT min=3072*2 max=3072*2

			ap_uint<1> ptIdx = loop % 2;
			uint16_t i = loop / 2;

			const uint16_t bank = i % BANK_COUNT;
			const uint16_t addr = i / BANK_COUNT;
			const ap_uint<42> state = rowState[bank][addr];
			const ap_uint<40> ptFrameTag = state.range(41, 2);
			if( (ptFrameTag == currFrameTag) && (out_counter < total_point_num))
			{
			    const uint16_t linePtNum = (uint16_t)state.range(1, 0);
				if( (linePtNum > 0) && (ptIdx == 0))
				{
					BuffCalibSubpix pt1 = interpolate_buf[i].pt1;
					outData.x = pt1.x;
					outData.y = pt1.y;
					outData.rid = pt1.rid;
					outData.value = pt1.value;
					outData.flag = pt1.flag;
					out_counter ++;
					if(out_counter == total_point_num)
						outData.Sync = 0b11;
					else
						outData.Sync = 0b00;
					outSubpixCalib.write(outData);
				}
				else if( (linePtNum == 2) && (ptIdx == 1))
				{
					BuffCalibSubpix pt2 = interpolate_buf[i].pt2;
					outData.x = pt2.x;
					outData.y = pt2.y;
					outData.rid = pt2.rid;
					outData.value = pt2.value;
					outData.flag = pt2.flag;
					out_counter ++;
					if(out_counter == total_point_num)
						outData.Sync = 0b11;
					else
						outData.Sync = 0b00;
					outSubpixCalib.write(outData);
				}
			}
		}
	}
}