// 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,
                                           float wavelength)
        : settings(in_settings), width(width), height(height), wavelength(wavelength) {
    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.GetBinCount() >= 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(),
                               experiment.GetWavelength_A()) {

    polarization_correction = experiment.GetPolarizationFactor();

    if (!experiment.IsGeometryTransformed())
        SetupRawGeom(experiment, mask.GetMaskRaw(experiment));
    else
        SetupConvGeom(experiment.GetDiffractionGeometry(),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);
    float phi_rad = geom.Phi_rad(x, y);

    pixel_resolution[pxl] = d;

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

    if (d > 0) {
        float q = 2.0f * static_cast<float>(M_PI) / d;
        pixel_to_bin[pxl] = settings.GetBin(q, phi_rad * 180.0 / M_PI);
    }
}

uint16_t AzimuthalIntegration::GetBinNumber() const {
    return settings.GetBinCount();
}

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

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

const std::vector<float> &AzimuthalIntegration::GetBinToD() const {
    return bin_to_d;
}

const std::vector<float> &AzimuthalIntegration::GetBinToTwoTheta() const {
    return bin_to_2theta;
}

const std::vector<float> &AzimuthalIntegration::GetBinToPhi() const {
    return bin_to_phi;
}

uint16_t AzimuthalIntegration::QToBin(float q) const {
    return settings.QToBin(q);
}

void AzimuthalIntegration::UpdateMaxBinNumber() {
    bin_to_q.resize(settings.GetBinCount());
    bin_to_d.resize(settings.GetBinCount());
    bin_to_2theta.resize(settings.GetBinCount());
    bin_to_phi.resize(settings.GetBinCount());

    for (int j = 0; j < settings.GetAzimuthalBinCount(); j++) {
        for (int i = 0; i < settings.GetQBinCount(); i++) {
            bin_to_q[j * settings.GetQBinCount() + i] = static_cast<float>(settings.GetQSpacing_recipA() * (i + 0.5) + settings.GetLowQ_recipA());
            bin_to_d[j * settings.GetQBinCount() + i] = 2.0f * static_cast<float>(M_PI) / bin_to_q[j * settings.GetQBinCount() + i];
            bin_to_2theta[j * settings.GetQBinCount() + i] = 2.0f * asinf(bin_to_q[i] * wavelength / (4.0f * static_cast<float>(M_PI))) * 180.0f /
                               static_cast<float>(M_PI);
            bin_to_phi[j * settings.GetQBinCount() + i] = static_cast<float>(j) * 360.0f / static_cast<float>(settings.GetAzimuthalBinCount());
        }
    }
}

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

int32_t AzimuthalIntegration::GetAzimuthalBinCount() const {
    return settings.GetAzimuthalBinCount();
}

int32_t AzimuthalIntegration::GetQBinCount() const {
    return settings.GetQBinCount();
}