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

#include "ImagePreprocessorCPU.h"

#include <algorithm>
#include <limits>
#include <type_traits>

ImagePreprocessorCPU::ImagePreprocessorCPU(const DiffractionExperiment &experiment, const PixelMask &mask)
    : ImagePreprocessor(experiment),
      mask_bits(mask.GetPackedMask()) {}

ImageStatistics ImagePreprocessorCPU::Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *image_ptr, CompressedImageMode image_mode) {
    if (processed_image.size() != npixels)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Processed image size mismatch");

    ImageStatistics ret{};
    AnalyzeBlock(processed_image, 0, image_ptr, npixels, image_mode, ret);
    return ret;
}

void ImagePreprocessorCPU::AnalyzeBlock(ImagePreprocessorBuffer &processed_image, size_t first, const uint8_t *input,
                                        size_t n, CompressedImageMode image_mode, ImageStatistics &stats) {
    switch (image_mode) {
        case CompressedImageMode::Int8:
            return AnalyzeBlock<int8_t>(processed_image, first, input, n, INT8_MIN, INT8_MAX, stats);
        case CompressedImageMode::Int16:
            return AnalyzeBlock<int16_t>(processed_image, first, input, n, INT16_MIN, INT16_MAX, stats);
        case CompressedImageMode::Int32:
            return AnalyzeBlock<int32_t>(processed_image, first, input, n, INT32_MIN, INT32_MAX, stats);
        case CompressedImageMode::Uint8:
            return AnalyzeBlock<uint8_t>(processed_image, first, input, n, UINT8_MAX, UINT8_MAX, stats);
        case CompressedImageMode::Uint16:
            return AnalyzeBlock<uint16_t>(processed_image, first, input, n, UINT16_MAX, UINT16_MAX, stats);
        case CompressedImageMode::Uint32:
            return AnalyzeBlock<uint32_t>(processed_image, first, input, n, UINT32_MAX, UINT32_MAX, stats);
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "RGB/float mode not supported");
    }
}

template<class T>
void ImagePreprocessorCPU::AnalyzeBlock(ImagePreprocessorBuffer &processed_image, size_t first, const uint8_t *input,
                                        size_t n, T err_pixel_val, T sat_pixel_val, ImageStatistics &ret) {
    auto image = reinterpret_cast<const T *>(input);
    int32_t *out = processed_image.data() + first;

    if (sat_pixel_val > saturation_limit)
        sat_pixel_val = static_cast<T>(saturation_limit);

    // Every pixel is classified without a branch, so that the loop over the pixels vectorises. The
    // extremes are taken over keys that order as the values do (a signed type's sign bit flipped), and a
    // pixel that does not count gives the key no maximum or minimum can be moved by; the counts are
    // integers. So the statistics are those of the pixel-by-pixel tests, whatever the order.
    using U = std::make_unsigned_t<T>;
    constexpr U flip = std::is_signed_v<T> ? static_cast<U>(U{1} << (8 * sizeof(T) - 1)) : U{0};
    uint64_t n_masked = 0, n_error = 0, n_saturated = 0;
    U key_max = 0, key_min = std::numeric_limits<U>::max();
    const auto pixel = [&](size_t i, bool masked) {
        const T v = image[i];
        // Error/invalid marker = the pixel type's extreme value (0xFFFFFFFF for EIGER uint32). Tested
        // before saturation, since for unsigned types the marker also exceeds sat_pixel_val (which is
        // clipped above to the HDF5 saturation_value).
        const bool error = !masked & (v == err_pixel_val);
        const bool saturated = !masked & !error & (v >= sat_pixel_val);
        const U counted = static_cast<U>(-static_cast<int>(!masked & !error & !saturated));   // all bits, or none
        out[i] = masked | error ? INT32_MIN : saturated ? INT32_MAX : static_cast<int32_t>(v);
        n_masked += masked;
        n_error += error;
        n_saturated += saturated;
        const U key = static_cast<U>(v) ^ flip;
        key_max = std::max<U>(key_max, key & counted);
        key_min = std::min<U>(key_min, key | static_cast<U>(~counted));
    };
    // Runs of whole mask words with no pixel masked - nearly all of a detector - take a loop of their
    // own, which the compiler vectorises; the rest is tested pixel by pixel.
    for (size_t i = 0; i < n;) {
        size_t end = i;
        if ((first + i) % 32 == 0)
            while (end + 32 <= n && mask_bits[(first + end) / 32] == 0)
                end += 32;
        if (end > i) {
            for (; i < end; i++)
                pixel(i, false);
        } else {
            pixel(i, (mask_bits[(first + i) / 32] >> ((first + i) % 32)) & 1U);
            i++;
        }
    }
    ret.masked_pixel_count += n_masked;
    ret.error_pixel_count += n_error;
    ret.saturated_pixel_count += n_saturated;
    if (n_masked + n_error + n_saturated < n) {   // some pixel counted
        ret.max_value = std::max<long long>(ret.max_value, static_cast<T>(key_max ^ flip));
        ret.min_value = std::min<long long>(ret.min_value, static_cast<T>(key_min ^ flip));
    }
}