Jungfraujoch/reader/JFJochReaderImage.cpp

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

#include "JFJochReaderImage.h"
#include "../common/PixelMask.h"
#include "JFJochDecompress.h"

#include <queue>
#include <algorithm>
#include <cmath>

JFJochReaderImage::JFJochReaderImage(const DataMessage &in_message,
                                     const std::shared_ptr<const JFJochReaderDataset> &in_dataset)
        : message(in_message),
          dataset(in_dataset),
          image(in_dataset->experiment.GetPixelsNum(), 0) {
    ProcessInputImage(in_message.image);

    message.image = CompressedImage(image, in_dataset->experiment.GetXPixelsNum(),
                                    in_dataset->experiment.GetYPixelsNum());
}

JFJochReaderImage::JFJochReaderImage(const JFJochReaderImage &other)
    : dataset(other.dataset),
      image(other.image),
      message(other.message),
      saturated_pixel(other.saturated_pixel),
      error_pixel(other.error_pixel),
      valid_pixel(other.valid_pixel),
      valid_min(other.valid_min),
      valid_max(other.valid_max),
      valid_count(other.valid_count),
      has_valid(other.has_valid),
      top_pixels_acc(other.top_pixels_acc),
      top_pixels(other.top_pixels),
      count_histogram(other.count_histogram),
      auto_foreground(other.auto_foreground) {
    // Need to make image use local copy
    message.image = CompressedImage(image, dataset->experiment.GetXPixelsNum(),
                                    dataset->experiment.GetYPixelsNum());
}

void JFJochReaderImage::ProcessInputImage(const CompressedImage &in_image) {
    size_t npixel = in_image.GetWidth() * in_image.GetHeight();
    if (npixel == 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Image size cannot be zero");

    std::vector<uint8_t> tmp;

    const auto compression_start_time = std::chrono::steady_clock::now();
    const uint8_t *image_ptr = in_image.GetUncompressedPtr(tmp);
    const auto compression_end_time = std::chrono::steady_clock::now();

    if (in_image.GetCompressionAlgorithm() != CompressionAlgorithm::NO_COMPRESSION)
        message.compression_time_s = std::chrono::duration<float>(compression_end_time - compression_start_time).count();

    switch (in_image.GetMode()) {
        case CompressedImageMode::Int8:
            ProcessInputImage<int8_t>(image_ptr, npixel, INT8_MAX, INT8_MIN);
            break;
        case CompressedImageMode::Int16:
            ProcessInputImage<int16_t>(image_ptr, npixel, INT16_MAX, INT16_MIN);
            break;
        case CompressedImageMode::Int32:
            ProcessInputImage<int32_t>(image_ptr, npixel, INT32_MAX, INT32_MIN);
            break;
        case CompressedImageMode::Uint8:
            ProcessInputImage<uint8_t>(image_ptr, npixel, UINT8_MAX, INT64_MAX);
            break;
        case CompressedImageMode::Uint16:
            ProcessInputImage<uint16_t>(image_ptr, npixel, UINT16_MAX, INT64_MAX);

            break;
        case CompressedImageMode::Uint32:
            ProcessInputImage<uint32_t>(image_ptr, npixel, INT32_MAX, INT64_MAX);
            break;
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Floating point images not supported");
    }
}

template<class T>
void JFJochReaderImage::ProcessInputImage(const void *data, size_t npixel, int64_t sat_value, int64_t special_value) {
    if (npixel != image.size())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Mismatch in input size");

    const T* img_ptr = reinterpret_cast<const T*>(data);

    // Reset per-image stats
    saturated_pixel.clear();
    error_pixel.clear();
    valid_count = 0;
    has_valid = false;

    top_pixels_acc.Clear();
    top_pixels.clear();
    top_pixels.reserve(top_pixels_acc.Capacity());

    bool has_input_mask = false;
    const auto &mask = dataset->pixel_mask->GetMask();
    if (mask.size() == npixel)
        has_input_mask = true;

    for (size_t i = 0; i < npixel; i++) {
        int32_t val;
        if (img_ptr[i] <= INT32_MAX)
            val = static_cast<int32_t>(img_ptr[i]);
        else
            val = INT32_MAX;

        uint32_t mask_val = 0;
        if (has_input_mask)
            mask_val = mask[i];

        if ((mask_val & (
                (1<<PixelMask::ModuleGapPixelBit)
                | (1<<PixelMask::ChipGapPixelBit)
                | (1<<PixelMask::ModuleEdgePixelBit))) != 0) {
            image[i] = GAP_PXL_VALUE;
        } else if ((mask_val & (1u << PixelMask::BeamStopPixelBit)) != 0) {
            image[i] = BEAM_STOP_PXL_VALUE;
        } else if ((mask_val != 0) || (img_ptr[i] == special_value)) {
            image[i] = ERROR_PXL_VALUE;
            error_pixel.emplace(static_cast<int64_t>(i));
        } else if (val >= sat_value) {
            image[i] = SATURATED_PXL_VALUE;
            saturated_pixel.emplace(static_cast<int64_t>(i));
        } else {
            image[i] = val;

            if (!has_valid) {
                has_valid = true;
                valid_min = val;
                valid_max = val;
            } else {
                valid_min = std::min(valid_min, val);
                valid_max = std::max(valid_max, val);
            }
            valid_count++;
            count_histogram.Add(val);
            top_pixels_acc.Add(val, static_cast<int32_t>(i));
        }
    }

    auto_foreground = count_histogram.Percentile(auto_foreground_range).value_or(10);

    // Export top pixels (already sorted descending) into the existing vector interface
    for (int i = 0; i < top_pixels_acc.Size(); i++) {
        const auto &e = top_pixels_acc[i];
        top_pixels.emplace_back(e.value, e.index);
    }
}

std::optional<std::pair<int32_t, int32_t>> JFJochReaderImage::ValidMinMax() const {
    if (!has_valid)
        return {};
    return std::make_pair(valid_min, valid_max);
}

const std::vector<std::pair<int32_t, int32_t>> &JFJochReaderImage::GetTopPixels() const {
    return top_pixels;
}

const DataMessage &JFJochReaderImage::ImageData() const {
    return message;
}

DataMessage &JFJochReaderImage::ImageData() {
    return message;
}

const std::vector<int32_t> &JFJochReaderImage::Image() const {
    return image;
}

const std::unordered_set<int64_t> &JFJochReaderImage::SaturatedPixels() const {
    return saturated_pixel;
}

const std::unordered_set<int64_t> &JFJochReaderImage::ErrorPixels() const {
    return error_pixel;
}

const JFJochReaderDataset &JFJochReaderImage::Dataset() const {
    if (!dataset)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Dataset not set");
    return *dataset;
}

void JFJochReaderImage::AddImage(const JFJochReaderImage &other) {
    if (other.Image().size() != image.size())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Mismatch in size");

    message.indexing_result = false;
    message.resolution_estimate = {};
    message.bkg_estimate = {};
    message.spots = {};

    error_pixel.clear();
    saturated_pixel.clear();
    valid_count = 0;
    has_valid = false;

    top_pixels_acc.Clear();
    top_pixels.clear();
    top_pixels.reserve(top_pixels_acc.Capacity());
    count_histogram.clear();

    for (size_t i = 0; i < image.size(); i++) {
        if (image[i] == GAP_PXL_VALUE || other.image[i] == GAP_PXL_VALUE) {
            image[i] = GAP_PXL_VALUE;
        } else if (image[i] == BEAM_STOP_PXL_VALUE || other.image[i] == BEAM_STOP_PXL_VALUE) {
            image[i] = BEAM_STOP_PXL_VALUE;
        } else if (image[i] == ERROR_PXL_VALUE || other.image[i] == ERROR_PXL_VALUE) {
            image[i] = ERROR_PXL_VALUE;
            error_pixel.emplace(static_cast<int64_t>(i));
        } else if (image[i] == SATURATED_PXL_VALUE || other.image[i] == SATURATED_PXL_VALUE) {
            image[i] = SATURATED_PXL_VALUE;
            saturated_pixel.emplace(static_cast<int64_t>(i));
        } else {
            int64_t sum = static_cast<int64_t>(image[i]) + static_cast<int64_t>(other.image[i]);
            if (sum < MIN_REAL_PXL_VALUE) [[unlikely]] {
                image[i] = ERROR_PXL_VALUE;
                error_pixel.emplace(static_cast<int64_t>(i));
            } else if (sum >= dataset->experiment.GetSaturationLimit()) [[unlikely]] {
                image[i] = SATURATED_PXL_VALUE;
                saturated_pixel.emplace(static_cast<int64_t>(i));
            } else {
                const int32_t val = static_cast<int32_t>(sum);
                image[i] = val;

                if (!has_valid) {
                    has_valid = true;
                    valid_min = val;
                    valid_max = val;
                } else {
                    valid_min = std::min(valid_min, val);
                    valid_max = std::max(valid_max, val);
                }
                valid_count++;
                count_histogram.Add(val);
                top_pixels_acc.Add(val, static_cast<int32_t>(i));
            }
        }
    }

    auto_foreground = count_histogram.Percentile(auto_foreground_range).value_or(10);

    for (int i = 0; i < top_pixels_acc.Size(); i++) {
        const auto &e = top_pixels_acc[i];
        top_pixels.emplace_back(e.value, e.index);
    }
}

std::vector<float> JFJochReaderImage::GetAzInt1D() const {
    if (dataset->azimuthal_bins <= 1) {
        return message.az_int_profile;
    } else if (message.az_int_profile.size() == dataset->azimuthal_bins * dataset->q_bins
               && dataset->azimuthal_bins * dataset->q_bins > 0 ) {
        std::vector<float> tmp(dataset->q_bins);
        for (int i = 0; i < message.az_int_profile.size(); i++)
            tmp[i % dataset->q_bins] += message.az_int_profile[i];
        return tmp;
    } else
        return {};
}

std::vector<float> JFJochReaderImage::GetAzInt1D_BinToQ() const {
    if (dataset->azimuthal_bins <= 1) {
        return dataset->az_int_bin_to_q;
    } else if (dataset->az_int_bin_to_q.size() == dataset->azimuthal_bins * dataset->q_bins
               && dataset->azimuthal_bins * dataset->q_bins > 0 ) {
        std::vector<float> tmp(dataset->q_bins);
        for (int i = 0; i < dataset->q_bins; i++)
            tmp[i] = dataset->az_int_bin_to_q[i];
        return tmp;
    } else
        return {};
}

std::shared_ptr<JFJochReaderDataset> JFJochReaderImage::CreateMutableDataset() {
    std::shared_ptr<JFJochReaderDataset> new_dataset = std::make_shared<JFJochReaderDataset>(*dataset);
    dataset = new_dataset;
    return new_dataset;
}

int32_t JFJochReaderImage::GetAutoContrastValue() const {
   return auto_foreground;
}

std::vector<float> JFJochReaderImage::GetHistogram() const {
    return count_histogram.GetCount();
}