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

#include "ErrorModel.h"

#include <algorithm>
#include <cmath>
#include <functional>
#include <future>

#include <Eigen/Dense>

#include "../../common/ParallelFor.h"
#include "../../common/ThreadAffinity.h"

ErrorModelFit FitErrorModel(const std::vector<ErrorModelSample> &pool, std::vector<ErrorModelBinned> &smp,
                            size_t nthreads) {
    constexpr int n_bins = ERROR_MODEL_BINS;
    ErrorModelFit fit;
    if (pool.size() < static_cast<size_t>(8 * n_bins))
        return fit;

    // The rank key is computed once, here, and carried with the sample: a key recomputed inside the
    // comparator can round differently at different inlined call sites, and nth_element then sees an
    // inconsistent order.
    smp.resize(pool.size());
    ParallelChunks(static_cast<int>(pool.size()), ThreadsForWork(pool.size(), nthreads), [&](int lo, int hi) {
        for (int i = lo; i < hi; ++i)
            smp[i] = {pool[i].I2 / pool[i].s2, pool[i].s2, pool[i].I2, pool[i].dev2, 0.0};
    });

    // Equal-count bins, each boundary put in place with nth_element; the split runs down the middle so
    // each level halves the range it works on. Total order, so which of two equal keys lands in which bin
    // is a property of the samples, not of the order they arrived in. The two halves of a split are
    // disjoint ranges, so they can run side by side while the range is still worth a thread.
    const size_t per = smp.size() / n_bins;
    const auto by_snr = [](const ErrorModelBinned &a, const ErrorModelBinned &b) {
        if (a.snr2 != b.snr2) return a.snr2 < b.snr2;
        if (a.s2 != b.s2) return a.s2 < b.s2;
        if (a.I2 != b.I2) return a.I2 < b.I2;
        return a.dev2 < b.dev2;
    };
    constexpr size_t SPLIT_MIN_PARALLEL = 1u << 17;
    const std::function<void(size_t, size_t, int, int, int)> split =
            [&](size_t lo, size_t hi, int b0, int b1, int depth) {
                if (b0 >= b1) return;
                const int mid = (b0 + b1) / 2;
                const size_t k = static_cast<size_t>(mid) * per;
                std::nth_element(smp.begin() + lo, smp.begin() + k, smp.begin() + hi, by_snr);
                if (depth > 0 && hi - lo > SPLIT_MIN_PARALLEL) {
                    auto left = std::async(std::launch::async,
                                           [&, lo, k, b0, mid, depth] {
                                               // Not the node of a merging engine that called this (ScopedGpuPin).
                                               RestoreThreadAffinity();
                                               split(lo, k, b0, mid, depth - 1);
                                           });
                    split(k, hi, mid + 1, b1, depth - 1);
                    left.get();
                } else {
                    split(lo, k, b0, mid, 0);
                    split(k, hi, mid + 1, b1, 0);
                }
            };
    split(0, smp.size(), 1, n_bins, nthreads > 1 ? 3 : 0);
    auto bin_lo = [&](int bin) { return static_cast<size_t>(bin) * per; };
    auto bin_hi = [&](int bin) { return bin == n_bins - 1 ? smp.size() : static_cast<size_t>(bin + 1) * per; };

    // b is identified only if the strongest bin reaches where b*<I> can outgrow the counting term: when
    // fewer reflections are strong than one bin holds, that bin sits at I/sigma ~ 1 and a two-term fit
    // hands b the bins' own noise instead (measured: b = 5.6, ISa 0.5, on a crystal XDS gives 6.1). Then
    // fit a alone and leave b at zero. Over the rotation battery the two crystals this fires on sit at
    // (I/sigma)^2 0.22 and 0.84 while the next is 31.7, so the threshold is not delicate.
    {
        const auto top_mid = smp.begin() + (bin_lo(n_bins - 1) + bin_hi(n_bins - 1)) / 2;
        std::nth_element(smp.begin() + bin_lo(n_bins - 1), top_mid, smp.end(), by_snr);
        fit.b_measured = top_mid->snr2 > ERROR_MODEL_B_LEVER_MIN;
    }

    std::vector<double> bs2(n_bins), bI2(n_bins), bd2(n_bins);
    double a = 1.0, b2 = 0.0;
    for (int iter = 0; iter < 30; ++iter) {
        // Per bin: the terms of var = a*s2 + b^2*I2, each divided by the current model's variance, and
        // the median normalised squared deviation, expressed as the sum it stands for. Reordering inside
        // a bin does not change which samples it holds. One bin per task, so the result does not depend
        // on threads.
        ParallelChunks(n_bins, ThreadsForWork(smp.size(), nthreads), [&](int blo, int bhi) {
            for (int bin = blo; bin < bhi; ++bin) {
                double s = 0.0, I = 0.0;
                for (size_t i = bin_lo(bin); i < bin_hi(bin); ++i) {
                    auto &x = smp[i];
                    const double v = a * x.s2 + b2 * x.I2;
                    s += x.s2 / v; I += x.I2 / v;
                    x.z2 = x.dev2 / v;
                }
                const size_t n = bin_hi(bin) - bin_lo(bin);
                const auto mid = smp.begin() + bin_lo(bin) + n / 2;
                std::nth_element(smp.begin() + bin_lo(bin), mid, smp.begin() + bin_hi(bin),
                                 [](const ErrorModelBinned &p, const ErrorModelBinned &q) { return p.z2 < q.z2; });
                bs2[bin] = s; bI2[bin] = I; bd2[bin] = mid->z2 / ERROR_MODEL_CHI2_1_MEDIAN * static_cast<double>(n);
            }
        });

        if (!ErrorModelUpdate(bs2, bI2, bd2, a, b2, fit))
            break;
    }
    return fit;
}

bool ErrorModelUpdate(const std::vector<double> &bs2, const std::vector<double> &bI2, const std::vector<double> &bd2,
                      double &a, double &b2, ErrorModelFit &fit) {
    const int n_bins = static_cast<int>(bs2.size());
    // Every bin weighs the same in RELATIVE terms, so the weak bins still constrain a.
    double Ass = 0, AsI = 0, AII = 0, Bs = 0, BI = 0;
    for (int bin = 0; bin < n_bins; ++bin) {
        if (!(bd2[bin] > 0.0)) continue;
        const double w = 1.0 / (bd2[bin] * bd2[bin]);
        Ass += w * bs2[bin] * bs2[bin]; AsI += w * bs2[bin] * bI2[bin]; AII += w * bI2[bin] * bI2[bin];
        Bs += w * bs2[bin] * bd2[bin]; BI += w * bI2[bin] * bd2[bin];
    }
    const double det = Ass * AII - AsI * AsI;
    double a_new, b2_new;
    if (!fit.b_measured) {
        if (!(Ass > 0.0)) return false;
        a_new = std::clamp(Bs / Ass, 0.25, 100.0);
        b2_new = 0.0;
        fit.b_resolved = false;
        fit.a_at_bound = a_new != Bs / Ass;
    } else {
        if (!(det > 1e-10 * Ass * AII)) return false;
        const double a_fit = (Bs * AII - BI * AsI) / det;
        const double b2_fit = (Ass * BI - AsI * Bs) / det;
        a_new = std::clamp(a_fit, 0.25, 100.0);
        b2_new = std::max(b2_fit, 0.0);
        // Leverage is not enough: b^2 can have it and still come out at zero within its own error,
        // and 1/b then reads the noise in b as an I/sigma. The standard error of b^2 comes from the
        // bins' own scatter about the fitted line; this only decides what is printed.
        // Nor is a standard error enough when a sits on a bound: the b^2 reported beside a clamped a
        // is not the solution of any fit. A bin whose median deviation lies far below the others' -
        // observations that repeat their mates to the bit - takes the whole fit through its 1/bd2^2
        // weight and pulls both terms towards zero; a stops at its bound, b^2 does not, and 1/b of
        // the remainder reads as an I/sigma in the thousands or millions.
        double chi = 0.0;
        for (int bin = 0; bin < n_bins; ++bin) {
            if (!(bd2[bin] > 0.0)) continue;
            const double r = (bd2[bin] - a_fit * bs2[bin] - b2_fit * bI2[bin]) / bd2[bin];
            chi += r * r;
        }
        fit.b_resolved = a_new == a_fit && b2_fit >= 2.0 * std::sqrt(chi / (n_bins - 2) * Ass / det);
        fit.a_at_bound = a_new != a_fit;
    }
    const bool settled = std::fabs(a_new - a) <= 1e-4 * a && std::fabs(b2_new - b2) <= 1e-4 * b2;
    a = a_new;
    b2 = b2_new;
    fit.active = true;
    fit.a = a;
    fit.b2 = b2;
    return !settled;
}

namespace {
    // SplitMix64 as a hash of (replicate, group) to draw from.
    // Following Steele, Lea & Flood (2014) OOPSLA 2014, 453-472
    uint64_t SplitMix64(uint64_t x) {
        x += 0x9E3779B97F4A7C15ull;
        x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ull;
        x = (x ^ (x >> 27)) * 0x94D049BB133111EBull;
        return x ^ (x >> 31);
    }

    // A Poisson count by inversion of its CDF from a uniform in [0, 1), p0 = exp(-lambda). The cap stops a
    // u that the rounded CDF never passes; for lambda <= 1, past 20 the probability is below 1e-19.
    int Poisson(double lambda, double p0, double u) {
        double p = p0, cdf = p0;
        int k = 0;
        while (u >= cdf && k < 20) {
            ++k;
            p *= lambda / k;
            cdf += p;
        }
        return k;
    }
}

std::optional<ErrorModelB2Interval> BootstrapErrorModelB2(const std::vector<ErrorModelSample> &pool, double b2,
                                                          int replicates, size_t max_samples, size_t nthreads) {
    // The bootstrap over groups with Poisson weights in place of a multinomial draw: Efron (1979) Ann.
    // Statist. 7, 1-26. A pool larger than max_samples is resampled thinned, each group drawn
    // Poisson(lambda) times, and the replicates' spread about b2 shrunk by sqrt(lambda) to the size of the
    // pool - the m-out-of-n bootstrap, for a statistic whose error goes as one over the root of its count.
    const double lambda = std::min(1.0, static_cast<double>(max_samples) / static_cast<double>(pool.size()));
    const double p0 = std::exp(-lambda);
    // One replicate per task, each fitted on one thread: every replicate's draws and fit are its own, so
    // the interval is the same on any number of threads.
    std::vector<double> b2_r(replicates, NAN);
    ParallelChunks(replicates, nthreads, [&](int lo, int hi) {
        std::vector<ErrorModelSample> replicate;
        std::vector<ErrorModelBinned> scratch;
        for (int r = lo; r < hi; ++r) {
            replicate.clear();
            for (const auto &s : pool) {
                const uint64_t h = SplitMix64((static_cast<uint64_t>(r) << 32) | static_cast<uint32_t>(s.group));
                const int copies = Poisson(lambda, p0, static_cast<double>(h >> 11) * 0x1.0p-53);
                for (int c = 0; c < copies; ++c) replicate.push_back(s);
            }
            const auto fit = FitErrorModel(replicate, scratch, 1);
            if (fit.active) b2_r[r] = fit.b2;
        }
    });
    std::erase_if(b2_r, [](double v) { return std::isnan(v); });
    if (b2_r.empty())
        return std::nullopt;
    std::sort(b2_r.begin(), b2_r.end());
    const double last = static_cast<double>(b2_r.size() - 1);
    const double q_low = b2_r[static_cast<size_t>(std::floor(0.025 * last))];
    const double q_high = b2_r[static_cast<size_t>(std::ceil(0.975 * last))];
    // A replicate's b^2 held at zero is b reaching zero, which no shrinking moves away from it.
    const double shrink = std::sqrt(lambda);
    return ErrorModelB2Interval{q_low > 0.0 ? std::max(0.0, b2 + shrink * (q_low - b2)) : 0.0,
                                std::max(0.0, b2 + shrink * (q_high - b2))};
}

StillsErrorModelFit FitStillsErrorModel(const std::vector<StillsErrorModelSample> &pool, size_t nthreads,
                                        const StillsErrorModelFit &start) {
    constexpr int n_bins = ERROR_MODEL_BINS;
    constexpr int n_sub = 4;
    StillsErrorModelFit fit;
    if (pool.size() < static_cast<size_t>(8 * n_bins * n_sub))
        return fit;

    // Bins of counting (I/sigma)^2, each split by the partiality term's share of the intensity term. The
    // boundaries are put in place with nth_element on (key, position in the pool) - a total order, so which
    // sample lands in which bin depends on the samples and the order the caller gave them in alone. The keys
    // are computed once and carried in the comparison, never recomputed inside it.
    struct Key { double v; size_t i; };
    const auto by_key = [](const Key &p, const Key &q) { return p.v < q.v || (p.v == q.v && p.i < q.i); };
    std::vector<Key> key(pool.size());
    ParallelChunks(static_cast<int>(pool.size()), ThreadsForWork(pool.size(), nthreads), [&](int lo, int hi) {
        for (int i = lo; i < hi; ++i)
            key[i] = {pool[i].I2 / pool[i].s2, static_cast<size_t>(i)};
    });
    const size_t per = pool.size() / n_bins;
    // The split runs down the middle, and the two halves of a split are disjoint ranges, so they run side
    // by side while the range is still worth a thread (as in FitErrorModel).
    constexpr size_t SPLIT_MIN_PARALLEL = 1u << 17;
    const std::function<void(size_t, size_t, int, int, int)> split = [&](size_t lo, size_t hi, int b0, int b1, int depth) {
        if (b0 >= b1) return;
        const int mid = (b0 + b1) / 2;
        const size_t k = static_cast<size_t>(mid) * per;
        std::nth_element(key.begin() + lo, key.begin() + k, key.begin() + hi, by_key);
        if (depth > 0 && hi - lo > SPLIT_MIN_PARALLEL) {
            auto left = std::async(std::launch::async, [&, lo, k, b0, mid, depth] {
                RestoreThreadAffinity();
                split(lo, k, b0, mid, depth - 1);
            });
            split(k, hi, mid + 1, b1, depth - 1);
            left.get();
        } else {
            split(lo, k, b0, mid, 0);
            split(k, hi, mid + 1, b1, 0);
        }
    };
    split(0, key.size(), 1, n_bins, nthreads > 1 ? 3 : 0);
    // b only where the strongest bin reaches past the counting term, as in FitErrorModel.
    std::vector<double> top;
    for (size_t j = (n_bins - 1) * per; j < key.size(); ++j)
        top.push_back(key[j].v);
    std::nth_element(top.begin(), top.begin() + top.size() / 2, top.end());
    const bool b_measured = top[top.size() / 2] > ERROR_MODEL_B_LEVER_MIN;
    // Each bin in four by the share, and the samples copied out in that order.
    std::vector<std::pair<size_t, size_t>> groups(n_bins * n_sub);
    std::vector<StillsErrorModelSample> x(pool.size());
    ParallelFor(n_bins, nthreads, [&](int bin) {
        const size_t lo = bin * per, hi = bin == n_bins - 1 ? key.size() : (bin + 1) * per;
        for (size_t j = lo; j < hi; ++j) {
            const auto &o = pool[key[j].i];
            key[j].v = o.I2 > 0.0 ? o.P2 / o.I2 : 0.0;
        }
        const size_t sub = (hi - lo) / n_sub;
        for (int q = 1; q < n_sub; ++q)
            std::nth_element(key.begin() + lo + (q - 1) * sub, key.begin() + lo + q * sub, key.begin() + hi, by_key);
        for (int q = 0; q < n_sub; ++q)
            groups[bin * n_sub + q] = {lo + q * sub, q == n_sub - 1 ? hi : lo + (q + 1) * sub};
        for (size_t j = lo; j < hi; ++j)
            x[j] = pool[key[j].i];
    });
    key = {};

    const int ng = static_cast<int>(groups.size());
    std::vector<double> z2(pool.size());
    Eigen::MatrixXd X(ng, 3);
    Eigen::VectorXd y(ng);
    double a = start.a, b2 = start.b2, c2 = start.c2;
    for (int iter = 0; iter < 30; ++iter) {
        // Per group: the three terms divided by the current variance, and the median normalised deviation
        // expressed as the sum it stands for.
        ParallelFor(ng, ThreadsForWork(pool.size(), nthreads), [&](int g) {
            const auto [lo, hi] = groups[g];
            double s = 0.0, I = 0.0, P = 0.0;
            for (size_t i = lo; i < hi; ++i) {
                const double v = a * x[i].s2 + b2 * x[i].I2 + c2 * x[i].P2;
                s += x[i].s2 / v; I += x[i].I2 / v; P += x[i].P2 / v;
                z2[i] = x[i].dev2 / v;
            }
            const auto mid = z2.begin() + lo + (hi - lo) / 2;
            std::nth_element(z2.begin() + lo, mid, z2.begin() + hi);
            X.row(g) << s, I, P;
            y(g) = *mid / ERROR_MODEL_CHI2_1_MEDIAN * static_cast<double>(hi - lo);
        });
        // Every group weighs the same in relative terms; a negative term is dropped and the rest refitted.
        const Eigen::VectorXd w = y.array().square().inverse();
        std::vector<int> terms = b_measured ? std::vector<int>{0, 1, 2} : std::vector<int>{0, 2};
        Eigen::VectorXd sol = Eigen::VectorXd::Zero(3);
        while (true) {
            Eigen::MatrixXd Xs(ng, terms.size());
            for (size_t j = 0; j < terms.size(); ++j)
                Xs.col(j) = X.col(terms[j]);
            const Eigen::MatrixXd A = Xs.transpose() * w.asDiagonal() * Xs;
            const Eigen::VectorXd s = A.ldlt().solve(Xs.transpose() * w.asDiagonal() * y);
            sol.setZero();
            for (size_t j = 0; j < terms.size(); ++j)
                sol(terms[j]) = s(j);
            const auto neg = std::find_if(terms.begin() + 1, terms.end(), [&](int t) { return sol(t) < 0.0; });
            if (neg == terms.end())
                break;
            terms.erase(neg);
        }
        if (!std::isfinite(sol(0)) || !std::isfinite(sol(1)) || !std::isfinite(sol(2)))
            break;
        const double a_new = std::clamp(sol(0), 0.25, 100.0), b2_new = sol(1), c2_new = sol(2);
        const bool settled = std::fabs(a_new - a) <= 1e-4 * a && std::fabs(b2_new - b2) <= 1e-4 * b2
                             && std::fabs(c2_new - c2) <= 1e-4 * c2;
        a = a_new; b2 = b2_new; c2 = c2_new;
        fit.active = true;
        fit.b_resolved = b_measured && b2 > 0.0 && a == sol(0);
        if (settled)
            break;
    }
    fit.a = a; fit.b2 = b2; fit.c2 = c2;
    return fit;
}