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

#include "RfreeFlags.h"

#include <algorithm>
#include <unordered_map>
#include <unordered_set>

#include "HKLKey.h"
#include "gemmi/twin.hpp"

namespace {
    // splitmix64 bit-mix of a key -> uniform double in [0, 1). Same key -> same value, so all
    // mates of a reflection (which share the Laue-ASU key) get the same draw. Same idiom as the
    // CC1/2 half-set split (HalfForImage in Merge.cpp).
    double UniformFromKey(uint64_t key) {
        uint64_t z = key + 0x9e3779b97f4a7c15ULL;
        z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
        z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
        z = z ^ (z >> 31);
        return static_cast<double>(z >> 11) * (1.0 / 9007199254740992.0);
    }


    // The rotations of the lattice's own point group - its holohedry - found from the metric as for
    // twin laws (Le Page two-folds, with phenix.xtriage's default obliquity of 3 degrees, so a
    // pseudo-merohedral lattice counts as its higher metric; grouping too much only clusters the
    // free set, grouping too little splits twin mates). The centring is taken as P: the lattice of
    // the cell's own basis vectors, the same whatever group a merge is in, so the merged file, the P1
    // cross-check and any re-merge get one holohedry. On an R lattice in its hexagonal cell that is
    // 6/mmm rather than -3m, which groups each present reflection with exactly its -3m mates, the
    // extra ones being centring-absent. Empty when the cell does not carry the merging group itself
    // (a space group forced on a metric that does not have it): there is no holohedry to key on.
    std::vector<gemmi::Op> LatticeHolohedry(const UnitCell &cell, const gemmi::SpaceGroup &space_group) {
        std::vector<gemmi::Op> rots;
        for (const auto &op : gemmi::find_lattice_symmetry(static_cast<gemmi::UnitCell>(cell), 'P', 3.0).sym_ops)
            rots.push_back(gemmi::Op{op.rot, {0, 0, 0}, 'x'});
        for (const auto &op : space_group.operations().sym_ops)
            if (std::none_of(rots.begin(), rots.end(), [&](const gemmi::Op &r) { return r.rot == op.rot; }))
                return {};
        return rots;
    }

    uint64_t PackIndex(const gemmi::Op::Miller &h) {
        constexpr int64_t bias = 1 << 20;
        return (static_cast<uint64_t>(h[0] + bias) << 42) | (static_cast<uint64_t>(h[1] + bias) << 21)
               | static_cast<uint64_t>(h[2] + bias);
    }
}

void AssignRfreeFlags(std::vector<MergedReflection> &merged, const gemmi::SpaceGroup &space_group,
                      double rfree_fraction, int min_free_reflections, const std::optional<UnitCell> &cell) {
    for (auto &r : merged)
        r.rfree_flag = false;
    if (rfree_fraction <= 0.0 || merged.empty())
        return;

    // The flag is a pure function of the Friedel-merged (Laue) ASU key: symmetry- and Friedel-
    // equivalent reflections collapse to one key and so share a flag (a Bijvoet pair I(+)/I(-) is
    // never split across the work and free sets), and the draw depends only on the reflection index
    // - not on this dataset's resolution range or which reflections it happens to contain. So every
    // dataset of one crystal form gets the SAME free set, which is what a multi-dataset campaign
    // (ensemble refinement, PanDDA) needs. A uniform hash draws ~rfree_fraction of the distinct
    // reflections free; a stratified per-shell draw would be tied to the dataset and break that.
    //
    // The key is the orbit of the reflection under the LATTICE HOLOHEDRY, not under the merging group.
    // A twin law is a lattice symmetry that the crystal lacks, so keying on the merging group puts a
    // free reflection's twin mate in the working set almost every time (measured: 97-98% of the free
    // reflections that have a mate), and a twin-refined R-free then reads the working set through
    // I_calc. Keyed on the holohedry, twin mates share a flag, and the free set no longer depends on
    // the space group a file is merged in - the merged MTZ, the P1 cross-check and a re-merge in any
    // subgroup carry one free set (where the small-data floor below lifts the fraction differently in
    // two of them, the smaller set is contained in the larger: each orbit has one draw). As
    // phenix.refine does by default (use_lattice_symmetry).
    const HKLKeyGenerator laue_key(/*merge_friedel=*/true, space_group);
    const std::vector<gemmi::Op> holohedry = cell ? LatticeHolohedry(*cell, space_group) : std::vector<gemmi::Op>{};
    auto key_of = [&](const MergedReflection &r) -> uint64_t {
        if (holohedry.empty())
            return laue_key(r).pack();
        // The largest index of the orbit, Friedel mates included, names it.
        gemmi::Op::Miller best{{r.h, r.k, r.l}};
        for (const auto &op : holohedry) {
            const gemmi::Op::Miller h = op.apply_to_hkl({{r.h, r.k, r.l}});
            best = std::max({best, h, gemmi::Op::Miller{{-h[0], -h[1], -h[2]}}});
        }
        return PackIndex(best);
    };

    // Count the distinct test-eligible reflections (distinct Laue-ASU keys of the merging group; mates
    // collapse to one) so the fraction can be floored to a usable test-set size on small data. Counted
    // in the merging group, not in holohedral orbits: R-free is a sum over the reflections of this file,
    // and counting orbits would lift the fraction on every dataset below its holohedry.
    std::unordered_set<uint64_t> distinct;
    distinct.reserve(merged.size());
    for (const auto &r : merged)
        distinct.insert(laue_key(r).pack());

    // Effective fraction: at least rfree_fraction, lifted toward min_free_reflections/N on small data
    // (so R-free is not sampling-noise dominated), but the floor's lift is capped at MAX_FRACTION so a
    // large test set never steals working data. An explicitly large rfree_fraction is always honoured.
    constexpr double MAX_FRACTION = 0.10;
    const double floor_fraction =
        std::min(min_free_reflections / static_cast<double>(distinct.size()), MAX_FRACTION);
    const double eff_fraction = std::max(rfree_fraction, floor_fraction);

    for (auto &r : merged)
        r.rfree_flag = UniformFromKey(key_of(r)) < eff_fraction;
}

size_t ApplyReferenceFreeFlags(std::vector<MergedReflection> &merged, const gemmi::SpaceGroup &space_group,
                               const std::vector<MergedReflection> &reference) {
    // Reference free/work partition keyed by the Friedel-merged (Laue) ASU index, so it transfers
    // regardless of which Bijvoet mate / symmetry equivalent each dataset happens to have measured.
    const HKLKeyGenerator laue_key(/*merge_friedel=*/true, space_group);
    std::unordered_map<uint64_t, bool> ref_flag;
    ref_flag.reserve(reference.size());
    for (const auto &r : reference)
        ref_flag[laue_key(r).pack()] = r.rfree_flag;

    size_t matched = 0;
    for (auto &r : merged) {
        const auto it = ref_flag.find(laue_key(r).pack());
        if (it != ref_flag.end()) {          // reflections absent from the reference keep their hash flag
            r.rfree_flag = it->second;
            ++matched;
        }
    }
    return matched;
}