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

#include "ReferenceFreeReindex.h"

#include <algorithm>
#include <array>
#include <cstdint>
#include <numbers>
#include <random>
#include <unordered_map>

#include "HKLKey.h"
#include "ReindexAmbiguity.h"
#include "../../common/CrystalLattice.h"
#include "../../common/ParallelFor.h"

namespace {
    constexpr size_t MIN_REFLECTIONS = 20;
    constexpr int RESOLUTION_BINS = 20;
    constexpr int RESTARTS = 4;
    constexpr int MAX_SWEEPS = 100;
    constexpr double DECIDED_SIGMA = 3.0;
    // Two fits closer than this (a 0.1% difference of one axis, squared) are the same fit.
    constexpr double METRIC_TIE = 1e-6;

    double CellDistance(const UnitCell &u, const UnitCell &ref) {
        const auto rel = [](double v, double r) { return (v - r) / r; };
        const auto ang = [](double v, double r) { return (v - r) * std::numbers::pi / 180.0; };
        return rel(u.a, ref.a) * rel(u.a, ref.a) + rel(u.b, ref.b) * rel(u.b, ref.b)
               + rel(u.c, ref.c) * rel(u.c, ref.c) + ang(u.alpha, ref.alpha) * ang(u.alpha, ref.alpha)
               + ang(u.beta, ref.beta) * ang(u.beta, ref.beta) + ang(u.gamma, ref.gamma) * ang(u.gamma, ref.gamma);
    }

    double Median(std::vector<double> v) {
        if (v.empty())
            return NAN;
        std::nth_element(v.begin(), v.begin() + v.size() / 2, v.end());
        return v[v.size() / 2];
    }
}

MetricAnchors AnchorOnCells(const std::vector<IntegrationOutcome> &outcomes, const UnitCell &cell,
                            const std::vector<gemmi::Op> &ops) {
    std::vector<gemmi::Op> indexings{gemmi::Op::identity()};
    indexings.insert(indexings.end(), ops.begin(), ops.end());
    const size_t K = indexings.size();

    // How much better an image's own refined cell fits the run's cell as indexed than after indexing g.
    const auto cell_preference = [&](const IntegrationOutcome &o, size_t g) {
        return CellDistance(ReindexLattice(o.latt, indexings[g]).GetUnitCell(), cell)
               - CellDistance(o.latt.GetUnitCell(), cell);
    };
    // Which indexings the cells decide: those the indexed images' cells prefer the identity over (a tie
    // counting half) for at least METRIC_DECIDED of them.
    MetricAnchors out;
    std::vector<bool> decided(K, false);
    for (size_t g = 1; g < K; ++g) {
        double prefer = 0, n = 0;
        for (const auto &o : outcomes) {
            if (o.reflections.empty())
                continue;
            const double p = cell_preference(o, g);
            prefer += std::fabs(p) < METRIC_TIE ? 0.5 : p > 0 ? 1.0 : 0.0;
            n += 1;
        }
        out.metric_preference.push_back(n > 0 ? prefer / n : 0.5);
        decided[g] = out.metric_preference.back() >= METRIC_DECIDED;
    }
    // Where the cells decide an indexing, an image whose own cell prefers the identity stays as indexed.
    out.allowed.assign(outcomes.size() * K, 1);
    for (size_t o = 0; o < outcomes.size(); ++o)
        for (size_t g = 1; g < K; ++g)
            if (decided[g] && !outcomes[o].reflections.empty() && cell_preference(outcomes[o], g) >= METRIC_TIE)
                out.allowed[o * K + g] = 0;
    return out;
}

ReferenceAnchorResult AnchorReferenceReindex(std::vector<IntegrationOutcome> &outcomes, const UnitCell &cell,
                                             const std::vector<gemmi::Op> &ops) {
    std::vector<gemmi::Op> indexings{gemmi::Op::identity()};
    indexings.insert(indexings.end(), ops.begin(), ops.end());
    const size_t K = indexings.size();
    // The indexing best correlated with the reference among those open: the one as indexed unless another
    // beats it; a score that is not finite never wins (the rule of the per-image resolver).
    const auto best = [&](const std::vector<double> &cc, const uint8_t *open) {
        size_t b = 0;
        for (size_t g = 1; g < K; ++g)
            if (open[g] && std::isfinite(cc[g]) && (!std::isfinite(cc[b]) || cc[g] > cc[b]))
                b = g;
        return b;
    };
    const std::vector<uint8_t> all_open(K, 1);

    // Each image back in the indexing it was indexed in.
    std::vector<size_t> by_reference(outcomes.size(), 0);
    for (size_t o = 0; o < outcomes.size(); ++o) {
        auto &out = outcomes[o];
        if (out.reference_cc.size() != K)
            continue;
        by_reference[o] = best(out.reference_cc, all_open.data());
        if (by_reference[o] != 0)
            ReindexImage(out.reflections, out.latt, indexings[by_reference[o]].inverse());
    }

    const MetricAnchors anchors = AnchorOnCells(outcomes, cell, ops);
    ReferenceAnchorResult result;
    result.metric_preference = anchors.metric_preference;
    for (size_t o = 0; o < outcomes.size(); ++o) {
        auto &out = outcomes[o];
        if (out.reference_cc.size() != K)
            continue;
        const size_t g = best(out.reference_cc, &anchors.allowed[o * K]);
        result.overruled += g != by_reference[o];
        if (g != 0)
            ReindexImage(out.reflections, out.latt, indexings[g]);
        out.reference_cc.clear();   // the scores were for the indexing as indexed, which is no longer known
    }
    return result;
}

ReferenceFreeReindexResult ResolveReindexReferenceFree(std::vector<IntegrationOutcome> &outcomes,
                                                       const DiffractionExperiment &x,
                                                       const UnitCell &cell,
                                                       const std::vector<gemmi::Op> &all_ops,
                                                       size_t nthreads, bool apply) {
    ReferenceFreeReindexResult out;
    out.choice.assign(outcomes.size(), -1);

    if (all_ops.empty())
        return out;
    const auto &s = x.GetScalingSettings();
    const gemmi::SpaceGroup &sg = x.GetSpaceGroupOrP1();
    std::vector<gemmi::Op> indexings{gemmi::Op::identity()};
    indexings.insert(indexings.end(), all_ops.begin(), all_ops.end());
    const int K = static_cast<int>(indexings.size());

    const MetricAnchors anchors = AnchorOnCells(outcomes, cell, all_ops);
    out.metric_preference = anchors.metric_preference;

    const auto accept = [&](const Reflection &r) {
        return !r.on_ice_ring && AcceptReflection(r, s.GetHighResolutionLimit_A(), s.GetLowResolutionLimit_A())
               && r.partiality >= s.GetMinPartiality() && r.partiality > 0.0f
               && std::isfinite(r.I) && std::isfinite(r.sigma) && r.sigma > 0.0f;
    };
    const auto corrected = [](const Reflection &r) {
        return static_cast<double>(r.I) * (r.prescaling_corr * r.qe_corr * r.flight_corr) / r.partiality;
    };

    // What each image has to offer: how many reflections, how fine, and how far its indices reach in
    // any of the indexings (which sizes the index table below).
    const int n_out = static_cast<int>(outcomes.size());
    std::vector<size_t> count(n_out, 0);
    std::vector<std::array<int, 3>> reach(n_out, {0, 0, 0});
    std::vector<double> s3_max(n_out, 0.0);
    ParallelFor(n_out, nthreads, [&](int o) {
        for (const auto &r : outcomes[o].reflections) {
            if (!accept(r))
                continue;
            ++count[o];
            s3_max[o] = std::max(s3_max[o], 1.0 / (static_cast<double>(r.d) * r.d * r.d));
            for (const auto &op : indexings) {
                const auto h = op.apply_to_hkl({{r.h, r.k, r.l}});
                for (int i = 0; i < 3; ++i)
                    reach[o][i] = std::max(reach[o][i], std::abs(h[i]));
            }
        }
    });

    // The crystals: images with enough reflections, in image order.
    std::vector<int> crystal_outcome;
    std::vector<size_t> begin{0};
    std::array<int, 3> bound{0, 0, 0};
    double s3_top = 0.0;
    for (int o = 0; o < n_out; ++o) {
        if (count[o] < MIN_REFLECTIONS)
            continue;
        crystal_outcome.push_back(o);
        begin.push_back(begin.back() + count[o]);
        for (int i = 0; i < 3; ++i)
            bound[i] = std::max(bound[i], reach[o][i]);
        s3_top = std::max(s3_top, s3_max[o]);
    }
    const int N = static_cast<int>(crystal_outcome.size());
    if (N < 2)
        return out;

    // The indexings each crystal may take (AnchorOnCells): a crystal its own cell anchors is one of those
    // the others are decided against.
    std::vector<uint8_t> allowed(static_cast<size_t>(N) * K);
    std::vector<bool> free(N, false);
    for (int c = 0; c < N; ++c)
        for (int g = 0; g < K; ++g) {
            allowed[c * K + g] = anchors.allowed[crystal_outcome[c] * K + g];
            free[c] = free[c] || (g > 0 && allowed[c * K + g]);
        }
    out.crystals = std::count(free.begin(), free.end(), true);
    if (out.crystals == 0)
        return out;
    const size_t nnz = begin.back();

    // The mean intensity in shells of equal reciprocal volume over the whole run, so that a crystal's
    // intensities are compared after the fall-off with resolution, which every indexing shares, is out.
    const auto bin_of = [&](const Reflection &r) {
        const double s3 = 1.0 / (static_cast<double>(r.d) * r.d * r.d);
        return std::min(RESOLUTION_BINS - 1, static_cast<int>(RESOLUTION_BINS * s3 / s3_top));
    };
    std::vector<std::array<double, 2 * RESOLUTION_BINS>> bin_part(N);
    ParallelFor(N, nthreads, [&](int c) {
        auto &b = bin_part[c];
        b.fill(0.0);
        for (const auto &r : outcomes[crystal_outcome[c]].reflections)
            if (accept(r)) {
                b[bin_of(r)] += corrected(r);
                b[RESOLUTION_BINS + bin_of(r)] += 1.0;
            }
    });
    std::array<double, RESOLUTION_BINS> bin_mean{};
    {
        std::array<double, 2 * RESOLUTION_BINS> total{};
        for (const auto &b : bin_part)   // in crystal order: the same sum at any thread count
            for (int i = 0; i < 2 * RESOLUTION_BINS; ++i)
                total[i] += b[i];
        for (int i = 0; i < RESOLUTION_BINS; ++i)
            bin_mean[i] = total[RESOLUTION_BINS + i] > 0 ? total[i] / total[RESOLUTION_BINS + i] : 1.0;
        if (std::any_of(bin_mean.begin(), bin_mean.end(), [](double m) { return !(m > 0.0); }))
            for (auto &m : bin_mean)
                m = 1.0;   // a shell with a non-positive mean: compare raw intensities instead
    }

    // Every index any indexing can produce, mapped to a dense number of its reflection in the
    // asymmetric unit (Friedel mates together: the anomalous difference would only add noise here).
    const HKLKeyGenerator key_gen(/*merge_friedel=*/true, sg);
    const int dh = 2 * bound[0] + 1, dk = 2 * bound[1] + 1, dl = 2 * bound[2] + 1;
    const auto entry = [&](int h, int k, int l) {
        return (static_cast<size_t>(h + bound[0]) * dk + (k + bound[1])) * dl + (l + bound[2]);
    };
    std::vector<uint64_t> table_key(static_cast<size_t>(dh) * dk * dl);
    ParallelFor(dh, nthreads, [&](int ih) {
        for (int k = -bound[1]; k <= bound[1]; ++k)
            for (int l = -bound[2]; l <= bound[2]; ++l)
                table_key[entry(ih - bound[0], k, l)] = key_gen(ih - bound[0], k, l).pack();
    });
    std::vector<int32_t> table(table_key.size());
    int32_t n_keys = 0;
    {
        std::unordered_map<uint64_t, int32_t> dense;
        for (size_t i = 0; i < table_key.size(); ++i) {
            const auto [it, inserted] = dense.try_emplace(table_key[i], n_keys);
            n_keys += inserted;
            table[i] = it->second;
        }
    }
    table_key = {};

    // Each crystal's intensities, divided by the shell mean and standardised within the crystal (the
    // unknown scale and offset of an image drop out), and where each lands in every indexing. Each is
    // weighted by how much of it is signal: a standardised intensity has unit variance, of which its
    // measurement error takes (sigma / sd)^2, so a weak or barely recorded reflection counts for little.
    std::vector<float> z(nnz);
    std::vector<std::vector<int32_t>> key(K, std::vector<int32_t>(nnz));
    ParallelFor(N, nthreads, [&](int c) {
        std::vector<double> err;
        size_t j = begin[c];
        double sum = 0, sum2 = 0;
        for (const auto &r : outcomes[crystal_outcome[c]].reflections) {
            if (!accept(r))
                continue;
            const double v = corrected(r) / bin_mean[bin_of(r)];
            z[j] = static_cast<float>(v);
            err.push_back(r.sigma * (r.prescaling_corr * r.qe_corr * r.flight_corr) / r.partiality
                          / bin_mean[bin_of(r)]);
            sum += v;
            sum2 += v * v;
            for (int g = 0; g < K; ++g) {
                const auto h = indexings[g].apply_to_hkl({{r.h, r.k, r.l}});
                key[g][j] = table[entry(h[0], h[1], h[2])];
            }
            ++j;
        }
        const double n = static_cast<double>(begin[c + 1] - begin[c]);
        const double mean = sum / n;
        const double sd = std::sqrt(std::max(0.0, sum2 / n - mean * mean));
        for (size_t i = begin[c]; i < begin[c + 1]; ++i) {
            const double e = err[i - begin[c]] / sd;
            z[i] = sd > 0 ? static_cast<float>((z[i] - mean) / sd / (1.0 + e * e)) : 0.0f;
        }
    });
    table = {};

    // Following Brehm & Diederichs (2014) Acta Cryst. D70, 101-109: each crystal in turn takes the
    // indexing in which it agrees best with all the others, until none changes. Agreement is the sum
    // over all pairs of crystals of the products of their standardised intensities on common
    // reflections, which is the dot product of a crystal with the sum of all the others - so a pass over
    // the crystals costs one pass over the reflections, not one per pair of crystals. The sum of all
    // pairs only ever grows, so it ends; a few starts are made and the one that ends highest is kept.
    std::vector<double> M(n_keys);
    const auto add = [&](int c, int g, double sign) {
        for (size_t i = begin[c]; i < begin[c + 1]; ++i)
            M[key[g][i]] += sign * z[i];
    };
    const auto dot = [&](int c, int g) {
        double d = 0;
        for (size_t i = begin[c]; i < begin[c + 1]; ++i)
            d += z[i] * M[key[g][i]];
        return d;
    };
    std::vector<int> best_choice;
    double best_objective = -1.0;
    for (int restart = 0; restart < RESTARTS; ++restart) {
        std::vector<int> choice(N, 0);
        std::mt19937 rng(restart);
        if (restart > 0)
            for (int c = 0; c < N; ++c) {
                const int g = static_cast<int>(rng() % K);
                choice[c] = allowed[c * K + g] ? g : 0;
            }
        std::fill(M.begin(), M.end(), 0.0);
        for (int c = 0; c < N; ++c)
            add(c, choice[c], 1.0);
        int sweeps = 0;
        // Run until a pass moves no more than one crystal in a thousand: what still moves then are crystals
        // too weak to be decided either way.
        for (int changed = N; changed > N / 1000 && sweeps < MAX_SWEEPS; ++sweeps) {
            changed = 0;
            for (int c = 0; c < N; ++c) {
                add(c, choice[c], -1.0);
                int best = choice[c];
                double best_dot = dot(c, best);
                for (int g = 0; g < K; ++g) {
                    if (g == choice[c] || !allowed[c * K + g])
                        continue;
                    const double d = dot(c, g);
                    if (d > best_dot) {
                        best_dot = d;
                        best = g;
                    }
                }
                changed += best != choice[c];
                choice[c] = best;
                add(c, best, 1.0);
            }
        }
        double objective = 0;
        for (double m : M)
            objective += m * m;
        if (objective > best_objective) {
            best_objective = objective;
            best_choice = choice;
            out.sweeps = sweeps;
        }
    }

    // How clearly each crystal is decided: its CC with the mean of all the other crystals, in the chosen
    // indexing and in the best other one, and the difference of the two in sigma (Fisher's z; each CC
    // on n reflections has an error of 1/sqrt(n - 3), and the two are nearly independent).
    std::fill(M.begin(), M.end(), 0.0);
    std::vector<double> C(n_keys, 0.0);
    for (int c = 0; c < N; ++c)
        for (size_t i = begin[c]; i < begin[c + 1]; ++i) {
            M[key[best_choice[c]][i]] += z[i];
            C[key[best_choice[c]][i]] += 1.0;
        }
    std::vector<double> sigma(N, NAN), cc_chosen(N, NAN), cc_other(N, NAN);
    ParallelFor(N, nthreads, [&](int c) {
        std::vector<double> cc(K, NAN);
        std::vector<double> n_used(K, 0.0);
        for (int g = 0; g < K; ++g) {
            double sx = 0, sy = 0, sxx = 0, syy = 0, sxy = 0, n = 0;
            for (size_t i = begin[c]; i < begin[c + 1]; ++i) {
                const int32_t k = key[g][i];
                const double own = g == best_choice[c] ? 1.0 : 0.0;
                const double cnt = C[k] - own;
                if (cnt <= 0)
                    continue;
                const double xv = z[i], yv = (M[k] - own * z[i]) / cnt;
                sx += xv; sy += yv; sxx += xv * xv; syy += yv * yv; sxy += xv * yv; n += 1;
            }
            const double vx = sxx - sx * sx / n, vy = syy - sy * sy / n;
            if (n > 3 && vx > 0 && vy > 0) {
                cc[g] = std::clamp((sxy - sx * sy / n) / std::sqrt(vx * vy), -0.999, 0.999);
                n_used[g] = n;
            }
        }
        const int g0 = best_choice[c];
        int other = -1;
        for (int g = 0; g < K; ++g)
            if (g != g0 && allowed[c * K + g] && std::isfinite(cc[g]) && (other < 0 || cc[g] > cc[other]))
                other = g;
        if (!std::isfinite(cc[g0]) || other < 0)
            return;
        cc_chosen[c] = cc[g0];
        cc_other[c] = cc[other];
        sigma[c] = (std::atanh(cc[g0]) - std::atanh(cc[other]))
                   * std::sqrt((std::min(n_used[g0], n_used[other]) - 3.0) / 2.0);
    });
    std::vector<double> sig, ccc, cco;
    for (int c = 0; c < N; ++c)
        if (free[c] && std::isfinite(sigma[c])) {
            sig.push_back(sigma[c]);
            ccc.push_back(cc_chosen[c]);
            cco.push_back(cc_other[c]);
        }
    out.median_sigma = Median(sig);
    out.cc_chosen = Median(ccc);
    out.cc_other = Median(cco);
    if (!sig.empty())
        out.decided_fraction = static_cast<double>(std::count_if(sig.begin(), sig.end(), [](double v) {
            return v > DECIDED_SIGMA; })) / static_cast<double>(sig.size());

    // Any consistent answer relabelled as a whole by one of the indexings is as consistent: take the
    // one that leaves the most crystals where they are - unless the cells decide an indexing, in which
    // case the crystals they anchor fix the whole. Two indexings are the same for the merge when they
    // send a general reflection to the same reflection of the asymmetric unit.
    const std::array<int, 3> general{2, 5, 11};
    const auto same = [&](const gemmi::Op::Miller &a, const gemmi::Op::Miller &b) {
        return key_gen(a[0], a[1], a[2]) == key_gen(b[0], b[1], b[2]);
    };
    const auto composed = [&](int global, int g) {   // indexing g, then the overall relabel
        return indexings[global].apply_to_hkl(indexings[g].apply_to_hkl(general));
    };
    int global = 0;
    size_t global_kept = 0;
    const bool anchored = std::any_of(out.metric_preference.begin(), out.metric_preference.end(),
                                      [](double p) { return p >= METRIC_DECIDED; });
    for (int h = 0; h < (anchored ? 1 : K); ++h) {
        size_t kept = 0;
        for (int c = 0; c < N; ++c)
            kept += same(composed(h, best_choice[c]), general);
        if (kept > global_kept) {
            global_kept = kept;
            global = h;
        }
    }
    for (int c = 0; c < N; ++c) {
        int final_g = best_choice[c];
        for (int g = 0; g < K; ++g)
            if (same(composed(global, best_choice[c]), indexings[g].apply_to_hkl(general))) {
                final_g = g;
                break;
            }
        out.choice[crystal_outcome[c]] = final_g;
        if (final_g == 0)
            continue;
        ++out.reindexed;
        if (apply) {
            auto &o = outcomes[crystal_outcome[c]];
            ReindexImage(o.reflections, o.latt, indexings[final_g]);
        }
    }
    return out;
}