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

#include <cmath>
#include <thread>
#include <future>

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

AzimuthalIntegrationMapping::AzimuthalIntegrationMapping(const DiffractionExperiment &experiment,
                                           const PixelMask& mask,
                                           size_t in_nthreads)
    : settings(experiment.GetAzimuthalIntegrationSettings()),
      wavelength(experiment.GetWavelength_A()),
      width(experiment.GetXPixelsNumConv()),
      height(experiment.GetYPixelsNumConv()) {

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

    polarization_factor = experiment.GetPolarizationFactor();

    if (in_nthreads == 0)
        nthreads = std::thread::hardware_concurrency();
    else
        nthreads = in_nthreads;

    nthreads = std::clamp<size_t>(nthreads, 1, 64);

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

    UpdateMaxBinNumber();
}

void AzimuthalIntegrationMapping::SetupConvGeomRows(const DiffractionGeometry &geom, const std::vector<uint32_t> &mask,
    size_t row0, size_t row_end) {
    for (size_t row = row0; row < row_end && row < height; row++) {
        for (size_t col = 0; col < width; col++)
            SetupPixel(geom, mask, row * width + col, col, row);
    }
}

void AzimuthalIntegrationMapping::SetupConvGeom(const DiffractionGeometry &geom,
                                                const std::vector<uint32_t> &mask,
                                                size_t nthreads) {
    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");

    if (nthreads <= 1) {
        SetupConvGeomRows(geom, mask, 0, height);
    } else {
        nthreads = std::min(nthreads, height);
        std::vector<std::future<void>> futures;

        for (size_t t = 0; t < nthreads; ++t)
            futures.emplace_back(std::async(std::launch::async,
                &AzimuthalIntegrationMapping::SetupConvGeomRows,
                this, std::cref(geom), std::cref(mask), t * height / nthreads, (t + 1) * height / nthreads));

        for (auto &f: futures)
            f.get();
    }
}

void AzimuthalIntegrationMapping::SetupRawGeom(const DiffractionExperiment &experiment,
                                               const std::vector<uint32_t> &mask, size_t nthreads) {
    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();

    if (nthreads <= 1) {
        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);
            }
        }
    } else {
        nthreads = std::min<size_t>(nthreads, experiment.GetModulesNum());
        std::vector<std::future<void>> futures;
        futures.reserve(nthreads);

        for (size_t t = 0; t < nthreads; ++t) {
            const size_t module_begin = t * experiment.GetModulesNum() / nthreads;
            const size_t module_end = (t + 1) * experiment.GetModulesNum() / nthreads;

            futures.emplace_back(std::async(std::launch::async, [&, module_begin, module_end] {
                for (size_t m = module_begin; m < module_end; ++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);
                    }
                }
            }));
        }

        for (auto &f: futures)
            f.get();
    }
}

void AzimuthalIntegrationMapping::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_factor)
        corr /= geom.CalcAzIntPolarizationCorr(x, y, polarization_factor.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 AzimuthalIntegrationMapping::GetBinNumber() const {
    return settings.GetBinCount();
}

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

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

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

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

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

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

void AzimuthalIntegrationMapping::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> &AzimuthalIntegrationMapping::Corrections() const {
    return corrections;
}

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

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

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

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

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

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

size_t AzimuthalIntegrationMapping::GetNThreads() const {
    return nthreads;
}