// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: CERN-OHL-S-2.0

#include "hls_jfjoch.h"

void spot_finder_mask(STREAM_768 &data_in,
                      STREAM_768 &data_out,
                      hls::stream<axis_completion > &s_axis_completion,
                      hls::stream<axis_completion > &m_axis_completion,
                      ap_uint<256> *d_hbm_p0,
                      ap_uint<256> *d_hbm_p1,
                      volatile ap_uint<32> &in_min_d_value,
                      volatile ap_uint<32> &in_max_d_value,
                      ap_uint<32> hbm_size_bytes) {
#pragma HLS INTERFACE ap_ctrl_none port=return

#pragma HLS INTERFACE register both axis port=data_in
#pragma HLS INTERFACE register both axis port=data_out
#pragma HLS INTERFACE register both axis port=m_axis_completion
#pragma HLS INTERFACE register both axis port=s_axis_completion
#pragma HLS INTERFACE register ap_none port=in_min_d_value
#pragma HLS INTERFACE register ap_none port=in_max_d_value
#pragma HLS INTERFACE register ap_none port=hbm_size_bytes

#pragma HLS INTERFACE m_axi port=d_hbm_p0 bundle=d_hbm_p0 depth=16384 offset=off \
        max_read_burst_length=16 max_write_burst_length=2 latency=120 num_write_outstanding=2 num_read_outstanding=8
#pragma HLS INTERFACE m_axi port=d_hbm_p1 bundle=d_hbm_p1 depth=16384 offset=off \
        max_read_burst_length=16 max_write_burst_length=2 latency=120 num_write_outstanding=2 num_read_outstanding=8

    packet_768_t packet_in;
    {
#pragma HLS PROTOCOL fixed
        data_in >> packet_in;
        ap_wait();
        data_out << packet_in;
        ap_wait();
    }
    ap_uint<32> offset_hbm_0         = 16 * hbm_size_bytes / 32;
    ap_uint<32> offset_hbm_1         = 17 * hbm_size_bytes / 32;

    ap_int<24> pixel_val[32];
    xray_d_t d[32];

    axis_completion cmpl;
    s_axis_completion >> cmpl;
    while (!cmpl.last) {
        m_axis_completion << cmpl;

        xray_d_t min_d_value = float_conv<xray_d_t>(in_min_d_value.read());
        xray_d_t max_d_value = float_conv<xray_d_t>(in_max_d_value.read());

        for (int i = 0; i < RAW_MODULE_SIZE / 32; i++) {
#pragma HLS PIPELINE II=1
            data_in >> packet_in;

            ap_uint<256> d_0 = d_hbm_p0[offset_hbm_0 + cmpl.module * RAW_MODULE_SIZE * sizeof(int16_t) / 64 + i];
            ap_uint<256> d_1 = d_hbm_p1[offset_hbm_1 + cmpl.module * RAW_MODULE_SIZE * sizeof(int16_t) / 64 + i];

            unpack_2xhbm_to_32x16bit(d_0, d_1, d);

            unpack32(packet_in.data, pixel_val);
            ap_uint<32> mask_val = 0;
            ap_uint<32> strong_pixel = 0;

            ap_uint<9> line = i / 32;
            ap_uint<5> col = i % 32;

            for (int j = 0; j < 32; j++) {
                if ((line == 0)
                    || (line == 255)
                    || (line == 256)
                    || (line == 511)
                    || (((col == 7) || (col == 15) || (col == 23) || (col == 31)) && (j == 31))
                    || (((col == 0) || (col == 8) || (col == 16) || (col == 24)) && (j ==  0))
                    || (pixel_val[j] == INT24_MIN)
                    || (d[j] < min_d_value)
                    || (d[j] > max_d_value)
                    || !packet_in.strb[j]) {
                    mask_val[j] = 0;
                    strong_pixel[j] = 0;
                } else {
                    // All the conditions above are false, but very high number (above UINT20_MAX) needs a special treatment.
                    // It is masked (not included in st. dev. calculation), but it is ALWAYS considered a strong pixel.
                    if (pixel_val[j] == INT24_MAX) {
                        mask_val[j] = 0;
                        strong_pixel[j] = 1;
                    } else {
                        mask_val[j] = 1;
                        strong_pixel[j] = 0;
                    }
                }
            }
            packet_in.keep = (strong_pixel, mask_val);

            data_out << packet_in;
        }
        s_axis_completion >> cmpl;
    }
    m_axis_completion << cmpl;

    data_in >> packet_in;
    data_out << packet_in;
}