//
// Created by jungfrau on 4/22/26.
//

#include "ImagePreprocessor.h"

ImagePreprocessor::ImagePreprocessor(const DiffractionExperiment &experiment,
                                     const AzimuthalIntegration &integration,
                                     const PixelMask &mask,
                                     std::vector<int32_t> &processed_image)
    : npixels(experiment.GetPixelsNum()),
      experiment(experiment),
      integration(integration),
      azint_sum(integration.GetBinNumber(), 0.0),
      azint_sum2(integration.GetBinNumber(), 0.0),
      azint_count(integration.GetBinNumber(), 0),
      processed_image(processed_image),
      mask_1bit(npixels, false),
      azint_bins(integration.GetBinNumber()),
      saturation_limit(experiment.GetSaturationLimit()) {

    if (processed_image.size() != npixels)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Processed image size mismatch");

    for (int i = 0; i < npixels; i++)
        mask_1bit[i] = (mask.GetMask().at(i) != 0);
}

void ImagePreprocessor::Update(AzimuthalIntegrationProfile &profile) const {
    profile.Clear(integration);
    profile.Add(azint_sum, azint_count);
}

ImageStatistics ImagePreprocessor::Analyze(const uint8_t *image_ptr, CompressedImageMode image_mode) {
    switch (image_mode) {
        case CompressedImageMode::Int8:
            return Analyze<int8_t>(image_ptr, INT8_MIN, INT8_MAX);
        case CompressedImageMode::Int16:
            return Analyze<int16_t>(image_ptr, INT16_MIN, INT16_MAX);
        case CompressedImageMode::Int32:
            return Analyze<int32_t>(image_ptr, INT32_MIN, INT32_MAX);
        case CompressedImageMode::Uint8:
            return Analyze<uint8_t>(image_ptr, UINT8_MAX, UINT8_MAX);
        case CompressedImageMode::Uint16:
            return Analyze<uint16_t>(image_ptr, UINT16_MAX, UINT16_MAX);
        case CompressedImageMode::Uint32:
            return Analyze<uint32_t>(image_ptr, UINT32_MAX, UINT32_MAX);
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "RGB/float mode not supported");
    }
}

template<class T>
ImageStatistics ImagePreprocessor::Analyze(const uint8_t *input, T err_pixel_val, T sat_pixel_val) {
    auto image = reinterpret_cast<const T *>(input);

    for (int i = 0; i < azint_count.size(); i++) {
        azint_sum[i] = 0.0f;
        azint_sum2[i] = 0.0f;
        azint_count[i] = 0;
    }

    ImageStatistics ret{};

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

    auto &pixel_to_bin = integration.GetPixelToBin();
    auto &corrections = integration.Corrections();

    for (int i = 0; i < npixels; i++) {
        if (mask_1bit[i] != 0) {
            processed_image[i] = INT32_MIN;
            ++ret.masked_pixel_count;
        } else if (image[i] >= sat_pixel_val) {
            processed_image[i] = INT32_MAX;
            ++ret.saturated_pixel_count;
        } else if (std::is_signed<T>::value && (image[i] == err_pixel_val)) {
            // Error pixels are possible only for signed types
            processed_image[i] = INT32_MIN;
            ++ret.error_pixel_count;
        } else {
            processed_image[i] = static_cast<int32_t>(image[i]);

            if (image[i] > ret.max_value)
                ret.max_value = image[i];
            if (image[i] < ret.min_value)
                ret.min_value = image[i];

            const uint16_t bin = pixel_to_bin[i];
            if (bin < azint_bins) {
                float val = image[i] * corrections[i];
                azint_sum[bin] += val;
                ++azint_count[bin];
            }
        }
    }
    return ret;
}