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

#include <cmath>

#include "AzimuthalIntegration.h"
#include "JFJochException.h"
#include "DiffractionGeometry.h"
#include "RawToConvertedGeometry.h"

AzimuthalIntegration::AzimuthalIntegration(const AzimuthalIntegrationSettings& in_settings, size_t width, size_t height)
        : settings(in_settings), width(width), height(height) {

    if (width <= 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Detector width must be above 0");

    if (height <= 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Detector height must be above 0");

    if (settings.GetBinNumber() >= UINT16_MAX)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Cannot handle more than 65534 az. int. bins");
}

AzimuthalIntegration::AzimuthalIntegration(const DiffractionExperiment &experiment,
                                           const PixelMask& mask)
        : AzimuthalIntegration(experiment.GetAzimuthalIntegrationSettings(),
                               experiment.GetXPixelsNum(),
                               experiment.GetYPixelsNum()) {

    if (!experiment.IsGeometryTransformed())
        SetupRawGeom(experiment, mask.GetMaskRaw(experiment));
    else
        SetupConvGeom(experiment.GetDiffractionGeometry(),mask.GetMask());

    UpdateMaxBinNumber();
}

AzimuthalIntegration::AzimuthalIntegration(const DiffractionGeometry &geom,
                                           const AzimuthalIntegrationSettings &in_settings,
                                           const PixelMask& mask,
                                           size_t in_width,
                                           size_t in_height)
        : AzimuthalIntegration(in_settings, in_width, in_height) {

    SetupConvGeom(geom, mask.GetMask());
    UpdateMaxBinNumber();
}

void AzimuthalIntegration::SetupConvGeom(const DiffractionGeometry &geom,
                                         const std::vector<uint32_t> &mask) {
    pixel_to_bin.resize(width * height, UINT16_MAX);
    pixel_resolution.resize(width * height, 0);
    corrections.resize(width * height, 0);

    if (mask.size() != width * height)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mask size invalid");
    for (int row = 0; row < height; row++)
        for (int col = 0; col < width; col++)
            SetupPixel(geom, mask, row * width + col, col, row);
}

void AzimuthalIntegration::SetupRawGeom(const DiffractionExperiment &experiment,
                                        const std::vector<uint32_t> &mask) {
    if (mask.size() != RAW_MODULE_SIZE * experiment.GetModulesNum())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mask size invalid");

    pixel_to_bin.resize(RAW_MODULE_SIZE * experiment.GetModulesNum(), UINT16_MAX);
    pixel_resolution.resize(RAW_MODULE_SIZE * experiment.GetModulesNum(), 0);
    corrections.resize(RAW_MODULE_SIZE * experiment.GetModulesNum(), 0);

    auto geom = experiment.GetDiffractionGeometry();

    for (int m = 0; m < experiment.GetModulesNum(); m++) {
        for (int pxl = 0; pxl < RAW_MODULE_SIZE; pxl++) {
            auto [x,y] = RawToConvertedCoordinate(experiment, m, pxl);
            SetupPixel(geom, mask, m * RAW_MODULE_SIZE + pxl, x, y);
        }
    }
}

void AzimuthalIntegration::SetupPixel(const DiffractionGeometry &geom,
                                      const std::vector<uint32_t> &mask,
                                      uint32_t pxl, uint32_t col, uint32_t row) {
    if (mask[pxl] != 0)
        return;

    auto x = static_cast<float>(col);
    auto y = static_cast<float>(row);
    float d = geom.PxlToRes(x, y);
    pixel_resolution[pxl] = d;

    float corr = 1.0;
    if (settings.IsSolidAngleCorrection())
        corr /= geom.CalcAzIntSolidAngleCorr(x, y);
    if (settings.GetPolarizationFactor())
        corr /= geom.CalcAzIntPolarizationCorr(x, y, settings.GetPolarizationFactor().value());
    corrections[pxl] = corr;

    if (d > 0) {
        float q = 2.0f * M_PI / d;
        if ((q >= settings.GetLowQ_recipA())
            && (q < settings.GetHighQ_recipA())) {
            uint16_t bin = std::floor((q - settings.GetLowQ_recipA()) / settings.GetQSpacing_recipA());
            if (bin > max_bin_number)
                max_bin_number = bin;
            pixel_to_bin[pxl] = bin;
        }
    }
}

uint16_t AzimuthalIntegration::GetBinNumber() const {
    return max_bin_number + 1;
}

const std::vector<uint16_t> &AzimuthalIntegration::GetPixelToBin() const {
    return pixel_to_bin;
}

const std::vector<float> &AzimuthalIntegration::GetBinToQ() const {
    return bin_to_q;
}

uint16_t AzimuthalIntegration::QToBin(float q) const {
    return std::min(max_bin_number, settings.QToBin(q));
}

void AzimuthalIntegration::UpdateMaxBinNumber() {
    bin_to_q.resize(max_bin_number + 1);
    for (int i = 0; i < max_bin_number + 1; i++)
        bin_to_q[i] = static_cast<float>(settings.GetQSpacing_recipA() * (i + 0.5) + settings.GetLowQ_recipA());
}

const std::vector<float> &AzimuthalIntegration::Corrections() const {
    return corrections;
}

const std::vector<float> &AzimuthalIntegration::Resolution() const {
    return pixel_resolution;
}

const AzimuthalIntegrationSettings &AzimuthalIntegration::Settings() const {
    return settings;
}

size_t AzimuthalIntegration::GetWidth() const {
    return width;
}

size_t AzimuthalIntegration::GetHeight() const {
    return height;
}