// 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"

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

    message.image = CompressedImage{
            .data = reinterpret_cast<const uint8_t *>(image.data()),
            .size = image.size() * sizeof(int32_t),
            .xpixel = (size_t) in_dataset->experiment.GetXPixelsNum(),
            .ypixel = (size_t) in_dataset->experiment.GetYPixelsNum(),
            .pixel_depth_bytes = 4,
            .pixel_is_signed = true,
            .pixel_is_float = false,
            .algorithm = CompressionAlgorithm::NO_COMPRESSION,
            .channel = "default"
    };
}

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

    std::vector<uint8_t> tmp;

    const uint8_t *image_ptr;
    if (in_image.algorithm == CompressionAlgorithm::NO_COMPRESSION)
        image_ptr = in_image.data;
    else {
        tmp.resize(in_image.xpixel * in_image.ypixel * in_image.pixel_depth_bytes);

        JFJochDecompressPtr(tmp.data(),
                            in_image.algorithm,
                            in_image.data,
                            in_image.size,
                            in_image.xpixel * in_image.ypixel,
                            in_image.pixel_depth_bytes);
        image_ptr = tmp.data();
    }

    // For signed types there is no special value, only mask
    switch (in_image.pixel_depth_bytes) {
        case 1:
            if (in_image.pixel_is_signed)
                ProcessInputImage<int8_t>(image_ptr, npixel, INT8_MAX, INT8_MIN);
            else
                ProcessInputImage<uint8_t>(image_ptr, npixel, UINT8_MAX, INT64_MAX);
            break;
        case 2:
            if (in_image.pixel_is_signed)
                ProcessInputImage<int16_t>(image_ptr, npixel, INT16_MAX, INT16_MIN);
            else
                ProcessInputImage<uint16_t>(image_ptr, npixel, UINT16_MAX, INT64_MAX);
            break;
        case 4:
            if (in_image.pixel_is_signed)
                ProcessInputImage<int32_t>(image_ptr, npixel, INT32_MAX, INT32_MIN);
            else
                ProcessInputImage<uint32_t>(image_ptr, npixel, INT32_MAX, INT64_MAX);
            break;
    }
}

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);

    size_t good_pixel = 0;

    for (int 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 (!dataset->pixel_mask.empty())
            mask_val = dataset->pixel_mask.at(i);

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

const DataMessage &JFJochReaderImage::ImageData() const {
    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 std::map<int32_t, int32_t> &JFJochReaderImage::ValidPixels() const {
    return valid_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_pixel.clear();

    for (int 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] == ERROR_PXL_VALUE || other.image[i] == ERROR_PXL_VALUE) {
            image[i] = ERROR_PXL_VALUE;
            error_pixel.emplace(i);
        } else if (image[i] == SATURATED_PXL_VALUE || other.image[i] == SATURATED_PXL_VALUE) {
            image[i] = SATURATED_PXL_VALUE;
            saturated_pixel.emplace(i);
        } else {
            int64_t sum = static_cast<int64_t>(image[i]) + static_cast<int64_t>(other.image[i]);
            if (sum <= INT32_MIN + 5) [[unlikely]] {
                image[i] = ERROR_PXL_VALUE;
                error_pixel.emplace(i);
            } else if (sum > dataset->experiment.GetSaturationLimit()) [[unlikely]] {
                image[i] = SATURATED_PXL_VALUE;
                saturated_pixel.emplace(i);
            } else {
                image[i] = static_cast<int32_t>(sum);
                valid_pixel.emplace(sum, i);
            }
        }
    }
}

void JFJochReaderImage::CalcROI(const ROIElement *input) {
    if (!input) {
        roi = {};
        return;
    }
    int64_t width = dataset->experiment.GetXPixelsNum();
    int64_t height = dataset->experiment.GetYPixelsNum();

    int64_t roi_val = 0;
    uint64_t roi_val_2 = 0;
    int64_t roi_max = INT64_MIN;
    uint64_t roi_npixel = 0;

    for (int64_t y = 0; y < height; y++) {
        for (int64_t x = 0; x < width; x++) {
            int32_t val = image[x + width * y];
            if (input->CheckROI(x, y, 0.0) && (val != SATURATED_PXL_VALUE)
                && (val != ERROR_PXL_VALUE) && (val != GAP_PXL_VALUE)) {
                roi_val += val;
                roi_val_2 += val * val;
                if (val > roi_max) roi_max = val;
                roi_npixel++;
            }
        }
    }

    roi = ROIMessage{
            .sum = roi_val,
            .sum_square = roi_val_2,
            .max_count = roi_max,
            .pixels = roi_npixel
    };
}

const std::optional<ROIMessage> &JFJochReaderImage::GetROI() const {
    return roi;
}