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

#include "ScaleOnTheFly.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;

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

    // The fixed-partiality residual for the Ceres path (used only when the B-factor is refined): the
    // stored partiality is a constant, so the model is G * partiality * exp(-B/(4 d^2)) * (1/rlp) * Itrue.
}

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;

    auto start = std::chrono::steady_clock::now();

    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;

    auto end = std::chrono::steady_clock::now();
    result.time_s = std::chrono::duration<float>(end - start).count();

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

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