// 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 "../image_analysis/bragg_integration/BraggIntegrate2D.h"

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) {
    // 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 uint8_t *image_ptr = in_image.GetUncompressedPtr(tmp);

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

    size_t good_pixel = 0;

    valid_pixel.reserve(image.size());

    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.GetMask().empty())
            mask_val = dataset->pixel_mask.GetMask().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_back(std::make_pair(val, i));
        }
    }
    // Sort based on the first element only
    std::sort(valid_pixel.begin(), valid_pixel.end(),
        [](const auto& a, const auto& b) {return a.first < b.first;});

    CalcAutoContrast();
}


void JFJochReaderImage::CalcAutoContrast() {
    if (valid_pixel.empty())
        auto_foreground = 10;
    else {
        auto it = valid_pixel.crbegin();
        const auto index = static_cast<size_t>(valid_pixel.size() * auto_foreground_range);
        std::advance(it, std::max<size_t>(1ULL, index) - 1);

        auto_foreground =  std::max(1, it->first);
    }
}

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 std::vector<std::pair<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();
    valid_pixel.reserve(image.size());
    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_back(std::make_pair(sum, i));
            }
        }
    }

    // Sort based on the first element only
    std::sort(valid_pixel.begin(), valid_pixel.end(),
        [](const auto& a, const auto& b) { return a.first < b.first; });

    CalcAutoContrast();
}

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