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

#include "AzimuthalIntegrationProfile.h"
#include "JFJochException.h"
#include "Definitions.h"

#include <algorithm>

inline float sum_to_count(float sum, uint64_t count) {
    if (count == 0)
        return NAN;
    return sum / (static_cast<float>(count));
}

inline float calc_std(float sum, float sum2, uint64_t count) {
    if (count == 0 || count == 1)
        return NAN;
    const auto fp_count = static_cast<float>(count);
    const float variance = (sum2 - sum * sum / fp_count) / (fp_count - 1);
    return std::sqrt(variance);
}

AzimuthalIntegrationProfile::AzimuthalIntegrationProfile(const AzimuthalIntegrationMapping &mapping)
    : sum(mapping.GetBinNumber(), 0),
      sum2(mapping.GetBinNumber(), 0),
      count(mapping.GetBinNumber(), 0),
      bin_to_q(mapping.GetBinToQ()),
      bin_to_d(mapping.GetBinToD()),
      bin_to_2theta(mapping.GetBinToTwoTheta()),
      bin_to_phi(mapping.GetBinToPhi()),
      q_bins(mapping.GetQBinCount()),
      azim_bins(mapping.GetAzimuthalBinCount()) {
}

void AzimuthalIntegrationProfile::Clear(const AzimuthalIntegrationMapping &mapping) {
    std::unique_lock ul(m);

    bin_to_d = mapping.GetBinToD();
    bin_to_q = mapping.GetBinToQ();
    bin_to_2theta = mapping.GetBinToTwoTheta();
    bin_to_phi = mapping.GetBinToPhi();
    q_bins = mapping.GetQBinCount();
    azim_bins = mapping.GetAzimuthalBinCount();

    sum = std::vector<float>(mapping.GetBinNumber(), 0);
    count = std::vector<uint64_t>(mapping.GetBinNumber(), 0);
}

void AzimuthalIntegrationProfile::Add(int64_t bin, int64_t value) {
    if (bin < 0 || bin >= sum.size())
        return;
    std::unique_lock ul(m);
    sum[bin] += static_cast<float>(value);
    sum2[bin] += static_cast<float>(value * value);
    count[bin]++;
}

void AzimuthalIntegrationProfile::Add(const std::vector<float> &in_sum,
                                      const std::vector<float> &in_sum2,
                                      const std::vector<uint32_t> &in_count) {
    std::unique_lock ul(m);
    if ((in_sum.size() == sum.size()) && (in_count.size() == count.size())) {
        for (int i = 0; i < sum.size(); i++) {
            sum[i] += in_sum[i];
            count[i] += in_count[i];
        }
        if (in_sum2.size() == sum2.size()) {
            for (int i = 0; i < sum.size(); i++) {
                sum2[i] += in_sum2[i];
            }
        }
    } else if (!in_sum.empty() && !in_count.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mismatch in size of sum/count datasets");
}

std::vector<float> AzimuthalIntegrationProfile::GetResult() const {
    std::unique_lock ul(m);

    std::vector<float> rad_int_profile(sum.size(), 0);
    for (int i = 0; i < sum.size(); i++)
        rad_int_profile[i] = sum_to_count(sum[i], count[i]);
    return rad_int_profile;
}

std::vector<float> AzimuthalIntegrationProfile::GetStd() const {
    std::unique_lock ul(m);

    std::vector<float> rad_int_profile(sum.size(), 0);
    for (int i = 0; i < sum.size(); i++)
        rad_int_profile[i] = calc_std(sum[i], sum2[i], count[i]);
    return rad_int_profile;
}

std::vector<uint64_t> AzimuthalIntegrationProfile::GetPixelCount() const {
    std::unique_lock ul(m);
    return count;
}

std::vector<float> AzimuthalIntegrationProfile::GetResult1D() const {
    std::unique_lock ul(m);

    std::vector<float> sum_q(q_bins, 0.0f);
    std::vector<uint64_t> count_q(q_bins, 0);

    for (int i = 0; i < sum.size(); i++) {
        const int q_bin = i % q_bins;
        sum_q[q_bin] += sum[i];
        count_q[q_bin] += count[i];
    }

    std::vector<float> rad_int_profile(q_bins, 0.0f);
    for (int q = 0; q < q_bins; q++)
        rad_int_profile[q] = sum_to_count(sum_q[q], count_q[q]);

    return rad_int_profile;
}
void AzimuthalIntegrationProfile::SetTitle(const std::string &input) {
    title = input;
}

const std::vector<float> &AzimuthalIntegrationProfile::GetXAxis(PlotAzintUnit unit) const {
    switch (unit) {
        case PlotAzintUnit::TwoTheta_deg:
            return bin_to_2theta;
        case PlotAzintUnit::d_A:
            return bin_to_d;
        default:
        case PlotAzintUnit::Q_recipA:
            return bin_to_q;
    }
}

MultiLinePlot AzimuthalIntegrationProfile::GetPlot(bool force_1d, PlotAzintUnit plot_unit) const {
    MultiLinePlot ret;

    const std::vector<float> &x_coord = GetXAxis(plot_unit);

    if (azim_bins == 1)
        ret.AddPlot(MultiLinePlotStruct{.title = title, .x = x_coord, .y = GetResult()});
    else {
        if (force_1d) {
            std::vector<float> x_shortened(q_bins);
            for (int i = 0; i < q_bins; i++)
                x_shortened[i] = x_coord[i];
            ret.AddPlot(MultiLinePlotStruct{.title = title, .x = x_shortened, .y = GetResult1D()});
        } else {
            ret.AddPlot(MultiLinePlotStruct{.title = title, .x = x_coord, .y= bin_to_phi, .z = GetResult()});
        }
    }
    return ret;
}

float AzimuthalIntegrationProfile::GetMeanValueOfBins(uint16_t min_bin, uint16_t max_bin) const {
    std::unique_lock ul(m);

    float ret_sum = 0;
    uint64_t ret_count = 0;

    for (int i = 0; i < sum.size(); i++) {
        uint16_t q_bin = i % q_bins;
        if (q_bin >= min_bin && q_bin <= max_bin) {
            ret_sum += sum[i];
            ret_count += count[i];
        }
    }

    return sum_to_count(ret_sum, ret_count);
}

float AzimuthalIntegrationProfile::GetBkgEstimate(const AzimuthalIntegrationSettings &settings) const {
    auto min_bin = settings.QToBin(settings.GetBkgEstimateLowQ_recipA());
    auto max_bin = settings.QToBin(settings.GetBkgEstimateHighQ_recipA());
    return GetMeanValueOfBins(min_bin, max_bin);
}

float AzimuthalIntegrationProfile::GetIceRingScore(const AzimuthalIntegrationSettings &settings,
                                                   float half_width_q) const {
    // Strongest hexagonal-ice ring's intensity relative to the background *under* it (1 = no ice). The
    // background is a smooth whole-profile estimate: a running median of the NON-ice bins, interpolated to
    // each ring position - not a couple of adjacent shoulder bins (the azint binning is coarser than the
    // ring width, so a local shoulder is only ~1 bin and a narrow ratio is noisy and can double-count the
    // ring's own edge). Clean profiles then sit at ~1 at every ring; ice makes the ring bin stand out.
    constexpr float two_pi = 6.283185307f;
    const std::vector<float> prof = GetResult1D();
    const int nq = static_cast<int>(prof.size());
    const float low_q = settings.GetLowQ_recipA();
    const float dq = settings.GetQSpacing_recipA();
    if (nq < 12 || !(dq > 0.0f))
        return 1.0f;

    auto q_of = [&](int i) { return low_q + (static_cast<float>(i) + 0.5f) * dq; };
    auto on_ice = [&](float q) {
        for (const float d : ICE_RING_RES_A)
            if (std::fabs(q - two_pi / d) < 1.5f * half_width_q)
                return true;
        return false;
    };

    // Non-ice, finite, positive bins (ascending q) carry the background.
    std::vector<int> base;
    for (int i = 0; i < nq; ++i)
        if (std::isfinite(prof[i]) && prof[i] > 0.0f && !on_ice(q_of(i)))
            base.push_back(i);
    if (base.size() < 8)
        return 1.0f;

    // Running median over the base bins => a smooth background robust to the ice peaks.
    constexpr int K = 4;
    std::vector<float> base_bg(base.size());
    std::vector<float> window;
    for (int j = 0; j < static_cast<int>(base.size()); ++j) {
        const int lo = std::max(0, j - K);
        const int hi = std::min(static_cast<int>(base.size()), j + K + 1);
        window.clear();
        for (int m = lo; m < hi; ++m)
            window.push_back(prof[base[m]]);
        std::sort(window.begin(), window.end());
        base_bg[j] = window[window.size() / 2];
    }

    float score = 1.0f;
    for (const float d : ICE_RING_RES_A) {
        const float qr = two_pi / d;
        const int b = static_cast<int>(std::lround((qr - low_q) / dq - 0.5f));
        if (b < 0 || b >= nq || !std::isfinite(prof[b]) || prof[b] <= 0.0f)
            continue;
        // Linear-interpolate the smooth background to the ring position.
        float bg;
        if (qr <= q_of(base.front()))
            bg = base_bg.front();
        else if (qr >= q_of(base.back()))
            bg = base_bg.back();
        else {
            bg = NAN;
            for (int j = 0; j + 1 < static_cast<int>(base.size()); ++j) {
                const float qa = q_of(base[j]), qb = q_of(base[j + 1]);
                if (qa <= qr && qr <= qb) {
                    const float t = (qb > qa) ? (qr - qa) / (qb - qa) : 0.0f;
                    bg = base_bg[j] + t * (base_bg[j + 1] - base_bg[j]);
                    break;
                }
            }
        }
        if (std::isfinite(bg) && bg > 0.0f)
            score = std::max(score, prof[b] / bg);
    }
    return score;
}

AzimuthalIntegrationProfile &AzimuthalIntegrationProfile::operator+=(const AzimuthalIntegrationProfile &other) {
    if ((other.bin_to_q != bin_to_q) || (sum.size() != other.sum.size())) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Error combining two radial integration profiles");
    }

    for (int i = 0; i < sum.size(); i++) {
        sum[i] += other.sum[i];
        sum2[i] += other.sum2[i];
        count[i] += other.count[i];
    }

    return *this;
}

void AzimuthalIntegrationProfile::Add(const DeviceOutput &result) {
    std::unique_lock ul(m);

    if (sum.size() > FPGA_INTEGRATION_BIN_COUNT )
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Error in getting result from FPGA");

    for (int i = 0; i < sum.size(); i++) {
        sum[i] += result.integration_result[i].sum;
        count[i] += result.integration_result[i].count;
    }
}