// 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>
#include <ceres/ceres.h>

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.
    struct IntensityFixedResidual {
        IntensityFixedResidual(const Reflection &r, double Itrue, double sigma)
            : Iobs(r.I),
              Itrue(Itrue),
              weight(SafeInv(sigma, 1.0)),
              lp(SafeInv(r.rlp, 1.0)),
              b_resolution_coeff(-SafeInv(4.0 * r.d * r.d, 0.0)),
              partiality(r.partiality) {
        }

        template<typename T>
        bool operator()(const T *const G, const T *const B, T *residual) const {
            const T B_term = ceres::exp(B[0] * T(b_resolution_coeff));
            residual[0] = (G[0] * T(partiality) * B_term * T(lp) * Itrue - T(Iobs)) * T(weight);
            return true;
        }

        const double Iobs;
        const double Itrue;
        const double weight;
        const double lp;
        const double b_resolution_coeff;
        const double partiality;
    };
}

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

std::pair<double, size_t> ScaleOnTheFly::CalculateGlobalCC(const std::vector<Reflection> &reflections) const {
    double sum_x = 0.0;
    double sum_y = 0.0;
    double sum_x2 = 0.0;
    double sum_y2 = 0.0;
    double sum_xy = 0.0;
    size_t n = 0;

    for (const auto &r: reflections) {
        if (r.on_ice_ring)
            continue;
        if (!AcceptReflection(r, s.GetHighResolutionLimit_A()))
            continue;
        if (r.partiality < s.GetMinPartiality())
            continue;
        if (!std::isfinite(r.I) || !std::isfinite(r.image_scale_corr) || r.image_scale_corr <= 0.0f)
            continue;
        if (!std::isfinite(r.sigma) || r.sigma <= 0.0f)
            continue;

        const HKLKey key = hkl_key_generator(r);
        const auto it = reference_data.find(key);
        if (it == reference_data.end())
            continue;

        const double image_i = static_cast<double>(r.I) * static_cast<double>(r.image_scale_corr);
        const double ref_i = it->second;

        if (!std::isfinite(image_i) || !std::isfinite(ref_i))
            continue;

        sum_x += image_i;
        sum_y += ref_i;
        sum_x2 += image_i * image_i;
        sum_y2 += ref_i * ref_i;
        sum_xy += image_i * ref_i;
        ++n;
    }

    if (n < MIN_REFLECTIONS)
        return {NAN, n};

    const double nd = static_cast<double>(n);
    const double cov = sum_xy - sum_x * sum_y / nd;
    const double var_x = sum_x2 - sum_x * sum_x / nd;
    const double var_y = sum_y2 - sum_y * sum_y / nd;

    if (!(var_x > 0.0 && var_y > 0.0))
        return {NAN, n};

    return {cov / std::sqrt(var_x * var_y), n};
}

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

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

    ScaleOnTheFlyResult result{ .B = 0.0, .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();
        integration_outcome.image_scale_b_factor_Ang2.reset();
    };

    // With B fixed 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) instead of a Ceres problem per image. Ceres is kept only when
    // the B-factor (exp(-B/...)) is refined.
    const bool linear_in_g = !s.GetRefineB();

    if (linear_in_g) {
        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 B_term = std::exp(result.B * -SafeInv(4.0 * r.d * r.d, 0.0));
            const double coeff = r.partiality * B_term * 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);
    } else {
        ceres::Problem problem;

        size_t n_reflections = 0;
        for (const auto &r: integration_outcome.reflections) {
            if (!Accept(r))
                continue;

            const HKLKey key = hkl_key_generator(r);
            if (!reference_data.contains(key))
                continue;

            ++n_reflections;

            auto *cost = new ceres::AutoDiffCostFunction<IntensityFixedResidual, 1, 1, 1>(
                new IntensityFixedResidual(r, reference_data.at(key), r.sigma));
            problem.AddResidualBlock(cost, new ceres::CauchyLoss(SCALE_ROBUST_K), &result.G, &result.B);
        }

        if (n_reflections < MIN_REFLECTIONS) {
            clear_scale();
            return;
        }

        problem.SetParameterLowerBound(&result.G, 0, 0.0);
        problem.SetParameterLowerBound(&result.B, 0, s.GetMinB());
        problem.SetParameterUpperBound(&result.B, 0, s.GetMaxB());

        ceres::Solver::Options options;
        options.linear_solver_type = ceres::DENSE_QR;
        options.minimizer_progress_to_stdout = false;
        options.num_threads = 1;

        ceres::Solver::Summary summary;
        ceres::Solve(options, &problem, &summary);
    }

    for (auto &r: integration_outcome.reflections) {
        const double B_term = exp(result.B * -SafeInv(4.0 * r.d * r.d, 0.0));
        const double denom = B_term * 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] = CalculateGlobalCC(integration_outcome.reflections);
    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();

    if (s.GetRefineB())
        integration_outcome.image_scale_b_factor_Ang2 = result.B;
    else
        integration_outcome.image_scale_b_factor_Ang2.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();
    }
}