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

#include "CentreOfSymmetry.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <map>
#include <random>

namespace {
    using Hkl = std::array<int, 3>;

    struct Refl {
        Hkl hkl;
        double I, sigma, d;
    };

    // Bounds of the decision, from the readings of untwinned small-molecule data processed by this
    // program: centric sets read f = 0.74-1.03 on <|L|>, acentric ones -0.22-0.27; the other two
    // statistics separate less well (one acentric set reads <|E^2-1|> fully centric on outliers), so
    // they only have to agree, not decide.
    constexpr double MAX_F_L_ACENTRIC = 0.3;
    constexpr double MAX_F_OTHER_ACENTRIC = 0.5;
    // The range a Wilson population can produce, with room for noise; outside it the statistic does not
    // describe the structure. A perfect twin of acentric data reads f ~ -1 on <|L|>, which is the twin
    // detection's business, not this; a hypercentric structure reads well above 1.
    constexpr double MIN_F_PHYSICAL = -0.3;
    constexpr double MAX_F_PHYSICAL = 1.3;
    constexpr double MIN_F_CONTROL_CENTRIC = 0.5;
    constexpr int MIN_REFLECTIONS = 50;
    constexpr int REFLECTIONS_PER_SHELL = 200;
    constexpr int SIMULATION_REPEATS = 10;

    struct Stats {
        double e2m1 = 0.0, nz01 = 0.0;
    };

    // E^2 = I / <I> in equal-count resolution shells, then <|E^2-1|> and N(0.1).
    // `order` lists the class's reflections by resolution, so the shells are the same for every call.
    Stats ClassStats(const std::vector<double>& I, const std::vector<size_t>& order) {
        const size_t n = order.size();
        const size_t n_shells = std::max<size_t>(1, n / REFLECTIONS_PER_SHELL);
        double sum_e2m1 = 0.0;
        size_t below = 0;
        for (size_t s = 0; s < n_shells; ++s) {
            const size_t lo = s * n / n_shells, hi = (s + 1) * n / n_shells;
            double mean = 0.0;
            for (size_t j = lo; j < hi; ++j)
                mean += I[order[j]];
            mean /= static_cast<double>(hi - lo);
            for (size_t j = lo; j < hi; ++j) {
                const double e2 = mean > 0.0 ? I[order[j]] / mean : 0.0;
                sum_e2m1 += std::fabs(e2 - 1.0);
                if (e2 < 0.1)
                    ++below;
            }
        }
        return {sum_e2m1 / static_cast<double>(n), static_cast<double>(below) / static_cast<double>(n)};
    }

    double MeanAbsL(const std::vector<double>& I, const std::vector<std::pair<size_t, size_t>>& pairs) {
        double sum = 0.0;
        size_t n = 0;
        for (const auto& [a, b] : pairs) {
            const double s = I[a] + I[b];
            if (s > 0.0) {
                sum += std::min(1.0, std::fabs((I[a] - I[b]) / s));
                ++n;
            }
        }
        return n > 0 ? sum / static_cast<double>(n) : 0.0;
    }

    // Uniform, exponential and normal deviates written out rather than taken from <random>'s
    // distributions, whose output the standard leaves to the implementation: the decision this feeds
    // must be the same on every platform. mt19937_64 itself is fully specified.
    struct Deviates {
        std::mt19937_64 engine{20261004};
        double Uniform() { return (static_cast<double>(engine() >> 11) + 0.5) * 0x1.0p-53; }
        double Exponential() { return -std::log(Uniform()); }
        double Normal() {   // Box-Muller, one of the pair
            return std::sqrt(-2.0 * std::log(Uniform())) * std::cos(2.0 * 3.14159265358979323846 * Uniform());
        }
    };

    CentreStatistic Calibrated(double measured, double acentric, double centric) {
        return {measured, acentric, centric,
                centric != acentric ? (measured - acentric) / (centric - acentric) : 0.0};
    }
}

// Following Howells, Phillips & Rogers (1950) Acta Cryst. 3, 210-214 (N(z)), and Padilla & Yeates
// (2003) Acta Cryst. D59, 1124-1130 (the L-test, read here against its centric value 2/pi).
CentreOfSymmetryResult AnalyzeCentreOfSymmetry(const std::vector<MergedReflection>& merged,
                                               const gemmi::SpaceGroup& space_group) {
    CentreOfSymmetryResult out;
    const gemmi::GroupOps gops = space_group.operations();

    // The Laue class: every proper rotation of the group, with and without the inversion.
    std::vector<gemmi::Op> proper, laue;
    for (const auto& op : gops.derive_symmorphic().sym_ops) {
        const gemmi::Op r = op.det_rot() > 0 ? gemmi::Op{op.rot, {0, 0, 0}}
                                              : gemmi::Op{op.negated_rot(), {0, 0, 0}};
        if (std::none_of(proper.begin(), proper.end(), [&](const gemmi::Op& p) { return p.rot == r.rot; }))
            proper.push_back(r);
    }
    for (const auto& r : proper) {
        laue.push_back(r);
        laue.push_back(gemmi::Op{r.negated_rot(), {0, 0, 0}});
    }

    // Fold onto the Laue class: one weighted mean per orbit, keyed by the orbit's largest member.
    struct Acc { double w = 0.0, wI = 0.0, d = 0.0; };
    std::map<Hkl, Acc> folded;
    for (const auto& r : merged) {
        if (!std::isfinite(r.I) || !std::isfinite(r.sigma) || !(r.sigma > 0.0f) || !(r.d > 0.0f))
            continue;
        const gemmi::Op::Miller hkl{{r.h, r.k, r.l}};
        if (gops.is_systematically_absent(hkl))
            continue;
        Hkl key{r.h, r.k, r.l};
        for (const auto& op : laue)
            key = std::max(key, Hkl(op.apply_to_hkl(hkl)));
        auto& a = folded[key];
        const double w = 1.0 / (static_cast<double>(r.sigma) * r.sigma);
        a.w += w;
        a.wI += w * r.I;
        a.d = r.d;
    }
    std::vector<Refl> refl;
    for (const auto& [hkl, a] : folded)
        refl.push_back({hkl, a.wI / a.w, 1.0 / std::sqrt(a.w), a.d});
    if (refl.size() < static_cast<size_t>(MIN_REFLECTIONS))
        return out;

    // Low resolution first, and the high-resolution end cut where 20 equal-count shells first fall
    // below <I/sigma> 3: past that the noise, not the structure, sets the distribution.
    std::sort(refl.begin(), refl.end(), [](const Refl& a, const Refl& b) {
        return a.d != b.d ? a.d > b.d : a.hkl < b.hkl;
    });
    size_t keep = refl.size();
    for (size_t s = 0; s < 20; ++s) {
        const size_t lo = s * refl.size() / 20, hi = (s + 1) * refl.size() / 20;
        double sum = 0.0;
        for (size_t j = lo; j < hi; ++j)
            sum += refl[j].I / refl[j].sigma;
        if (hi > lo && sum / static_cast<double>(hi - lo) < 3.0) {
            keep = lo + 1;
            break;
        }
    }
    refl.resize(keep);
    out.d_min_A = refl.back().d;

    // The two classes, in resolution order.
    std::vector<size_t> general, control;
    std::map<Hkl, size_t> general_index;
    for (size_t i = 0; i < refl.size(); ++i) {
        const gemmi::Op::Miller hkl{{refl[i].hkl[0], refl[i].hkl[1], refl[i].hkl[2]}};
        const gemmi::Op::Miller minus{{-hkl[0], -hkl[1], -hkl[2]}};
        int fixed_by = 0;
        for (const auto& op : laue)
            if (op.apply_to_hkl(hkl) == hkl)
                ++fixed_by;
        bool to_friedel_mate = false;
        for (const auto& op : proper)
            if (op.apply_to_hkl(hkl) == minus)
                to_friedel_mate = true;
        if (fixed_by == 1) {
            general_index[refl[i].hkl] = general.size();
            general.push_back(i);
        } else if (to_friedel_mate) {
            control.push_back(i);
        }
    }
    // L-test pairs: general reflections a step of 0 or 2 apart in each index, the even steps keeping a
    // half-integer pseudo-translation from correlating the pair.
    std::vector<std::pair<size_t, size_t>> pairs;
    for (size_t g : general)
        for (int a = -2; a <= 2; a += 2)
            for (int b = -2; b <= 2; b += 2)
                for (int c = -2; c <= 2; c += 2) {
                    if (Hkl{a, b, c} <= Hkl{0, 0, 0})
                        continue;
                    const Hkl& h = refl[g].hkl;
                    const auto it = general_index.find(Hkl{h[0] + a, h[1] + b, h[2] + c});
                    if (it != general_index.end())
                        pairs.emplace_back(g, general[it->second]);
                }
    out.n_general = static_cast<int>(general.size());
    out.n_control = static_cast<int>(control.size());
    out.l_pairs = static_cast<int>(pairs.size());
    if (out.n_general < MIN_REFLECTIONS || out.n_control < MIN_REFLECTIONS || out.l_pairs < MIN_REFLECTIONS)
        return out;

    std::vector<double> I(refl.size()), sigma(refl.size());
    for (size_t i = 0; i < refl.size(); ++i) {
        I[i] = refl[i].I;
        sigma[i] = refl[i].sigma;
    }
    // The true mean each simulated intensity is drawn about: the measured shell mean of its class.
    std::vector<double> mean_true(refl.size(), 0.0);
    for (const auto* cls : {&general, &control}) {
        const size_t n = cls->size();
        const size_t n_shells = std::max<size_t>(1, n / REFLECTIONS_PER_SHELL);
        for (size_t s = 0; s < n_shells; ++s) {
            const size_t lo = s * n / n_shells, hi = (s + 1) * n / n_shells;
            double mean = 0.0;
            for (size_t j = lo; j < hi; ++j)
                mean += I[(*cls)[j]];
            mean = std::max(mean / static_cast<double>(hi - lo), 1e-9);
            for (size_t j = lo; j < hi; ++j)
                mean_true[(*cls)[j]] = mean;
        }
    }

    const Stats g = ClassStats(I, general), c = ClassStats(I, control);
    const double l = MeanAbsL(I, pairs);

    Deviates dev;
    Stats sim_g[2], sim_c[2];
    double sim_l[2] = {0.0, 0.0};
    std::vector<double> I_sim(refl.size());
    for (int centric = 0; centric < 2; ++centric)
        for (int rep = 0; rep < SIMULATION_REPEATS; ++rep) {
            for (size_t i = 0; i < refl.size(); ++i) {
                double x;
                if (centric) {
                    const double z = dev.Normal();
                    x = z * z;
                } else {
                    x = dev.Exponential();
                }
                I_sim[i] = mean_true[i] * x + dev.Normal() * sigma[i];
            }
            const Stats sg = ClassStats(I_sim, general), sc = ClassStats(I_sim, control);
            sim_g[centric].e2m1 += sg.e2m1 / SIMULATION_REPEATS;
            sim_g[centric].nz01 += sg.nz01 / SIMULATION_REPEATS;
            sim_c[centric].e2m1 += sc.e2m1 / SIMULATION_REPEATS;
            sim_l[centric] += MeanAbsL(I_sim, pairs) / SIMULATION_REPEATS;
        }

    out.mean_abs_e2_minus_1 = Calibrated(g.e2m1, sim_g[0].e2m1, sim_g[1].e2m1);
    out.n_z_01 = Calibrated(g.nz01, sim_g[0].nz01, sim_g[1].nz01);
    out.mean_abs_l = Calibrated(l, sim_l[0], sim_l[1]);
    out.control_mean_abs_e2_minus_1 = Calibrated(c.e2m1, sim_c[0].e2m1, sim_c[1].e2m1);

    const auto physical = [](const CentreStatistic& s) { return s.f >= MIN_F_PHYSICAL && s.f <= MAX_F_PHYSICAL; };
    out.applicable = physical(out.mean_abs_e2_minus_1) && physical(out.n_z_01) && physical(out.mean_abs_l)
                     && out.control_mean_abs_e2_minus_1.f >= MIN_F_CONTROL_CENTRIC;
    out.reads_acentric = out.applicable
                         && out.mean_abs_l.f <= MAX_F_L_ACENTRIC
                         && out.mean_abs_e2_minus_1.f <= MAX_F_OTHER_ACENTRIC
                         && out.n_z_01.f <= MAX_F_OTHER_ACENTRIC;
    return out;
}