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

#include "StillsPartialityRefine.h"

#include <atomic>
#include <cmath>
#include <future>
#include <thread>
#include <vector>

#include <ceres/ceres.h>
#include <ceres/rotation.h>

#include "Merge.h"

namespace {
    constexpr size_t MIN_FIT_REFLECTIONS = 20;
    constexpr double kRadToDeg = 180.0 / 3.14159265358979323846;

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

    // One accepted reflection reduced to the physical partiality fit: the base reciprocal vector q (crystal
    // frame at the stored per-image orientation), the reference full intensity, the measured intensity, the
    // Lorentz factor and the weight. dist_ewald / partiality are recomputed from q under the refined tilt.
    struct FitObs {
        double qx, qy, qz;
        double Iref;
        double Iobs;
        double lp;           // 1 / rlp
        double weight;       // 1 / sigma
    };

    double VecLen(double x, double y, double z) { return std::sqrt(x * x + y * y + z * z); }

    // Analytic partiality for a reflection whose base reciprocal vector is q, tilted by (psi_x, psi_y).
    // Mirrors BraggPrediction: dist_ewald = |S| - 1/lambda with S = q_rot + S0, and
    // p = exp(-dist_ewald^2 / 2 sigma^2), sigma^2 = gamma0^2 + (gamma_e*d*)^2 + (bw*|q_z|)^2 - three
    // independent broadenings added in quadrature: the reciprocal-lattice point's own radius (gamma0,
    // ~1/domain size, resolution-INdependent), the mosaic/divergence spread (gamma_e*d*, proportional to
    // d*) and the bandwidth smear along the beam. In practice the fit below returns gamma_e ~ 0 and the
    // width is essentially gamma0 - see there.
    double ComputeP(double qx, double qy, double qz,
                    double psi_x, double psi_y,
                    double s0x, double s0y, double s0z, double inv_lambda,
                    double gamma0, double gamma_e, double bw) {
        double aa[3] = {psi_x, psi_y, 0.0};
        double q[3] = {qx, qy, qz};
        double qr[3];
        ceres::AngleAxisRotatePoint(aa, q, qr);
        const double Sx = qr[0] + s0x, Sy = qr[1] + s0y, Sz = qr[2] + s0z;
        const double de = std::sqrt(Sx * Sx + Sy * Sy + Sz * Sz) - inv_lambda;
        const double dstar = std::sqrt(qr[0] * qr[0] + qr[1] * qr[1] + qr[2] * qr[2]);
        const double sig_ang = gamma_e * dstar;
        const double sbw = bw * std::fabs(qr[2]);
        const double sig2 = gamma0 * gamma0 + sig_ang * sig_ang + sbw * sbw;
        if (!(sig2 > 0.0))
            return 1.0;
        return std::exp(-0.5 * de * de / sig2);
    }

    // Robust per-crystal scale G (linear in G given the model coefficients), identical objective to
    // ScaleOnTheFly::SolveScaleIRLS: minimise sum Cauchy_k( w (G*coeff - Iobs) ) over G >= 0.
    double SolveScaleIRLS(const std::vector<double> &coeff, const std::vector<double> &Iobs,
                          const std::vector<double> &weight, double robust_k) {
        auto weighted_scale = [&](auto robust_weight) {
            double num = 0.0, den = 0.0;
            for (size_t i = 0; i < coeff.size(); ++i) {
                const double rw = robust_weight(i);
                const double w2 = weight[i] * weight[i];
                num += rw * w2 * coeff[i] * Iobs[i];
                den += rw * w2 * coeff[i] * coeff[i];
            }
            return den > 0.0 ? num / den : NAN;
        };

        double G = weighted_scale([](size_t) { 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([&](size_t i) {
                const double res = weight[i] * (G * coeff[i] - Iobs[i]);
                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;
    }

    // Ceres residual: refine the orientation tilt (psi_x, psi_y) holding the scale G fixed. The tilt
    // rotates the base reciprocal vector q; partiality follows analytically. Residual is the intensity
    // mismatch weighted by 1/sigma, exactly matching ScaleOnTheFly's intensity-space objective.
    struct PsiResidual {
        double qx, qy, qz;
        double s0x, s0y, s0z, inv_lambda;
        double gamma0, gamma_e, bw;
        double G, lp, Iref, Iobs, weight;

        template<typename T>
        bool operator()(const T *const psi, T *residual) const {
            T q[3] = {T(qx), T(qy), T(qz)};
            T aa[3] = {psi[0], psi[1], T(0.0)};
            T qr[3];
            ceres::AngleAxisRotatePoint(aa, q, qr);
            const T Sx = qr[0] + T(s0x), Sy = qr[1] + T(s0y), Sz = qr[2] + T(s0z);
            const T de = ceres::sqrt(Sx * Sx + Sy * Sy + Sz * Sz) - T(inv_lambda);
            const T dstar = ceres::sqrt(qr[0] * qr[0] + qr[1] * qr[1] + qr[2] * qr[2]);
            const T sig_ang = T(gamma_e) * dstar;
            const T sbw = T(bw) * ceres::abs(qr[2]);
            const T sig2 = T(gamma0) * T(gamma0) + sig_ang * sig_ang + sbw * sbw;
            const T p = ceres::exp(T(-0.5) * de * de / sig2);
            residual[0] = T(weight) * (T(G) * p * T(lp) * T(Iref) - T(Iobs));
            return true;
        }
    };

    // Gaussian prior N(0, sigma_prior^2) on the tilt. The data residuals above are (model-obs)/sigma, a
    // proper chi^2, so the MAP prior residual is simply dpsi/sigma_prior - no scale calibration needed.
    struct PsiPrior {
        double inv_sigma;
        template<typename T>
        bool operator()(const T *const psi, T *residual) const {
            residual[0] = T(inv_sigma) * psi[0];
            residual[1] = T(inv_sigma) * psi[1];
            return true;
        }
    };
}

StillsPartialityRefine::StillsPartialityRefine(const DiffractionExperiment &x)
    : experiment_(x),
      hkl_key_generator_(x.GetScalingSettings().GetMergeFriedel(), x.GetSpaceGroupNumber().value_or(1)),
      d_min_limit_(x.GetScalingSettings().GetHighResolutionLimit_A()),
      min_partiality_(x.GetScalingSettings().GetMinPartiality()),
      bandwidth_sigma_(x.GetBandwidthFWHM().value_or(0.0f) / 2.3548f) {}

double StillsPartialityRefine::RefineOne(IntegrationOutcome &outcome,
                                         const std::map<HKLKey, double> &reference) const {
    if (outcome.reflections.empty())
        return 0.0;

    const Coord Astar = outcome.latt.Astar();
    const Coord Bstar = outcome.latt.Bstar();
    const Coord Cstar = outcome.latt.Cstar();
    const Coord S0 = outcome.geom.GetScatteringVector();
    const double inv_lambda = 1.0 / outcome.geom.GetWavelength_A();
    const double bw = bandwidth_sigma_;

    auto base_q = [&](const Reflection &r) {
        return Astar * static_cast<float>(r.h) + Bstar * static_cast<float>(r.k)
               + Cstar * static_cast<float>(r.l);
    };

    // Collect the reflections that constrain the fit (accepted, non-ice, finite, present in the reference).
    std::vector<FitObs> obs;
    obs.reserve(outcome.reflections.size());
    // Moments of the excitation error against resolution: de^2 ~ gamma0^2 + gamma_e^2 * d*^2, fitted per
    // crystal by ordinary least squares on (d*^2, de^2). Both components come out of the data.
    double m_n = 0.0, m_x = 0.0, m_xx = 0.0, m_y = 0.0, m_xy = 0.0;
    size_t n_de = 0;
    for (const Reflection &r: outcome.reflections) {
        if (r.on_ice_ring || !AcceptReflection(r, d_min_limit_))
            continue;
        if (!std::isfinite(r.I) || !std::isfinite(r.sigma) || r.sigma <= 0.0f)
            continue;
        const auto it = reference.find(hkl_key_generator_(r));
        if (it == reference.end() || !std::isfinite(it->second))
            continue;

        const Coord q = base_q(r);
        obs.push_back(FitObs{
            .qx = q.x, .qy = q.y, .qz = q.z,
            .Iref = it->second,
            .Iobs = static_cast<double>(r.I),
            .lp = SafeInv(r.rlp, 1.0),
            .weight = SafeInv(r.sigma, 1.0),
        });

        // Excitation error at the stored orientation (psi = 0), collected as the moments of de^2 against
        // d*^2, so BOTH width components are fitted rather than one being forced to zero. Forcing the
        // width to be purely angular (gamma0 = 0) pins it to the high-resolution edge - it is fitted over
        // a d*^2-dense population - and it then collapses at low d*, giving p ~ 0 for reflections that
        // were plainly recorded, which inflated the merged low-resolution intensity scale ~3.6x.
        const double dstar = VecLen(q.x, q.y, q.z);
        const double de0 = VecLen(q.x + S0.x, q.y + S0.y, q.z + S0.z) - inv_lambda;
        if (dstar > 1e-9) {
            const double x = dstar * dstar, y = de0 * de0;
            m_n += 1.0;
            m_x += x;
            m_xx += x * x;
            m_y += y;
            m_xy += x * y;
            ++n_de;
        }
    }

    if (obs.size() < MIN_FIT_REFLECTIONS || n_de == 0)
        return 0.0;

    // Solve the 2x2 normal equations for de^2 = A + B d*^2. A degenerate spread in d* (all reflections in
    // one shell) leaves B undetermined, so fall back to the pure angular width there; a negative fitted
    // component is unphysical and is clamped to zero, which reduces to the previous model.
    //
    // Measured outcome, worth knowing before touching this: the fit does NOT split the width between the
    // two terms - it returns gamma0 ~ 4e-4 1/A and gamma_e ~ 0 (their cross-over sits at d = 0.3 A, far
    // outside any measured range), i.e. a width constant in the LINEAR Ewald distance. That is
    // structural, not a fluke: prediction accepts reflections on a fixed linear |dist_ewald| cutoff, so
    // the accepted population's de^2 is flat in d*^2 by construction and the slope is genuinely ~0. The
    // truncated population cannot constrain an angular term; the resolution-independent one is what the
    // data actually support.
    const double det = m_n * m_xx - m_x * m_x;
    double A = 0.0, B = 0.0;
    if (std::fabs(det) > 1e-30) {
        A = (m_xx * m_y - m_x * m_xy) / det;
        B = (m_n * m_xy - m_x * m_y) / det;
    } else {
        B = m_x > 0.0 ? m_y / m_x : 0.0;
    }
    const double gamma0 = std::sqrt(std::max(0.0, A));
    const double gamma_e_fit = std::max(std::sqrt(std::max(0.0, B)), 1e-9);
    const double gamma_e = settings_.gamma_e > 0.0 ? settings_.gamma_e : gamma_e_fit;

    double psi[2] = {0.0, 0.0};
    double G = 1.0;

    const bool refine_tilt = obs.size() >= settings_.min_reflections;
    const int inner = refine_tilt ? settings_.inner_iterations : 1;

    for (int it = 0; it < inner; ++it) {
        // (1) Solve G given the current partialities.
        std::vector<double> coeff(obs.size()), Iobs(obs.size()), weight(obs.size());
        for (size_t j = 0; j < obs.size(); ++j) {
            const double p = ComputeP(obs[j].qx, obs[j].qy, obs[j].qz, psi[0], psi[1],
                                      S0.x, S0.y, S0.z, inv_lambda, gamma0, gamma_e, bw);
            coeff[j] = p * obs[j].lp * obs[j].Iref;
            Iobs[j] = obs[j].Iobs;
            weight[j] = obs[j].weight;
        }
        G = SolveScaleIRLS(coeff, Iobs, weight, settings_.robust_k);
        if (!(G > 0.0) || !std::isfinite(G))
            return 0.0;

        if (!refine_tilt)
            break;

        // (2) Refine the tilt holding G fixed.
        ceres::Problem problem;
        for (const auto &o: obs) {
            auto *cost = new ceres::AutoDiffCostFunction<PsiResidual, 1, 2>(new PsiResidual{
                .qx = o.qx, .qy = o.qy, .qz = o.qz,
                .s0x = S0.x, .s0y = S0.y, .s0z = S0.z, .inv_lambda = inv_lambda,
                .gamma0 = gamma0, .gamma_e = gamma_e, .bw = bw,
                .G = G, .lp = o.lp, .Iref = o.Iref, .Iobs = o.Iobs, .weight = o.weight});
            problem.AddResidualBlock(cost, new ceres::CauchyLoss(settings_.robust_k), psi);
        }
        if (settings_.prior_sigma_deg > 0.0) {
            const double inv_sigma = kRadToDeg / settings_.prior_sigma_deg;  // 1 / sigma_prior (rad)
            problem.AddResidualBlock(new ceres::AutoDiffCostFunction<PsiPrior, 2, 2>(
                                         new PsiPrior{inv_sigma}), nullptr, psi);
        }
        problem.SetParameterLowerBound(psi, 0, -settings_.max_tilt_rad);
        problem.SetParameterUpperBound(psi, 0, settings_.max_tilt_rad);
        problem.SetParameterLowerBound(psi, 1, -settings_.max_tilt_rad);
        problem.SetParameterUpperBound(psi, 1, settings_.max_tilt_rad);

        ceres::Solver::Options options;
        options.linear_solver_type = ceres::DENSE_QR;
        options.minimizer_progress_to_stdout = false;
        options.num_threads = 1;
        options.max_num_iterations = 25;
        ceres::Solver::Summary summary;
        const double psi_before[2] = {psi[0], psi[1]};
        ceres::Solve(options, &problem, &summary);
        // A solve that failed leaves whatever the minimiser last wrote in psi, and that tilt would go
        // straight onto every reflection's partiality below. Keep the tilt this crystal came in with
        // and stop refining it instead - G alone is still a usable model.
        if (!summary.IsSolutionUsable()) {
            psi[0] = psi_before[0];
            psi[1] = psi_before[1];
            break;
        }
    }

    // Write the refined partiality + scale correction onto every reflection of the crystal (not only the
    // fit subset), so the merge sees a consistent model. image_scale_corr = rlp / (partiality * G), the
    // same composition ScaleOnTheFly writes.
    for (auto &r: outcome.reflections) {
        const Coord q = base_q(r);
        const double p = ComputeP(q.x, q.y, q.z, psi[0], psi[1], S0.x, S0.y, S0.z, inv_lambda,
                                  gamma0, gamma_e, bw);
        r.partiality = static_cast<float>(p);
        const double denom = p * G;
        r.image_scale_corr = (std::isfinite(r.rlp) && std::isfinite(denom) && denom > 0.0)
                                 ? static_cast<float>(r.rlp / denom)
                                 : NAN;
    }
    outcome.image_scale_g = static_cast<float>(G);

    // The corrections just changed, so the CC that ScaleOnTheFly measured no longer describes them.
    // Refresh it here: it is reported per image and --min-image-cc drops images by it, so it has to be
    // the CC of the data that is actually merged.
    const auto [cc, cc_n] = ImageReferenceCC(outcome.reflections, reference, hkl_key_generator_,
                                             d_min_limit_, min_partiality_);
    outcome.image_scale_cc = cc;
    outcome.image_scale_cc_n = cc_n;

    const double tilt_deg = std::sqrt(psi[0] * psi[0] + psi[1] * psi[1]) * kRadToDeg;
    return tilt_deg;
}

double StillsPartialityRefine::Run(std::vector<IntegrationOutcome> &outcomes, size_t nthreads) const {
    if (nthreads == 0)
        nthreads = std::thread::hardware_concurrency();
    nthreads = std::max<size_t>(1, nthreads);

    double last_mean_tilt = 0.0;

    for (int outer = 0; outer < settings_.outer_iterations; ++outer) {
        // Reference full intensities from the current corrections.
        const std::vector<MergedReflection> merged = MergeAll(experiment_, outcomes);
        std::map<HKLKey, double> reference;
        for (const auto &m: merged)
            reference[hkl_key_generator_(m)] = m.I;

        std::atomic<double> tilt_sum{0.0};
        std::atomic<size_t> tilt_n{0};
        std::atomic<size_t> next{0};

        auto worker = [&]() {
            size_t i = next.fetch_add(1);
            while (i < outcomes.size()) {
                const double t = RefineOne(outcomes[i], reference);
                if (t > 0.0) {
                    double prev = tilt_sum.load();
                    while (!tilt_sum.compare_exchange_weak(prev, prev + t)) {}
                    tilt_n.fetch_add(1);
                }
                i = next.fetch_add(1);
            }
        };

        const size_t nt = std::min(nthreads, std::max<size_t>(1, outcomes.size()));
        if (nt <= 1) {
            worker();
        } else {
            std::vector<std::future<void>> futures;
            futures.reserve(nt);
            for (size_t t = 0; t < nt; ++t)
                futures.emplace_back(std::async(std::launch::async, worker));
            for (auto &f: futures)
                f.get();
        }

        last_mean_tilt = tilt_n > 0 ? tilt_sum.load() / static_cast<double>(tilt_n.load()) : 0.0;
    }

    return last_mean_tilt;
}