// 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"

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) {
    ReindexChoice best;
    best.identity_score = score(merged);
    best.score = best.identity_score;

    const auto ops = ReindexAmbiguityOperators(cell, space_group_number, max_obliquity_deg);
    best.n_candidates = 1 + static_cast<int>(ops.size());
    for (const auto &op : ops) {
        const double s = score(ReindexReflections(merged, op));
        if (s > best.score) {
            best.score = s;
            best.op = op;
            best.is_identity = false;
        }
    }
    return best;
}

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;
}