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

#include "BraggIntegrate2D.h"
#include "BraggPrediction.h"

template<class T>
bool IntegrateReflection(Reflection &r, const T *image, size_t xpixel, size_t ypixel,
                         int64_t special_value, int64_t saturation,
                         float r_3, float r_1_sq, float r_2_sq, float r_3_sq) {
    r.I = 0;
    r.bkg = 0;

    int64_t x0 = std::floor(r.predicted_x - r_3 - 1.0);
    int64_t x1 = std::ceil(r.predicted_x + r_3 + 1.0);
    int64_t y0 = std::floor(r.predicted_y - r_3 - 1.0);
    int64_t y1 = std::ceil(r.predicted_y + r_3 + 1.0);
    if (x0 < 0)
        x0 = 0;
    if (y0 < 0)
        y0 = 0;
    if (x1 >= xpixel)
        x1 = xpixel - 1;
    if (y1 >= ypixel)
        y1 = ypixel - 1;

    int64_t I_sum = 0;
    int64_t bkg_sum = 0;
    int64_t bkg_npixel = 0;
    int64_t I_npixel_inner = 0;
    int64_t I_npixel_integrated = 0;

    for (int64_t y = y0; y <= y1; y++) {
        for (int64_t x = x0; x <= x1; x++) {
            float dist_sq = (x - r.predicted_x) * (x - r.predicted_x) + (y - r.predicted_y) * (y - r.predicted_y);

            if (dist_sq < r_1_sq)
                I_npixel_inner++;

            if (image[y * xpixel + x] == special_value || image[y * xpixel + x] == saturation)
                continue;

            if (dist_sq < r_1_sq) {
                I_sum += image[y * xpixel + x];
                I_npixel_integrated++;
            } else if (dist_sq >= r_2_sq && dist_sq < r_3_sq) {
                bkg_sum += image[y * xpixel + x];
                bkg_npixel++;
            }
        }
    }

    if ((I_npixel_integrated == I_npixel_inner) && (bkg_npixel > 5)) {
        r.bkg = static_cast<float>(bkg_sum) / static_cast<float>(bkg_npixel);
        r.I = static_cast<float>(I_sum) - static_cast<float>(I_npixel_integrated) * r.bkg;
        // minimum sigma is 1!
        if (I_sum >= 1)
            r.sigma = std::sqrt(static_cast<float>(I_sum));
        else
            r.sigma = 1;
        return true;
    }
    return false;
}


template<class T>
std::vector<Reflection> IntegrateInternal(const DiffractionExperiment &experiment,
                                          const CompressedImage &image,
                                          const CrystalLattice &latt,
                                          BraggPrediction &prediction,
                                          float dist_from_ewald_sphere,
                                          int64_t special_value,
                                          int64_t saturation,
                                          int64_t image_number,
                                          char centering) {
    auto settings = experiment.GetBraggIntegrationSettings();
    auto geom = experiment.GetDiffractionGeometry();

    std::vector<uint8_t> buffer;
    auto ptr = reinterpret_cast<const T *>(image.GetUncompressedPtr(buffer));

    BraggPredictionSettings settings_prediction{
        .high_res_A = settings.GetDMinLimit_A(),
        .ewald_dist_cutoff = dist_from_ewald_sphere,
        .max_hkl = 100,
        .centering = centering
    };

    auto count = prediction.Calc(experiment, latt, settings_prediction);

    std::vector<Reflection> ret;
    float r_3 = settings.GetR3();
    float r_1_sq = settings.GetR1() * settings.GetR1();
    float r_2_sq = settings.GetR2() * settings.GetR2();
    float r_3_sq = settings.GetR3() * settings.GetR3();

    for (int i = 0; i < count; i++) {
        Reflection r = prediction.GetReflections().at(i);
        if (IntegrateReflection(r, ptr, image.GetWidth(), image.GetHeight(), special_value, saturation,
                                r_3, r_1_sq, r_2_sq, r_3_sq)) {
            if (experiment.GetPolarizationFactor()) {
                float pol = geom.CalcAzIntPolarizationCorr(r.predicted_x, r.predicted_y,
                                                           experiment.GetPolarizationFactor().value());
                r.I /= pol;
                r.bkg /= pol;
                r.sigma /= pol;
            }

            r.image_number = static_cast<int>(image_number);
            ret.push_back(r);
        }
    }

    return ret;
}

std::vector<Reflection> BraggIntegrate2D(const DiffractionExperiment &experiment,
                                         const CompressedImage &image,
                                         const CrystalLattice &latt,
                                         BraggPrediction &prediction,
                                         float dist_from_ewald_sphere,
                                         int64_t image_number,
                                         char centering) {
    if (image.GetCompressedSize() == 0)
        return {};

    switch (image.GetMode()) {
        case CompressedImageMode::Int8:
            return IntegrateInternal<int8_t>(experiment, image, latt, prediction, dist_from_ewald_sphere, INT8_MIN,
                                             INT8_MAX, image_number, centering);
        case CompressedImageMode::Int16:
            return IntegrateInternal<int16_t>(experiment, image, latt, prediction, dist_from_ewald_sphere, INT16_MIN,
                                              INT16_MAX, image_number, centering);
        case CompressedImageMode::Int32:
            return IntegrateInternal<int32_t>(experiment, image, latt, prediction, dist_from_ewald_sphere, INT32_MIN,
                                              INT32_MAX, image_number, centering);
        case CompressedImageMode::Uint8:
            return IntegrateInternal<uint8_t>(experiment, image, latt, prediction, dist_from_ewald_sphere, UINT8_MAX,
                                              UINT8_MAX, image_number, centering);
        case CompressedImageMode::Uint16:
            return IntegrateInternal<uint16_t>(experiment, image, latt, prediction, dist_from_ewald_sphere, UINT16_MAX,
                                               UINT16_MAX, image_number, centering);
        case CompressedImageMode::Uint32:
            return IntegrateInternal<uint16_t>(experiment, image, latt, prediction, dist_from_ewald_sphere, UINT32_MAX,
                                               UINT32_MAX, image_number, centering);
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Image mode not supported");
    }
}