#include "compute3D.h"
#include "stdlib.h"
#include <cstdint>

void compute3D(
	hls::stream<LRMatchedPoint> &inLRPt,
	hls::stream<Pnt3D> 			&outPt3D,
	float uL_offset,
	float u0_f,
	float v0_f,
	float F_f,
	float T_f,
	float plane_a_f,  //d = a*UL + b*V + c
	float plane_b_f,
	float plane_c_f,
	float plane_r,  //d = r * (a*UR + b*V + c)
	ap_uint<8> noise_filter_len
)
{
#if !(VITIS_HLS)
#pragma HLS DATA_PACK variable=inLRPt
#pragma HLS DATA_PACK variable=outPt3D
#else
#pragma HLS aggregate variable=inLRPt compact=bit
#pragma HLS aggregate variable=outPt3D compact=bit
#endif
#pragma HLS INTERFACE axis register both port=inLRPt
#pragma HLS INTERFACE axis register both port=outPt3D
#pragma HLS INTERFACE mode=ap_none port=uL_offset register
#pragma HLS INTERFACE mode=ap_none port=u0_f register
#pragma HLS INTERFACE mode=ap_none port=v0_f register
#pragma HLS INTERFACE mode=ap_none port=F_f register
#pragma HLS INTERFACE mode=ap_none port=T_f register
#pragma HLS INTERFACE mode=ap_none port=plane_a_f register
#pragma HLS INTERFACE mode=ap_none port=plane_b_f register
#pragma HLS INTERFACE mode=ap_none port=plane_c_f register
#pragma HLS INTERFACE mode=ap_none port=plane_r register
#pragma HLS INTERFACE mode=ap_none port=noise_filter_len register

    // Capacity is independent of MAX_PT_NUM/CAM_MAX_Y (which may be indices).
	constexpr unsigned POINT_CAPACITY = ROW_COUNT; // 3072
	constexpr unsigned MAX_SEGMENTS_PER_FRAME = ROW_COUNT;
	constexpr unsigned SEGMENT_CAPACITY = MAX_SEGMENTS_PER_FRAME + 1u; // 3073

    static LRMatchedPoint matchPt_buf[POINT_CAPACITY];
#pragma HLS BIND_STORAGE variable=matchPt_buf type=ram_2p impl=bram
#pragma HLS RESET variable=matchPt_buf off

    // Power-on initialization only; no ordinary per-frame clearing.
    // [55:16] internal frame generation, [15:0] point count.
    static ap_uint<56> segmentState[SEGMENT_CAPACITY] = {};
#pragma HLS BIND_STORAGE variable=segmentState type=ram_2p impl=bram
#pragma HLS RESET variable=segmentState off
    static ap_uint<40> frameGeneration = 0;
#pragma HLS RESET variable=frameGeneration off
    // Keep BOTH generation and state across ap_rst to avoid reusing old tags.

	uint32_t frmNo, timeStamp, encInfo;
	uint16_t minY, maxY;

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

	UintFloat inTmp;
	inTmp.fdata = InData.xL;
	frmNo = inTmp.unData;
	inTmp.fdata = InData.xR;
	timeStamp = inTmp.unData;

	InData = inLRPt.read();
	inTmp.fdata = InData.xL;
	encInfo = inTmp.unData;
	inTmp.fdata = InData.xR;
	minY = inTmp.unData & 0xffff;
	maxY = inTmp.unData >> 16;

	InData = inLRPt.read();
	inTmp.fdata = InData.xL;
    // Contract: header count <= POINT_CAPACITY; rows may contain two points.
    // Over-capacity frames are consumed but emitted as empty (no partial cloud).
    const uint32_t declared_count = inTmp.unData;
    const bool capacity_ok = declared_count <= POINT_CAPACITY;

    const ap_uint<40> max_generation = 0xffffffffff;
    if (frameGeneration == max_generation)
    {
        // Only on generation wrap (once per 2^40-1 accepted frames).
        // Publish generation 1 AFTER clearing, so interrupted clearing retries.
GENERATION_WRAP_CLEAR:
        for (unsigned s = 0; s < SEGMENT_CAPACITY; ++s)
        {
#pragma HLS PIPELINE II=1
            segmentState[s] = 0;
        }
        frameGeneration = 1;
    }
    else
        frameGeneration = frameGeneration + 1;

    const ap_uint<40> currFrameTag = frameGeneration;
    uint16_t received_count = 0;

RECEIVE_POINTS:
    for (uint32_t i = 0; i < declared_count; ++i)
    {
#pragma HLS LOOP_TRIPCOUNT min=0 max=3072
#pragma HLS PIPELINE II=3
        InData = inLRPt.read();
        if (capacity_ok)
        {
            const uint16_t seg_id = InData.match_seg_id;
            if (seg_id != 0 && seg_id < SEGMENT_CAPACITY)
            {
                const ap_uint<56> old_state = segmentState[seg_id];
                const ap_uint<40> old_tag = old_state.range(55, 16);
                uint16_t seg_size = 0;
                if (old_tag == currFrameTag)
                    seg_size = (uint16_t)old_state.range(15, 0);
    
                ap_uint<56> new_state = 0;
                new_state.range(55, 16) = currFrameTag;
                new_state.range(15, 0) = (uint16_t)(seg_size + 1);
                segmentState[seg_id] = new_state;
            }
            else
            {
                // Invalid segment: consume the point, but do not emit it.
                InData.matchFlag = 0;
            }
            matchPt_buf[i] = InData;
            received_count = i + 1;
        }
        if (InData.Sync == 0b11)
            break;
    }

	uint16_t outSize = 0;
FILTER_POINTS:
	for(uint16_t i = 0; i < received_count; i ++)
	{
#pragma HLS LOOP_TRIPCOUNT min=0 max=3072
#pragma HLS PIPELINE II=1

		LRMatchedPoint a_matchPt = matchPt_buf[i];
		uint16_t seg_id = a_matchPt.match_seg_id;
        uint16_t seg_size = 0;
        if (seg_id != 0 && seg_id < SEGMENT_CAPACITY)
        {
            const ap_uint<56> state = segmentState[seg_id];
            seg_size = (uint16_t)state.range(15, 0);
        }
		if(seg_size <= noise_filter_len)
			matchPt_buf[i].matchFlag = 0;
		else if(a_matchPt.matchFlag != 0)
			outSize ++;
	}

	Pnt3D outData = {0.0f, 0.0f, 0.0f, 0, 0, 0, 0};
	UintFloat outTmp;
	outTmp.unData = frmNo;
	outData.x = outTmp.fdata;
	outTmp.unData = timeStamp;
	outData.y = outTmp.fdata;
	outTmp.unData = encInfo;
	outData.z = outTmp.fdata;
	outData.Sync = 0b10;
	outPt3D.write(outData);

	outTmp.unData = outSize;
	outData.x = outTmp.fdata;
	outData.y = 0.0f;
	outData.z = 0.0f;
	if(outSize == 0)
		outData.Sync = 0b11;
	else
		outData.Sync = 0b00;
	outPt3D.write(outData);

	if(outSize > 0)
	{
		uint16_t outCounter = 0;
OUTPUT_POINTS:
		for(uint16_t i = 0; i < received_count; i ++)
		{
#pragma HLS LOOP_TRIPCOUNT min=0 max=3072
#pragma HLS PIPELINE II=1

			LRMatchedPoint a_matchPt = matchPt_buf[i];
			float y = (float)a_matchPt.y;
			float disparity = 0;
			if(a_matchPt.matchFlag == 1) //L-R match
				disparity = a_matchPt.xL + uL_offset - a_matchPt.xR;
			else if(a_matchPt.matchFlag == 2) //left single
				disparity = (a_matchPt.xL + uL_offset) * plane_a_f + y * plane_b_f + plane_c_f;
			else if(a_matchPt.matchFlag == 3) //left single
				disparity = plane_r * (a_matchPt.xR * plane_a_f + y * plane_b_f + plane_c_f);

			if(a_matchPt.matchFlag > 0)
			{
				double d = disparity;
				float u;
				if(a_matchPt.matchFlag == 3)
					u = a_matchPt.xR - u0_f + disparity;
				else
					u = a_matchPt.xL - u0_f;// + uL_offset;
				float v = y - v0_f;
				double TP = (double)T_f / d;
				float TP_f = (float)TP;

				float out_x = u * TP_f;
				float out_y = v * TP_f;
				float out_z = F_f * TP_f;

				outData.x = out_x;
				outData.y = out_y;
				outData.z = out_z;
				outData.valueL = a_matchPt.valueL;
				outData.valueR = a_matchPt.valueR;
				outData.matchFlag = a_matchPt.matchFlag;
				outCounter ++;
				if(outCounter == outSize)
					outData.Sync = 0b11;
				else
					outData.Sync = 0b00;

				outPt3D.write(outData);
			}
		}
	}
}