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

#include "ReindexAmbiguity.h"

#include <cmath>
#include <unordered_map>

#include "gemmi/twin.hpp"

#include "HKLKey.h"

namespace {
    constexpr size_t MIN_REFLECTIONS = 20;

    struct BestReindex {
        gemmi::Op op = gemmi::Op::identity();
        bool is_identity = true;
        double score = 0.0;
        double identity_score = 0.0;
    };

    // Among identity (the baseline, scored `identity_score`) and `ops`, return the operator with the
    // highest `score_op(op)`. A non-finite score never wins (too little overlap to decide). Shared by the
    // stills per-image resolver and the rotation post-merge ChooseReindex.
    BestReindex PickBestReindex(double identity_score, const std::vector<gemmi::Op> &ops,
                                const std::function<double(const gemmi::Op &)> &score_op) {
        BestReindex best;
        best.identity_score = identity_score;
        best.score = identity_score;
        for (const auto &op : ops) {
            const double sc = score_op(op);
            if (std::isfinite(sc) && (!std::isfinite(best.score) || sc > best.score)) {
                best.score = sc;
                best.op = op;
                best.is_identity = false;
            }
        }
        return best;
    }
}

std::vector<gemmi::Op> ReindexAmbiguityOperators(const UnitCell &cell, int space_group_number,
                                                 double max_obliquity_deg) {
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_number(space_group_number);
    if (sg == nullptr)
        return {};
    return gemmi::find_twin_laws(static_cast<gemmi::UnitCell>(cell), sg, max_obliquity_deg,
                                 /*all_ops=*/false);
}

std::vector<MergedReflection> ReindexReflections(const std::vector<MergedReflection> &merged,
                                                 const gemmi::Op &op) {
    std::vector<MergedReflection> out = merged;
    for (auto &r : out) {
        const gemmi::Op::Miller h = op.apply_to_hkl({{static_cast<int>(r.h),
                                                      static_cast<int>(r.k),
                                                      static_cast<int>(r.l)}});
        r.h = h[0];
        r.k = h[1];
        r.l = h[2];
    }
    return out;
}

ReindexChoice ChooseReindex(const std::vector<MergedReflection> &merged,
                            const UnitCell &cell, int space_group_number,
                            const std::function<double(const std::vector<MergedReflection> &)> &score,
                            double max_obliquity_deg) {
    const auto ops = ReindexAmbiguityOperators(cell, space_group_number, max_obliquity_deg);
    const BestReindex best = PickBestReindex(
        score(merged), ops,
        [&](const gemmi::Op &op) { return score(ReindexReflections(merged, op)); });

    ReindexChoice choice;
    choice.op = best.op;
    choice.is_identity = best.is_identity;
    choice.score = best.score;
    choice.identity_score = best.identity_score;
    choice.n_candidates = 1 + static_cast<int>(ops.size());
    return choice;
}

double ReferenceIntensityCC(const std::vector<MergedReflection> &merged,
                            const std::vector<MergedReflection> &reference,
                            int space_group_number) {
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_number(space_group_number);
    if (sg == nullptr)
        return 0.0;
    const HKLKeyGenerator key(/*merge_friedel=*/true, *sg);

    std::unordered_map<uint64_t, double> ref;
    ref.reserve(reference.size());
    for (const auto &r : reference)
        if (std::isfinite(r.I))
            ref[key(r).pack()] = r.I;

    double sx = 0, sy = 0, sxx = 0, syy = 0, sxy = 0;
    int n = 0;
    for (const auto &m : merged) {
        if (!std::isfinite(m.I))
            continue;
        const auto it = ref.find(key(m).pack());
        if (it == ref.end())
            continue;
        const double x = m.I, y = it->second;
        sx += x; sy += y; sxx += x * x; syy += y * y; sxy += x * y;
        ++n;
    }
    if (n < 10)
        return 0.0;
    const double cov = n * sxy - sx * sy;
    const double vx = n * sxx - sx * sx;
    const double vy = n * syy - sy * sy;
    return (vx > 0 && vy > 0) ? cov / std::sqrt(vx * vy) : 0.0;
}

ReindexAmbiguityResolver::ReindexAmbiguityResolver(const DiffractionExperiment &x,
                                                   const std::vector<MergedReflection> &reference)
    : s(x.GetScalingSettings()),
      hkl_key_generator(s.GetMergeFriedel(), x.GetSpaceGroupNumber().value_or(1)) {
    for (const auto &r : reference)
        reference_data[hkl_key_generator(r)] = r.I;
    if (x.GetUnitCell().has_value() && x.GetSpaceGroupNumber().has_value())
        ops = ReindexAmbiguityOperators(*x.GetUnitCell(), static_cast<int>(*x.GetSpaceGroupNumber()));
}

bool ReindexAmbiguityResolver::Accept(const Reflection &r) const {
    if (r.on_ice_ring)   // ice-contaminated intensity would bias the correlation; keep it out
        return false;
    return AcceptReflection(r, s.GetHighResolutionLimit_A());
}

double ReindexAmbiguityResolver::ReferenceCC(const std::vector<Reflection> &reflections,
                                             const gemmi::Op &op) const {
    double sx = 0, sy = 0, sxx = 0, syy = 0, sxy = 0;
    size_t n = 0;
    for (const auto &r : reflections) {
        if (!Accept(r) || r.partiality < s.GetMinPartiality())
            continue;
        if (!std::isfinite(r.I) || !std::isfinite(r.sigma) || r.sigma <= 0.0f || r.partiality <= 0.0f)
            continue;
        const gemmi::Op::Miller h = op.apply_to_hkl({{r.h, r.k, r.l}});
        const auto it = reference_data.find(hkl_key_generator(h[0], h[1], h[2]));
        if (it == reference_data.end())
            continue;
        const double x = static_cast<double>(r.I) * r.rlp / r.partiality;
        const double y = it->second;
        if (!std::isfinite(x) || !std::isfinite(y))
            continue;
        sx += x; sy += y; sxx += x * x; syy += y * y; sxy += x * y;
        ++n;
    }
    if (n < MIN_REFLECTIONS)
        return NAN;
    const double nd = static_cast<double>(n);
    const double cov = sxy - sx * sy / nd;
    const double vx = sxx - sx * sx / nd;
    const double vy = syy - sy * sy / nd;
    return (vx > 0 && vy > 0) ? cov / std::sqrt(vx * vy) : NAN;
}

// Serial stills index each crystal in one of the merohedrally-equivalent hands at random; pick, for this
// image alone, the reindexing whose intensities correlate best with the external reference and apply it.
void ReindexAmbiguityResolver::Resolve(std::vector<Reflection> &reflections) const {
    if (ops.empty())
        return;

    const BestReindex best = PickBestReindex(
        ReferenceCC(reflections, gemmi::Op::identity()), ops,
        [&](const gemmi::Op &op) { return ReferenceCC(reflections, op); });
    if (best.is_identity)
        return;

    for (auto &r : reflections) {
        const gemmi::Op::Miller h = best.op.apply_to_hkl({{r.h, r.k, r.l}});
        r.h = h[0]; r.k = h[1]; r.l = h[2];
    }
}