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

#pragma once

#include <limits>
#include <type_traits>

#include "ROIIntegration.h"
#include "../../common/JFJochException.h"

class ROIIntegrationCPU : public ROIIntegration {
public:
    explicit ROIIntegrationCPU(const DiffractionExperiment &experiment);

    // image is anything indexable with operator[] and size(). Templated on the
    // pixel type so a future narrow-integer (e.g. 16-bit) path works as well.
    template <class T>
    void RunROI(const T &image, std::map<std::string, ROIMessage> &out) {
        if (image.size() != npixel)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "ROIIntegration: mismatch in image size");

        for (uint16_t r = 0; r < roi_count; r++) {
            roi_sum[r] = 0;
            roi_sum2[r] = 0;
            roi_pixels[r] = 0;
            roi_x_weighted[r] = 0;
            roi_y_weighted[r] = 0;
            roi_max[r] = INT64_MIN;
        }

        using pixel_t = std::remove_cv_t<std::remove_reference_t<decltype(image[0])>>;

        for (size_t i = 0; i < npixel; i++) {
            const uint16_t mask = roi_map[i];
            if (mask == 0)
                continue;

            const pixel_t v = image[i];
            // masked/bad pixels (signed types only) are excluded entirely
            if constexpr (std::is_signed_v<pixel_t>) {
                if (v == std::numeric_limits<pixel_t>::min())
                    continue;
            }
            // saturated pixels still count towards the max, but not the sums
            const bool saturated = (v == std::numeric_limits<pixel_t>::max());

            const int64_t val = static_cast<int64_t>(v);
            const int64_t x = static_cast<int64_t>(i % width);
            const int64_t y = static_cast<int64_t>(i / width);

            for (uint16_t r = 0; r < roi_count; r++) {
                if (!(mask & (1u << r)))
                    continue;
                if (!saturated) {
                    roi_sum[r] += val;
                    roi_sum2[r] += static_cast<uint64_t>(val * val);
                    roi_pixels[r] += 1;
                    roi_x_weighted[r] += val * x;
                    roi_y_weighted[r] += val * y;
                }
                if (val > roi_max[r])
                    roi_max[r] = val;
            }
        }
        Export(out);
    }

    void Run(const ImagePreprocessorBuffer &image, std::map<std::string, ROIMessage> &out) override;
};