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

#include <algorithm>
#include <cmath>

#include "AzimuthalIntegrationSettings.h"
#include "JFJochException.h"

#define check_max(param, val, max) if ((val) > (max)) throw JFJochException(JFJochExceptionCategory::InputParameterAboveMax, param)
#define check_min(param, val, min) if ((val) < (min)) throw JFJochException(JFJochExceptionCategory::InputParameterBelowMin, param)
#define check_finite(param, val) if (!std::isfinite(val)) throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, param)

AzimuthalIntegrationSettings::AzimuthalIntegrationSettings() {
    UpdateBinCount();
}

AzimuthalIntegrationSettings &AzimuthalIntegrationSettings::SolidAngleCorrection(bool input) {
    solid_angle_correction = input;
    return *this;
}

AzimuthalIntegrationSettings &AzimuthalIntegrationSettings::QRange_recipA(float low, std::optional<float> high) {
    check_finite("Low Q for azimuthal integration", low);
    check_min("Low Q for azimuthal integration", low, minQ_recipA);
    // The low limit has to leave room for at least one bin below maxQ. This used to follow from the
    // high limit being mandatory (high <= maxQ and high > low); once "unset" became legal, ResolveHighQ
    // was left computing std::clamp(q, low + spacing, maxQ) with the lower bound above the upper one.
    check_max("Low Q for azimuthal integration", low, maxQ_recipA - GetRequestedQSpacing_recipA().value_or(minQ_recipA));
    if (high.has_value()) {
        check_finite("High Q for azimuthal integration", *high);
        check_max("High Q for azimuthal integration", *high, maxQ_recipA);
        if (*high <= low)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "High Q must be higher than low Q");
    }

    requested_high_q_recipA = high;
    low_q_recipA = low;
    // Until ResolveHighQ runs, an unset limit keeps the value the bins were last built from.
    high_q_recipA = high.value_or(high_q_recipA);
    UpdateBinCount();
    return *this;
}

void AzimuthalIntegrationSettings::ResolveHighQ(float detector_max_q_recipA) {
    if (requested_high_q_recipA.has_value())
        return;

    high_q_recipA = std::clamp(detector_max_q_recipA, low_q_recipA + GetRequestedQSpacing_recipA().value_or(minQ_recipA),
                               maxQ_recipA);
    UpdateBinCount();
}

AzimuthalIntegrationSettings &AzimuthalIntegrationSettings::QSpacing_recipA(float input) {
    return QBinning(QBinSpacing::Linear, input, std::nullopt);
}

AzimuthalIntegrationSettings &AzimuthalIntegrationSettings::QBinning(QBinSpacing type, std::optional<float> spacing,
                                                                     std::optional<int32_t> q_bin_count) {
    if (spacing.has_value()) {
        check_finite("Q spacing for azimuthal integration", *spacing);
        check_min("Q spacing for azimuthal integration", *spacing, minQ_recipA);
    }
    if (q_bin_count.has_value()) {
        check_min("Number of Q bins for azimuthal integration", *q_bin_count, 1);
        check_max("Number of Q bins for azimuthal integration", *q_bin_count, UINT16_MAX - 1);
    }
    if (type == QBinSpacing::Linear && spacing.has_value() && q_bin_count.has_value())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Linear Q binning takes either the Q spacing or the number of Q bins, not both");
    if (type != QBinSpacing::Linear && spacing.has_value())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Q spacing applies only to linear Q binning; give the number of Q bins instead");
    if (type != QBinSpacing::Linear && !q_bin_count.has_value())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Logarithmic Q binning needs the number of Q bins");

    q_bin_spacing = type;
    q_spacing = spacing.value_or(q_spacing);
    requested_q_bins = q_bin_count;
    UpdateBinCount();
    return *this;
}

bool AzimuthalIntegrationSettings::IsSolidAngleCorrection() const {
    return solid_angle_correction;
}

float AzimuthalIntegrationSettings::GetHighQ_recipA() const {
    return high_q_recipA;
}

std::optional<float> AzimuthalIntegrationSettings::GetRequestedHighQ_recipA() const {
    return requested_high_q_recipA;
}

float AzimuthalIntegrationSettings::GetLowQ_recipA() const {
    return low_q_recipA;
}

float AzimuthalIntegrationSettings::GetQSpacing_recipA() const {
    return q_bin_spacing == QBinSpacing::Linear ? bin_step : q_spacing;
}

QBinSpacing AzimuthalIntegrationSettings::GetQBinSpacing() const {
    return q_bin_spacing;
}

std::optional<float> AzimuthalIntegrationSettings::GetRequestedQSpacing_recipA() const {
    if (q_bin_spacing == QBinSpacing::Linear && !requested_q_bins.has_value())
        return q_spacing;
    return std::nullopt;
}

std::optional<int32_t> AzimuthalIntegrationSettings::GetRequestedQBinCount() const {
    return requested_q_bins;
}

std::vector<double> AzimuthalIntegrationSettings::GetQBinCentres_recipA() const {
    std::vector<double> ret(q_bins);
    for (int32_t i = 0; i < q_bins; i++) {
        if (q_bin_spacing == QBinSpacing::Linear)
            ret[i] = bin_step * (i + 0.5) + low_q_recipA;
        else
            ret[i] = std::pow(10.0, std::log10(static_cast<double>(low_q_recipA)) + bin_step * (i + 0.5));
    }
    return ret;
}

AzimuthalIntegrationSettings &AzimuthalIntegrationSettings::BkgEstimateQRange_recipA(float low, float high) {
    check_finite("Low Q for background estimation", low);
    check_finite("High Q for background estimation", high);
    check_max("High Q for background estimation", high, maxQ_recipA);
    check_min("Low Q for background estimation", low, minQ_recipA);
    if (high <= low)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "High Q must be higher than low Q");

    bkg_estimate_low_q_recipA = low;
    bkg_estimate_high_q_recipA = high;

    return *this;
}

float AzimuthalIntegrationSettings::GetBkgEstimateLowQ_recipA() const {
    return bkg_estimate_low_q_recipA;
}

float AzimuthalIntegrationSettings::GetBkgEstimateHighQ_recipA() const {
    return bkg_estimate_high_q_recipA;
}

AzimuthalIntegrationSettings &AzimuthalIntegrationSettings::PolarizationCorrection(bool input) {
    polarization_correction = input;
    return *this;
}

bool AzimuthalIntegrationSettings::IsPolarizationCorrection() const {
    return polarization_correction;
}

AzimuthalIntegrationSettings &AzimuthalIntegrationSettings::AzimuthalBinCount(int32_t input) {
    check_min("Azimuthal bin count", input, 1);
    check_max("Azimuthal bin count", input, 512);
    azim_bins = input;
    UpdateBinCount();
    return *this;
}

AzimuthalIntegrationSettings &AzimuthalIntegrationSettings::ForceCPUinFPGAWorkflow(bool input) {
    force_cpu_in_fpga_workflow = input;
    return *this;
}

bool AzimuthalIntegrationSettings::IsForceCPUinFPGAWorkflow() const {
    return force_cpu_in_fpga_workflow;
}

int32_t AzimuthalIntegrationSettings::GetQBinCount() const {
    return q_bins;
}

int32_t AzimuthalIntegrationSettings::GetAzimuthalBinCount() const {
    return azim_bins;
}

int32_t AzimuthalIntegrationSettings::GetBinCount() const {
    return total_bins;
}

float AzimuthalIntegrationSettings::BinCoordinate(float q) const {
    if (q_bin_spacing == QBinSpacing::Linear)
        return (q - low_q_recipA) / bin_step;
    return (std::log10(q) - std::log10(low_q_recipA)) / bin_step;
}

uint16_t AzimuthalIntegrationSettings::QToBin(float q) const {
    return std::min<uint16_t>(GetBinCount() - 1, std::floor(std::max(0.0f, BinCoordinate(q))));
}

uint16_t AzimuthalIntegrationSettings::GetBin(float q, float phi_deg) const {
    if (q < low_q_recipA || q >= high_q_recipA)
        return UINT16_MAX;
    if (phi_deg < 0.0 || phi_deg >= 360)
        return UINT16_MAX;

    // Rounding can put a q just below high Q on the far edge of the last bin.
    int32_t q_bin = std::min<int32_t>(std::floor(BinCoordinate(q)), q_bins - 1);
    int32_t phi_bin = std::floor(phi_deg / 360.0f * azim_bins);
    return q_bin + phi_bin * GetQBinCount();
}

void AzimuthalIntegrationSettings::UpdateBinCount() {
    if (!requested_q_bins.has_value()) {
        bin_step = q_spacing;
        q_bins = std::ceil((high_q_recipA - low_q_recipA) / q_spacing);
    } else {
        q_bins = requested_q_bins.value();
        if (q_bin_spacing == QBinSpacing::Linear)
            bin_step = (high_q_recipA - low_q_recipA) / static_cast<float>(q_bins);
        else
            bin_step = (std::log10(high_q_recipA) - std::log10(low_q_recipA)) / static_cast<float>(q_bins);
    }
    total_bins = q_bins * azim_bins;
}