// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "ImageSpotFinderCPU.h"
#include "StrongPixelSet.h"

ImageSpotFinderCPU::ImageSpotFinderCPU(const std::vector<bool> &mask, size_t in_width, size_t in_height)
    : width(in_width), height(in_height), pxl_mask(mask) {
    if (mask.size() != width * height)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Coordinate mismatch of mask size");
}

void ImageSpotFinderCPU::nbx(int input) {
    NBX = input;
}

std::vector<DiffractionSpot> ImageSpotFinderCPU::Run(const int32_t *input, const SpotFindingSettings &settings,
                                     const std::vector<bool> &res_mask) {
    return Run<int32_t>(input, INT32_MIN, settings, res_mask);
}

std::vector<DiffractionSpot> ImageSpotFinderCPU::Run(const int16_t *input, const SpotFindingSettings &settings,
                                     const std::vector<bool> &res_mask) {
    return Run<int16_t>(input, INT16_MIN, settings, res_mask);
}

std::vector<DiffractionSpot> ImageSpotFinderCPU::Run(const int8_t *input, const SpotFindingSettings &settings,
                                     const std::vector<bool> &res_mask) {
    return Run<int8_t>(input, INT8_MIN, settings, res_mask);
}

std::vector<DiffractionSpot> ImageSpotFinderCPU::Run(const uint32_t *input, const SpotFindingSettings &settings,
                                     const std::vector<bool> &res_mask) {
    return Run<uint32_t>(input, UINT32_MAX, settings, res_mask);
}

std::vector<DiffractionSpot> ImageSpotFinderCPU::Run(const uint16_t *input, const SpotFindingSettings &settings,
                                     const std::vector<bool> &res_mask) {
    return Run<uint16_t>(input, UINT16_MAX, settings, res_mask);
}

std::vector<DiffractionSpot> ImageSpotFinderCPU::Run(const uint8_t *input, const SpotFindingSettings &settings,
                                     const std::vector<bool> &res_mask) {
    return Run<uint8_t>(input, UINT8_MAX, settings, res_mask);
}

template<class T>
std::vector<DiffractionSpot> ImageSpotFinderCPU::Run(const T *input,
                                                  const T special_value,
                                                  const SpotFindingSettings &settings,
                                                  const std::vector<bool> &res_mask) {
    std::vector<DiffractionSpot> output;
    StrongPixelSet strong_pixel_set;

    if (settings.signal_to_noise_threshold <= 0.0) {
        if (settings.photon_count_threshold > 0) {
            for (int line = 0; line < height; line++) {
                for (int col = 0; col < width; col++) {
                    int32_t pxl_val = input[line * width + col];
                    if (pxl_val > settings.photon_count_threshold && !pxl_mask[line * width + col] && !res_mask[line * width + col])
                        strong_pixel_set.AddStrongPixel(col, line, pxl_val);
                }
            }
            strong_pixel_set.FindSpotsImage(settings, output);
        }
        return output;
    }

    float strong2 = settings.signal_to_noise_threshold * settings.signal_to_noise_threshold;

    std::vector<uint8_t> mask(width * height, 0);

    // Sum and sum of squares of (2*NBY+1) vertical elements
    // These are updated after each line is finished
    // 64-bit integer guarantees calculations are made without rounding errors
    std::vector<int64_t> sum_vert(width, 0);
    std::vector<int64_t> sum2_vert(width, 0);
    std::vector<uint16_t> valid_vert(width, 0);

    for (int line = 0; line < NBX; line++) {
        for (int col = 0; col < width; col++) {
            auto pxl = line * width + col;
            mask[pxl] = (input[pxl] == special_value) || pxl_mask[pxl];

            if (mask[pxl] == 0) {
                int64_t tmp = input[pxl];
                sum_vert[col] += tmp;
                sum2_vert[col] += tmp * tmp;
                valid_vert[col] += 1;
            }
        }
    }

    for (int line = 0; line < height; line++) {
        for (int col = 0; col < width; col++) {
            if (line < height - NBX) {
                auto pxl = (line + NBX) * width + col;
                mask[pxl] = (input[pxl] == special_value) || pxl_mask[pxl];

                if (mask[pxl] == 0) {
                    int64_t tmp = input[pxl];
                    sum_vert[col] += tmp;
                    sum2_vert[col] += tmp * tmp;
                    valid_vert[col] += 1;
                }
            }

            if (line >= NBX + 1) {
                auto pxl = (line - (NBX + 1)) * width + col;
                if (mask[pxl] == 0) {
                    int64_t tmp = input[pxl];
                    sum_vert[col] -= tmp;
                    sum2_vert[col] -= tmp * tmp;
                    valid_vert[col] -= 1;
                }
            }
        }

        int64_t sum = 0;
        int64_t sum2 = 0;
        int64_t valid = 0;

        for (int col = 0; col < NBX; col++) {
            sum += sum_vert[col];
            sum2 += sum2_vert[col];
            valid += valid_vert[col];
        }

        for (int col = 0; col < width; col++) {
            if (col < width - NBX) {
                sum += sum_vert[col + NBX];
                sum2 += sum2_vert[col + NBX];
                valid += valid_vert[col + NBX];
            }

            if (col >= NBX + 1) {
                sum -= sum_vert[col - NBX - 1];
                sum2 -= sum2_vert[col - NBX - 1];
                valid -= valid_vert[col - NBX - 1];
            }

            T pxl_val = input[line * width + col];
            int64_t var = valid * sum2 - (sum * sum);
            int64_t in_minus_mean = pxl_val * valid - sum;

            if ((pxl_val > settings.photon_count_threshold) && // pixel is above count threshold
                (in_minus_mean > 0) && // pixel value is larger than mean
                (in_minus_mean * in_minus_mean > std::ceil(var * strong2)) && // pixel is above SNR threshold
                (mask[line * width + col] == 0) && // pixel is not bad pixel
                (res_mask[line * width + col] == 0)) // pixel is within a valid resolution range
                strong_pixel_set.AddStrongPixel(col, line, pxl_val);
        }
    }

    strong_pixel_set.FindSpotsImage(settings, output);
    return output;
}