Jungfraujoch/image_analysis/scale_merge/ReindexAmbiguity.h

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

#pragma once

#include <functional>
#include <cmath>
#include <map>
#include <optional>
#include <span>
#include <vector>

#include "gemmi/symmetry.hpp"

#include "HKLKey.h"
#include "../../common/DiffractionExperiment.h"
#include "../../common/Reflection.h"
#include "../../common/UnitCell.h"
#include "../../common/CrystalLattice.h"

// Resolving the indexing ("alternative indexing" / reindexing) ambiguity: when the lattice symmetry
// is higher than the crystal Laue symmetry (merohedral / pseudo-merohedral), a crystal can be validly
// indexed in several ways related by reindexing operators. The right one is picked by agreement with a
// reference (reference-MTZ intensities or a model), so no reference-free ambigator (cosym) is needed.
//
// This is distinct from the screw-axis / enantiomorph ambiguity (P4_1 vs P4_3): that leaves the merged
// intensities unchanged, so it is never resolved here — it is adopted from the reference/model.

// The candidate reindexing operators = twin laws = cosets of the crystal point group in the lattice
// symmetry. Empty when the lattice symmetry equals the crystal Laue symmetry (holohedral crystal),
// i.e. there is no indexing ambiguity. Identity is NOT included in the returned list.
std::vector<gemmi::Op> ReindexAmbiguityOperators(const UnitCell &cell, const gemmi::SpaceGroup &space_group,
                                                 double max_obliquity_deg = 2.0);

// A copy of `merged` with each Miller index transformed by `op` (the intensities are unchanged).
// The indices are left where `op` puts them, which is what the scoring callers want - they reduce to
// the ASU themselves.
// The lattice an image has once its indices are moved by `op` (h' = h R): the reciprocal axes follow as
// R^-1 times the old ones, so h' (a*, b*, c*)' is the same reciprocal vector as h (a*, b*, c*).
CrystalLattice ReindexLattice(const CrystalLattice &latt, const gemmi::Op &op);

std::vector<MergedReflection> ReindexReflections(const std::vector<MergedReflection> &merged,
                                                 const gemmi::Op &op);

// As above, but the reindexed reflections are brought back into the ASU of `space_group` -
// what a reflection file has to be written in. Where the reduction reaches a reflection through its
// Friedel mate, the Bijvoet halves travel with it. `merge_friedel` is the merge's own setting, since
// it fixes where the minus hand is stored (at -hkl when the mates are kept apart).
std::vector<MergedReflection> ReindexMergedIntoAsu(const std::vector<MergedReflection> &merged,
                                                   const gemmi::Op &op, const gemmi::SpaceGroup &space_group,
                                                   bool merge_friedel);

struct ReindexChoice {
    gemmi::Op op = gemmi::Op::identity();   // the winning operator (identity = keep the current indexing)
    bool is_identity = true;
    double score = 0.0;                     // score of the winning operator
    double identity_score = 0.0;            // score with no reindex, for comparison / logging
    int n_candidates = 1;                   // identity + twin laws that were tried
};

// Pick the reindexing operator (identity or one of the twin laws) that MAXIMIZES `score(reindexed)`.
// `score` is higher-is-better (e.g. correlation against a reference, or minus the R-free against a
// model). Returns identity when there is no ambiguity or identity scores best.
ReindexChoice ChooseReindex(const std::vector<MergedReflection> &merged,
                            const UnitCell &cell, const gemmi::SpaceGroup &space_group,
                            const std::function<double(const std::vector<MergedReflection> &)> &score,
                            double max_obliquity_deg = 2.0);

// Ready-made scorer: Pearson correlation of merged intensities against a reference set, matched by the
// Friedel-merged (Laue) ASU index in `space_group`. Returns NaN when too few reflections match.
double ReferenceIntensityCC(const std::vector<MergedReflection> &merged,
                            const std::vector<MergedReflection> &reference,
                            const gemmi::SpaceGroup &space_group);

// Which axes a reference MTZ is written on is not the run's to choose: its free-R flags and its
// intensities are keyed by its own indices. So the merge is put into the reference's frame before
// anything is matched against it. The candidates are every change of basis that carries the data cell
// onto the reference cell (an axis permutation, a sign flip, I- against C-centred; CellMappingOperators),
// each also combined with the data's own twin laws (the alternative indexings of the same metric), and
// the identity. Each is scored by ReferenceIntensityCC in the reference's group and the best is taken;
// the identity wins ties. `cob` is the change of basis from the data's frame to the reference's.
struct ReferenceSettingChoice {
    gemmi::Op cob = gemmi::Op::identity();
    double cc = NAN;               // CC with the reference under `cob`
    double identity_cc = NAN;      // ... and with the data left where they are
    int n_candidates = 1;
};
ReferenceSettingChoice ChooseReferenceSetting(const std::vector<MergedReflection> &merged,
                                              const UnitCell &data_cell,
                                              const gemmi::SpaceGroup &data_space_group,
                                              const std::vector<MergedReflection> &reference,
                                              const std::optional<UnitCell> &reference_cell,
                                              const gemmi::SpaceGroup &reference_space_group);

// How well a merge, in the reference's frame, matches the reference: the intensity CC and the
// fraction of the merged reflections inside the reference's resolution range that it holds.
struct ReferenceMatch {
    double cc = NAN;
    double matched_fraction = 0.0;
};
ReferenceMatch MatchReference(const std::vector<MergedReflection> &merged,
                              const std::vector<MergedReflection> &reference,
                              const gemmi::SpaceGroup &space_group);

// Below either bound the reference is not the same crystal form on the same axes, and its free-R
// flags would mark reflections unrelated to its own test set - they are then not inherited. A
// reference in the right frame reads a CC of 0.8-0.95 and matches nearly everything in its range;
// one on permuted axes pairs reflections of different resolution and reads far lower.
constexpr double MIN_REFERENCE_CC = 0.5;
constexpr double MIN_REFERENCE_MATCHED_FRACTION = 0.5;

// Per-image resolver for serial stills. Each crystal is indexed independently in one of the
// merohedrally-equivalent hands at random, so the ambiguity has to be broken per image (unlike rotation
// data, which resolve it once for the whole merged set with ChooseReindex above). Holds the external
// reference intensities and the crystal's twin-law operators, and reindexes one image's reflections into
// the hand best-correlated with the reference. This runs in the per-image pipeline at integration time,
// decoupled from scaling.
class ReindexAmbiguityResolver {
    const ScalingSettings s;
    const HKLKeyGenerator hkl_key_generator;
    std::map<HKLKey, double> reference_data;
    // Twin-law reindexing operators (identity excluded). Empty = no ambiguity (holohedral crystal) or the
    // cell/space group is unknown, in which case Resolve() is a no-op.
    std::vector<gemmi::Op> ops;

    bool Accept(const Reflection &r) const;
    // Pearson correlation of this image's partiality/Lorentz-corrected intensities against the reference,
    // matched by ASU index under `op`. Scale-invariant (the unknown per-image scale cancels). NaN when
    // too few reflections match.
    double ReferenceCC(std::span<const Reflection> reflections, const gemmi::Op &op) const;

public:
    // The reference must be an external, correctly-handed dataset (not the data's own running merge).
    ReindexAmbiguityResolver(const DiffractionExperiment &x, const std::vector<MergedReflection> &reference);

    // Reindex this image's reflections in place into the hand best-correlated with the reference (done
    // once per image, for good). No-op when there is no ambiguity.
    // Relabels the image's reflections and moves its lattice onto the same axes, so that a spot's new
    // index still predicts it where it was measured.
    void Resolve(std::span<Reflection> reflections, CrystalLattice &latt) const;
};