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

#include "ScaleOnTheFly.h"

#include "../../common/Logger.h"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <future>
#include <vector>

namespace {
    // Robust loss scale (in sigma units) for the per-image scale fit: a few outlier reflections
    // (zingers, overlaps, a mis-predicted spot) must not drag a frame's G/B into a bad optimum -
    // that is the stochastic per-frame mis-scaling that elevates R-meas and collapses CC1/2 at low
    // symmetry. Cauchy down-weights residuals beyond ~this many sigma without a hard cut.
    constexpr double SCALE_ROBUST_K = 3.0;

    // Smallest per-image scale, relative to the run's median, that is still believable as a scale
    // rather than a failed fit. The same ratio guards the rotation path's per-frame scales.
    constexpr double MIN_CREDIBLE_SCALE_RATIO = 0.02;

    double SafeInv(double x, double fallback) {
        if (!std::isfinite(x) || x == 0.0)
            return fallback;
        return 1.0 / x;
    }

    // One reflection reduced to the 1-D scale fit: predicted intensity is G * coeff (coeff is constant
    // while B is fixed), measured is Iobs, weighted by 1/sigma.
    struct ScaleObs {
        double coeff;
        double Iobs;
        double weight;
    };

    // Robust per-image scale: minimise sum_i Cauchy_k( weight_i (G*coeff_i - Iobs_i) ) over G >= 0. The
    // model is linear in G, so this M-estimate is a few reweighted-least-squares steps (each a closed-form
    // weighted ratio) - the same objective the Ceres path solves, without a per-image problem/autodiff/
    // trust-region. Seeded from the plain weighted-LS solution; Cauchy weight is 1/(1 + (res/k)^2).
    double SolveScaleIRLS(const std::vector<ScaleObs> &obs, double robust_k) {
        auto weighted_scale = [&obs](auto robust_weight) {
            double num = 0.0, den = 0.0;
            for (const auto &o: obs) {
                const double rw = robust_weight(o);
                const double w2 = o.weight * o.weight;
                num += rw * w2 * o.coeff * o.Iobs;
                den += rw * w2 * o.coeff * o.coeff;
            }
            return den > 0.0 ? num / den : NAN;
        };

        double G = weighted_scale([](const ScaleObs &) { return 1.0; });
        if (!std::isfinite(G))
            return 1.0;
        G = std::max(0.0, G);

        const double k2 = robust_k * robust_k;
        for (int iter = 0; iter < 30; ++iter) {
            const double G_prev = G;
            const double G_next = weighted_scale([&](const ScaleObs &o) {
                const double res = o.weight * (G * o.coeff - o.Iobs);
                return 1.0 / (1.0 + res * res / k2);
            });
            if (!std::isfinite(G_next))
                break;
            G = std::max(0.0, G_next);
            if (std::abs(G - G_prev) <= 1e-7 * std::max(G, 1.0))
                break;
        }
        return G;
    }

}

ScaleOnTheFly::ScaleOnTheFly(const DiffractionExperiment &x, const std::vector<MergedReflection> &ref)
    : s(x.GetScalingSettings()),
      hkl_key_generator(s.GetMergeFriedel(), x.GetSpaceGroupNumber().value_or(1)) {
    for (const auto &r: ref) {
        const auto key = hkl_key_generator(r);
        reference_data[key] = r.I;
    }
}

bool ScaleOnTheFly::Accept(const Reflection &r) const {
    if (r.on_ice_ring)   // ice-contaminated intensity would drag the per-image scale; keep it out of the fit
        return false;

    return AcceptReflection(r, s.GetHighResolutionLimit_A());
}

void ScaleOnTheFly::Scale(IntegrationOutcome &integration_outcome) const {
    if (integration_outcome.reflections.empty())
        return;

    ScaleOnTheFlyResult result{ .G = 1.0 };

    auto clear_scale = [&]() {
        integration_outcome.image_scale_cc.reset();
        integration_outcome.image_scale_cc_n.reset();
        integration_outcome.image_scale_g.reset();
    };

    // The fixed-partiality model G * coeff is linear in G, so the robust per-image scale is a 1-D
    // M-estimate solved directly (IRLS) rather than a Ceres problem per image.
    {
        std::vector<ScaleObs> obs;
        obs.reserve(integration_outcome.reflections.size());
        for (const auto &r: integration_outcome.reflections) {
            if (!Accept(r))
                continue;
            const auto it = reference_data.find(hkl_key_generator(r));
            if (it == reference_data.end())
                continue;
            const double coeff = r.partiality * SafeInv(r.rlp, 1.0) * it->second;
            obs.push_back({coeff, static_cast<double>(r.I), SafeInv(r.sigma, 1.0)});
        }

        if (obs.size() < MIN_REFLECTIONS) {
            clear_scale();
            return;
        }

        result.G = SolveScaleIRLS(obs, SCALE_ROBUST_K);
    }

    for (auto &r: integration_outcome.reflections) {
        const double denom = r.partiality * result.G;
        r.image_scale_corr = (std::isfinite(r.rlp) && std::isfinite(denom) && denom > 0.0)
                                 ? static_cast<float>(r.rlp / denom)
                                 : NAN;
    }

    const auto [cc, cc_n] = ImageReferenceCC(integration_outcome.reflections, reference_data,
                                            hkl_key_generator, s.GetHighResolutionLimit_A(),
                                            s.GetMinPartiality());
    result.cc = cc;
    result.cc_n = cc_n;

    integration_outcome.image_scale_cc = cc;
    integration_outcome.image_scale_cc_n = cc_n;
    integration_outcome.image_scale_g = result.G;
    integration_outcome.image_scale_wedge_deg.reset();
}

// A per-image scale that has collapsed toward zero multiplies that image's intensities by 1/G - and its
// sigmas by the same factor, so nothing downstream can recognise it: the merge's n-sigma outlier test
// scales with the very number that is wrong. The rotation path already refuses a per-frame scale this
// far below its neighbours; the stills path had no such guard.
//
// Such an image is DROPPED from the merge, not merged unscaled. Substituting G = 1 looks conservative
// but is the more damaging of the two errors: if the collapsed value was a failed fit, the image goes in
// mis-scaled by an unknown factor, and if it was a real scale (per-crystal scales on serial stills
// genuinely span orders of magnitude, unlike frames of one rotation sweep) then G = 1 divides both its
// intensities AND its sigmas by 1/G, so it enters at 1/G^2 times the weight it deserves - the merge
// cannot down-weight it, because the number that is wrong is the same number the weight is built from.
// An image whose scale is not believable has no usable scale, so it contributes nothing instead.
void ScaleOnTheFly::RejectCollapsedScales(std::vector<IntegrationOutcome> &integration) {
    std::vector<double> fitted;
    fitted.reserve(integration.size());
    for (const auto &i: integration)
        if (i.image_scale_g && std::isfinite(*i.image_scale_g) && *i.image_scale_g > 0.0)
            fitted.push_back(*i.image_scale_g);
    if (fitted.size() < 2)
        return;

    const size_t mid = fitted.size() / 2;
    std::nth_element(fitted.begin(), fitted.begin() + mid, fitted.end());
    const double g_floor = fitted[mid] * MIN_CREDIBLE_SCALE_RATIO;

    int64_t n_rejected = 0;
    for (auto &i: integration) {
        if (!i.image_scale_g || !std::isfinite(*i.image_scale_g) || *i.image_scale_g >= g_floor)
            continue;
        // A non-finite correction is what every merge path already skips on (Merge.cpp), so this drops
        // the image consistently from the merged intensities, the error model and the statistics.
        for (auto &r: i.reflections)
            r.image_scale_corr = NAN;
        i.image_scale_cc.reset();
        i.image_scale_cc_n.reset();
        i.image_scale_g.reset();
        ++n_rejected;
    }
    if (n_rejected > 0)
        Logger("ScaleOnTheFly").Warning(
            "Dropped {} image(s) from the merge: their per-image scale collapsed more than {:.0f}x below "
            "the run median, so it is not a scale the intensities can be put on",
            n_rejected, 1.0 / MIN_CREDIBLE_SCALE_RATIO);
}

void ScaleOnTheFly::Scale(std::vector<IntegrationOutcome> &integration, size_t nthreads) const {
    if (nthreads == 0)
        nthreads = std::thread::hardware_concurrency();

    if (nthreads <= 1) {
        for (auto & i : integration)
            Scale(i);
    } else {
        auto local_nthreads = std::min(nthreads, integration.size());
        std::vector<std::future<void>> futures;
        futures.reserve(local_nthreads);
        std::atomic<size_t> curr_image = 0;

        for (size_t t = 0; t < local_nthreads; ++t)
            futures.emplace_back(std::async(std::launch::async, [&] {
                size_t i = curr_image.fetch_add(1);
                while (i < integration.size()) {
                    Scale(integration[i]);
                    i = curr_image.fetch_add(1);
                }
            }));

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

    RejectCollapsedScales(integration);
}