// Copyright (2019-2023) Paul Scherrer Institute

#include <iostream>
#include <catch2/catch.hpp>

#include "../fpga/hls/hls_jfjoch.h"
#include "../fpga/hls/spot_finder.h"

TEST_CASE("Pack32_Unpack32","[FPGA][SpotFinder]") {
    ap_int<16> value[32];
    for (int i = 0; i < 32; i++)
        value[i] = i;

    ap_int<512> packed = pack32(value);
    ap_int<16> restore[32];
    unpack32(packed, restore);
    for (int i = 0; i < 32; i++) {
        REQUIRE(value[i] == restore[i]);
    }

}

TEST_CASE("Pack32_Unpack32_36","[FPGA][SpotFinder]") {
    ap_int<36> value[32];
    for (int i = 0; i < 32; i++)
        value[i] = i;

    ap_int<36*32> packed = pack32(value);
    ap_int<36> restore[32];
    unpack32(packed, restore);
    for (int i = 0; i < 32; i++) {
        REQUIRE(value[i] == restore[i]);
    }
}

TEST_CASE("FPGA_calc_sum","[FPGA][SpotFinder]") {
    ap_int<16> old_value[32];
    ap_int<16> new_value[32];

    ap_int<SUM_BITWIDTH> new_sum[32];
    ap_int<SUM2_BITWIDTH> new_sum2[32];

    for (int i = 0; i < 32; i++) {
        old_value[i] = i;
        new_value[i] = 2;
    }

    ap_uint<SUM_BITWIDTH*32> diff_sum;
    ap_uint<SUM2_BITWIDTH*32> diff_sum2;

    calc_sum(diff_sum, pack32(old_value), pack32(new_value));
    calc_sum2(diff_sum2, pack32(old_value), pack32(new_value));
    unpack32(diff_sum, new_sum);
    unpack32(diff_sum2, new_sum2);

    for (int i = 0; i < 32; i++) {
        REQUIRE(new_sum[i] == 2-i);
        REQUIRE(new_sum2[i] == 4-i*i);
    }
}

TEST_CASE("FPGA_calc_valid","[FPGA][SpotFinder]") {
    ap_uint<32> old_mask = UINT32_MAX;
    ap_uint<32> new_mask = UINT32_MAX;

    old_mask[15] = 0;
    new_mask[13] = 0;

    ap_uint<MASK_SUM_BITWIDTH*32> diff_mask;
    calc_mask_diff(diff_mask, old_mask, new_mask);

    ap_int<MASK_SUM_BITWIDTH> diff_mask_32[32];
    unpack32(diff_mask, diff_mask_32);

    REQUIRE(diff_mask_32[0] == 0);
    REQUIRE(diff_mask_32[1] == 0);

    REQUIRE(diff_mask_32[13] == -1);
    REQUIRE(diff_mask_32[15] == 1);
}

TEST_CASE("FPGA_calc_mask","[FPGA][SpotFinder]") {
    ap_int<16> value_in[32], value_out[32];

    for (int i = 0; i < 32; i++)
        value_in[i] = 154 + i;

    value_in[15] = INT16_MAX;
    value_in[0] = INT16_MIN;
    value_in[1] = INT16_MIN + 1;
    value_in[2] = INT16_MAX - 1;

    ap_uint<512> input = pack32(value_in);
    ap_uint<512> output = 0;
    ap_uint<32> mask = 0;

    calc_mask(input, output, mask);

    REQUIRE(mask[0] == 0);
    REQUIRE(mask[15] == 0);
    REQUIRE(mask[1] == 1);
    REQUIRE(mask[2] == 1);
    REQUIRE(mask[3] == 1);
    REQUIRE(mask[4] == 1);

    unpack32(output, value_out);

    REQUIRE(value_out[0] == 0);
    REQUIRE(value_out[15] == 0);
    REQUIRE(value_out[1] == value_in[1]);
    REQUIRE(value_out[2] == value_in[2]);
    REQUIRE(value_out[3] == value_in[3]);
    REQUIRE(value_out[4] == value_in[4]);
}

TEST_CASE("FPGA_update_sum" , "[FPGA][SpotFinder]") {
    ap_int<SUM2_BITWIDTH> arr_val1[32], arr_val2[32], arr_out[32];

    for (int i = 0; i < 32; i++) {
        arr_val1[i] = 5 * i;
        arr_val2[i] = 3 * i + 2;
    }
    ap_uint<SUM2_BITWIDTH*32> val1 = pack32(arr_val1);
    ap_uint<SUM2_BITWIDTH*32> val2 = pack32(arr_val2);
    update_sum<SUM2_BITWIDTH>(val1, val2);
    unpack32(val1, arr_out);
    for (int i = 0; i < 32; i++)
        REQUIRE(arr_out[i] == 8 * i + 2);
}

int sum_consecutive(int x0, int n) {
    int ret = 0;
    for (int i = 0; i < n; i++)
        ret += x0 + i;
    return ret;
}

TEST_CASE("FPGA_prefix_sum" , "[FPGA][SpotFinder]") {
    ap_uint<16*(32+2*FPGA_NBX)> input = 0;
    for (int i = 0; i < 32+2*FPGA_NBX; i++)
        input(i*16+15, i*16) = i;

    auto output = prefix_sum<16>(input);
    ap_uint<16> output_unpacked[32];
    unpack32(output, output_unpacked);
    REQUIRE(output_unpacked[0] == sum_consecutive(0, 2 * FPGA_NBX + 1));
    REQUIRE(output_unpacked[2] == sum_consecutive(2, 2 * FPGA_NBX + 1));
    REQUIRE(output_unpacked[7] == sum_consecutive(7, 2 * FPGA_NBX + 1));
}

bool Isigma_cpu(double val, double sum, double sum2, float threshold) {
    double mean = sum / ((2*FPGA_NBX+1) * (2*FPGA_NBX+1));
    double mean2 = sum2 / ((2*FPGA_NBX+1) * (2*FPGA_NBX+1));
    double variance = mean2 - mean * mean;
    double sigma = sqrt(variance);
    double i_over_sigma =  (val - mean) / sigma;
    return (i_over_sigma > threshold);
}

bool Isigma_fpga(ap_int<16> val, ap_int<SUM_BITWIDTH> sum, ap_uint<SUM2_BITWIDTH> sum2, float threshold) {
    return check_threshold(val, sum, sum2, (FPGA_NBX *2 + 1) * (FPGA_NBX *2 + 1),
                           threshold * threshold, -1);
}

TEST_CASE("FPGA_spot_check_threshold","[FPGA][SpotFinder]") {
    std::vector<float> threshold_values = {1.0, 3.0, 6.0};

    for (auto threshold: threshold_values) {
        uint32_t uniform_val = 10;
        uint32_t diff = 0;
        for (int16_t val = -100; val < INT16_MAX; val++) {
            uint32_t sum = val + ((2 * FPGA_NBX + 1) * (2 * FPGA_NBX + 1) - 1) * uniform_val;
            uint32_t sum2 = val * val + ((2 * FPGA_NBX + 1) * (2 * FPGA_NBX + 1) - 1) * uniform_val * uniform_val;
            bool cpu = Isigma_cpu(val, sum, sum2, threshold);
            bool fpga = Isigma_fpga(val, sum, sum2, threshold);
            if (cpu != fpga) {
                std::cout <<  "WRONG!!! " << threshold << " " << val << " " << sum << " " << sum2 << " CPU: " << cpu << " FPGA: " << fpga << std::endl;
                diff++;
            }
        }
        REQUIRE(diff == 0);
    }
}

TEST_CASE("FPGA_spot_finder_update_sum","[FPGA][SpotFinder]") {
    STREAM_512 input;
    hls::stream<spot_finder_packet> output;

    hls::stream<ap_uint<(SUM_BITWIDTH*32)>> sum_out;
    hls::stream<ap_uint<SUM2_BITWIDTH*32>> sum2_out;
    hls::stream<ap_uint<MASK_SUM_BITWIDTH*32>> valid_out;

    std::vector<int16_t> input_frame(RAW_MODULE_SIZE), output_frame(RAW_MODULE_SIZE);
    for (int i = 0; i < RAW_MODULE_SIZE; i++) {
        if (i % RAW_MODULE_COLS == 1023)
            input_frame[i] = INT16_MIN;
        else
            input_frame[i] = i % RAW_MODULE_COLS;
    }
    auto input_frame_512 = (ap_uint<512> *) input_frame.data();
    auto output_frame_512 = (ap_uint<512> *) output_frame.data();

    input << packet_512_t{.user = 0};
    for (int i = 0; i < RAW_MODULE_SIZE * sizeof(uint16_t) / 64; i++)
        input << packet_512_t{.data = input_frame_512[i], .user = 0};

    input << packet_512_t{.user = 1};

    spot_finder_col_sum(input, output, sum_out, sum2_out, valid_out);

    REQUIRE(input.size() == 0);

    REQUIRE(output.size()    == RAW_MODULE_SIZE * sizeof(uint16_t) / 64 + 2);
    REQUIRE(sum_out.size()   == RAW_MODULE_SIZE * sizeof(uint16_t) / 64);
    REQUIRE(sum2_out.size()  == RAW_MODULE_SIZE * sizeof(uint16_t) / 64);
    REQUIRE(valid_out.size() == RAW_MODULE_SIZE * sizeof(uint16_t) / 64);

    std::vector<int64_t> sum(RAW_MODULE_SIZE);
    std::vector<int64_t> sum2(RAW_MODULE_SIZE);
    std::vector<int64_t> valid(RAW_MODULE_SIZE);

    for (int i = 0; i < RAW_MODULE_SIZE * sizeof(uint16_t) / 64; i++) {
        ap_uint<32 * SUM_BITWIDTH> tmp_sum;
        ap_uint<32 * SUM2_BITWIDTH> tmp_sum2;
        ap_uint<32 * MASK_SUM_BITWIDTH> tmp_valid;
        sum_out   >> tmp_sum;
        sum2_out  >> tmp_sum2;
        valid_out >> tmp_valid;

        ap_uint<SUM_BITWIDTH> tmp_sum_unpacked[32];
        ap_uint<SUM2_BITWIDTH> tmp_sum2_unpacked[32];
        ap_uint<MASK_SUM_BITWIDTH> tmp_valid_unpacked[32];
        unpack32(tmp_sum, tmp_sum_unpacked);
        unpack32(tmp_sum2, tmp_sum2_unpacked);
        unpack32(tmp_valid, tmp_valid_unpacked);

        for (int j = 0; j < 32; j++) {
            sum[i * 32 + j] = tmp_sum_unpacked[j];
            sum2[i * 32 + j] = tmp_sum2_unpacked[j];
            valid[i * 32 + j] = tmp_valid_unpacked[j];
        }
    }

    CHECK(sum[1] == (FPGA_NBX+1) * 1);
    CHECK(sum[3] == (FPGA_NBX+1) * 3);
    CHECK(sum[1023] == 0);
    CHECK(sum[1022+200*1024] == (2 * FPGA_NBX+1) * 1022);

    CHECK(sum2[3] == (FPGA_NBX+1) * 3 * 3);
    CHECK(sum2[1023] == 0);
    CHECK(sum2[1022+200*1024] == (2 * FPGA_NBX+1) * 1022 * 1022);

    CHECK(valid[1] == FPGA_NBX + 1);
    CHECK(valid[3] == FPGA_NBX + 1);
    CHECK(valid[1023] == 0);
    CHECK(valid[1023 + 323*1024] == 0);

    CHECK(valid[1+1024] == FPGA_NBX + 1 + 1);
    CHECK(valid[1+1024] == FPGA_NBX + 1 + 1);
    CHECK(valid[1+3*1024] == FPGA_NBX + 1 + 3);
    CHECK(valid[1+200*1024] == 2 * FPGA_NBX + 1);

    CHECK(valid[1+509*1024] == FPGA_NBX + 1 + 2);
    CHECK(valid[1+510*1024] == FPGA_NBX + 1 + 1);
    CHECK(valid[1+511*1024] == FPGA_NBX + 1);

    spot_finder_packet packet_out;
    output >> packet_out;
    for (int i = 0; i < RAW_MODULE_SIZE * sizeof(uint16_t) / 64; i++) {
        output >> packet_out;
        output_frame_512[i] = packet_out.data;
    }
    REQUIRE(output_frame == input_frame);
}

TEST_CASE("FPGA_spot_finder_line_sum","[FPGA][SpotFinder]") {
    hls::stream<spot_finder_packet> input;
    hls::stream<spot_finder_packet> stream_0;
    hls::stream<spot_finder_packet> output;

    hls::stream<ap_uint<SUM_BITWIDTH*32>> sum_in;
    hls::stream<ap_uint<SUM2_BITWIDTH*32>> sum2_in;
    hls::stream<ap_uint<MASK_SUM_BITWIDTH*32>> valid_in;

    hls::stream<ap_uint<SUM_BITWIDTH*32>> sum_stream;
    hls::stream<ap_uint<SUM2_BITWIDTH*32>> sum2_stream;
    hls::stream<ap_uint<MASK_SUM_BITWIDTH*32>> valid_stream;

    hls::stream<ap_uint<SUM_BITWIDTH*32>> sum_out;
    hls::stream<ap_uint<SUM2_BITWIDTH*32>> sum2_out;
    hls::stream<ap_uint<MASK_SUM_BITWIDTH*32>> valid_out;

    ap_int<SUM_BITWIDTH> sum_unpacked[32];
    ap_int<SUM2_BITWIDTH> sum2_unpacked[32];
    ap_int<MASK_SUM_BITWIDTH> valid_unpacked[32];

    std::vector<int16_t> input_frame(RAW_MODULE_COLS), output_frame(RAW_MODULE_COLS);
    for (int i = 0; i < RAW_MODULE_COLS; i++) {
        if (i % RAW_MODULE_COLS == 1023)
            input_frame[i] = INT16_MIN;
        else
            input_frame[i] = i % RAW_MODULE_COLS;
    }
    auto input_frame_512 = (ap_uint<512> *) input_frame.data();
    auto output_frame_512 = (ap_uint<512> *) output_frame.data();

    input << spot_finder_packet{.user = 0};
    for (int i = 0; i < 32; i++)
        input << spot_finder_packet{.data = input_frame_512[i], .user = 0};

    input << spot_finder_packet{.user = 1};

    for (int i = 0; i < 32; i++) {
        for (int j = 0; j < 32; j++) {
            sum_unpacked[j] = i * 32 + j;
            sum2_unpacked[j] = 8934 + (i * 32 + j);
            valid_unpacked[j] = 1;
        }
        sum_in << pack32(sum_unpacked);
        sum2_in << pack32(sum2_unpacked);
        valid_in << pack32(valid_unpacked);
    }

    spot_finder_line_sum(input, stream_0, sum_in, sum2_in, valid_in, sum_stream, sum2_stream, valid_stream);
    REQUIRE(input.size() == 0);

    REQUIRE(stream_0.size()     == 32 + 2);
    REQUIRE(sum_stream.size()   == 33);
    REQUIRE(sum2_stream.size()  == 33);
    REQUIRE(valid_stream.size() == 33);

    spot_finder_line_sum_align(stream_0, output, sum_stream, sum2_stream, valid_stream, sum_out, sum2_out, valid_out);

    REQUIRE(stream_0.size() == 0);

    REQUIRE(output.size()    == 32 + 2);
    REQUIRE(sum_out.size()   == 32);
    REQUIRE(sum2_out.size()  == 32);
    REQUIRE(valid_out.size() == 32);

    std::vector<int64_t> sum_output(RAW_MODULE_COLS);
    std::vector<int64_t> sum2_output(RAW_MODULE_COLS);
    std::vector<int64_t> valid_output(RAW_MODULE_COLS);

    for (int i = 0; i < 32; i++) {
        ap_uint<32 * SUM_BITWIDTH> sum_tmp;
        sum_out >> sum_tmp;
        unpack32(sum_tmp, sum_unpacked);
        for (int j = 0; j < 32; j++)
            sum_output[i * 32 + j] = sum_unpacked[j];

        ap_uint<32 * SUM2_BITWIDTH> sum2_tmp;
        sum2_out >> sum2_tmp;
        unpack32(sum2_tmp, sum2_unpacked);
        for (int j = 0; j < 32; j++)
            sum2_output[i * 32 + j] = sum2_unpacked[j];

        ap_uint<32 * MASK_SUM_BITWIDTH> valid_tmp;
        valid_out >> valid_tmp;
        unpack32(valid_tmp, valid_unpacked);
        for (int j = 0; j < 32; j++)
            valid_output[i * 32 + j] = valid_unpacked[j];
    }

    CHECK(sum_output[0] == sum_consecutive(0, FPGA_NBX+1));
    CHECK(sum_output[1] == sum_consecutive(0, FPGA_NBX+2));
    CHECK(sum_output[2] == sum_consecutive(0, FPGA_NBX+3));
    CHECK(sum_output[FPGA_NBX] == sum_consecutive(0, FPGA_NBX*2+1));
    CHECK(sum_output[FPGA_NBX+1] == sum_consecutive(1, FPGA_NBX*2+1));
    CHECK(sum_output[FPGA_NBX+5] == sum_consecutive(5, FPGA_NBX*2+1));
    CHECK(sum_output[1023] == sum_consecutive(1023-FPGA_NBX, FPGA_NBX+1));

    CHECK(sum2_output[0] == sum_consecutive(8934, FPGA_NBX+1));

    CHECK(valid_output[0] == FPGA_NBX+1);
    CHECK(valid_output[1] == FPGA_NBX+2);

    CHECK(valid_output[1022] == FPGA_NBX+2);
    CHECK(valid_output[1023] == FPGA_NBX+1);

    spot_finder_packet packet_out;
    output >> packet_out;
    for (int i = 0; i < 32; i++) {
        output >> packet_out;
        output_frame_512[i] = packet_out.data;
    }
}

TEST_CASE("FPGA_spot_finder_core","[FPGA][SpotFinder]") {
    STREAM_512 input;
    STREAM_512 output;
    hls::stream<ap_axiu<32,1,1,1>> strong_pixel;

    ap_int<16> in_photon_count_threshold = 8;
    ap_uint<16> in_strong_pixel_threshold = 16;

    std::vector<int16_t> input_frame(RAW_MODULE_SIZE), output_frame(RAW_MODULE_SIZE);
    for (int i = 0; i < RAW_MODULE_SIZE; i++) {
        if (i % RAW_MODULE_COLS == 1023)
            input_frame[i] = INT16_MIN;
        else
            input_frame[i] = i % RAW_MODULE_COLS;
    }
    auto input_frame_512 = (ap_uint<512> *) input_frame.data();
    auto output_frame_512 = (ap_uint<512> *) output_frame.data();

    input << packet_512_t{.user = 0};
    for (int i = 0; i < RAW_MODULE_SIZE * sizeof(uint16_t) / 64; i++)
        input << packet_512_t{.data = input_frame_512[i], .user = 0};

    input << packet_512_t{.user = 1};

    spot_finder(input,
                output,
                strong_pixel,
                in_photon_count_threshold,
                in_strong_pixel_threshold);

    REQUIRE(input.size() == 0);
    REQUIRE(output.size() == RAW_MODULE_SIZE * sizeof(uint16_t) / 64 + 2);
    REQUIRE(strong_pixel.size() == RAW_MODULE_SIZE * sizeof(uint16_t) / 64 + 1);
}