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

#ifndef JUNGFRAUJOCH_CPUSPOTFINDER_H
#define JUNGFRAUJOCH_CPUSPOTFINDER_H

#include "../../common/DiffractionExperiment.h"
#include "SpotFindingSettings.h"
#include "StrongPixelSet.h"

#define NUM_PASS 1

template <class Tpixel = int16_t>
void FindSpots(DeviceOutput &output,
               int big_column, int row,
               int N, int LINES_TO_GO,
               const SpotFindingSettings& settings,
               const float *d_array,
               float *arr_mean,
               float *arr_stddev,
               uint32_t *arr_valid_count,
               uint32_t *arr_strong_pixel) {
    const auto image = reinterpret_cast<const Tpixel *>(output.pixels);

    std::vector<uint8_t> mask((2 * LINES_TO_GO + 1) * N, 0);

    for (int y = 0; y <= 2 * LINES_TO_GO; y++) {
        int line = row + y - LINES_TO_GO;
        if (line < 0 || line >= RAW_MODULE_LINES)
            continue;

        for (int x = 0; x < N; x++) {
            size_t col = big_column * N + x;
            size_t coord = line * RAW_MODULE_COLS + col;

            uint8_t bad_pixel = 0;
            if ((line == 255) || (line == 256)
                || (col == 255) || (col == 256)
                || (col == 511) || (col == 512)
                || (col == 767) || (col == 768))
                bad_pixel = 1;

            if ((d_array[coord] < settings.high_resolution_limit)
                || (d_array[coord] > settings.low_resolution_limit))
                bad_pixel = 1;

            if ((std::is_signed<Tpixel>::value
                 && (image[coord] == std::numeric_limits<Tpixel>::min()))
                || (image[coord] > std::numeric_limits<Tpixel>::max() - 10))
                bad_pixel = 1;

            mask[y * N + x] = bad_pixel;
        }
    }

    for (int i = 0; i < NUM_PASS; i++) {
        int64_t sum = 0;
        int64_t sum2 = 0;
        int64_t valid_count = 0;


        for (int y = 0; y <= 2 * LINES_TO_GO; y++) {
            int line = row + y - LINES_TO_GO;
            if (line < 0 || line >= RAW_MODULE_LINES)
                continue;

            for (int x = 0; x < N; x++) {
                size_t coord = line * RAW_MODULE_COLS + big_column * N + x;
                if (mask[y * N + x] == 0) {
                    sum += image[coord];
                    sum2 += image[coord] * image[coord];
                    valid_count++;
                }
            }
        }

        int64_t variance = valid_count * sum2 - sum * sum;
        float threshold = variance * settings.signal_to_noise_threshold * settings.signal_to_noise_threshold;

        for (int x = 0; x < N; x++) {
            size_t col = big_column * N + x;
            size_t coord = row * RAW_MODULE_COLS + col;

            arr_valid_count[coord] = valid_count;
            if (valid_count > 0) {
                arr_mean[coord] = static_cast<float>(sum) / static_cast<float>(valid_count);
                arr_stddev[coord] = sqrtf(
                        static_cast<float>(variance) / static_cast<float>(valid_count * valid_count));
            } else {
                arr_mean[coord] = -1;
                arr_stddev[coord] = -1;
            }

            bool strong_pixel = true;

            if (mask[LINES_TO_GO * N + x])
                strong_pixel = false;

            if ((settings.photon_count_threshold < 0)
                && (settings.signal_to_noise_threshold <= 0))
                strong_pixel = false;

            if ((settings.photon_count_threshold >= 0)
                && (image[coord] <= settings.photon_count_threshold)) {
                strong_pixel = false;
            }

            if (settings.signal_to_noise_threshold > 0) {
                int64_t in_minus_mean = image[coord] * valid_count - sum;
                if ((in_minus_mean * in_minus_mean <= threshold)
                    || (in_minus_mean <= 0)
                    || (valid_count <= N * N / 2))
                    strong_pixel = false;
            }

            if (strong_pixel) {
                mask[LINES_TO_GO * N + x] = 1;
                arr_strong_pixel[coord]++;
                output.spot_finding_result.strong_pixel[coord / 8] |= (1 << (coord % 8));
                output.spot_finding_result.strong_pixel_count++;
            }
        }
    }

}

template <class Tpixel = int16_t>
void FindSpots(DeviceOutput &output,
               const SpotFindingSettings& settings,
               const float *d_array,
               float *arr_mean,
               float *arr_stddev,
               uint32_t *arr_valid_count,
               uint32_t *arr_strong_pixel) {
    for (auto &i: output.spot_finding_result.strong_pixel)
        i = 0;

    int N = 32;
    int LINES_TO_GO = 15;

    for (int i = 0; i < RAW_MODULE_LINES; i++) {
        for (int j = 0; j < RAW_MODULE_COLS / N; j++)
            FindSpots<Tpixel>(output, j, i, N, LINES_TO_GO, settings, d_array, arr_mean, arr_stddev, arr_valid_count, arr_strong_pixel);
    }
}

#endif //JUNGFRAUJOCH_CPUSPOTFINDER_H