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

#include "SublatticeRelabel.h"

#include <algorithm>
#include <cmath>
#include <limits>

namespace {
    // Largest coefficient tried: a reference axis up to four frame axes long, which covers the 2x, 3x
    // and 4x axis multiples a few-spot frame gets indexed in.
    constexpr int MAX_COEFFICIENT = 4;
    // Largest volume ratio frame / reference accepted: the frame's lattice a 2x, 3x or 4x supercell (or
    // the reference itself in another setting). A denser lattice holds near-fits of almost any cell.
    constexpr int MAX_INDEX = 4;

    float AngleDeg(const Coord &u, const Coord &v) {
        const float c = (u * v) / (u.Length() * v.Length());
        return std::acos(std::clamp(c, -1.0f, 1.0f)) * (180.0f / 3.14159265f);
    }

    int Det(const std::array<std::array<int, 3>, 3> &m) {
        return m[0][0] * (m[1][1] * m[2][2] - m[1][2] * m[2][1])
             - m[0][1] * (m[1][0] * m[2][2] - m[1][2] * m[2][0])
             + m[0][2] * (m[1][0] * m[2][1] - m[1][1] * m[2][0]);
    }
}

std::vector<std::array<std::array<int, 3>, 3>> FindSublatticeBases(const CrystalLattice &frame,
                                                                   const UnitCell &reference,
                                                                   const IndexingSettings &settings) {
    const CrystalLattice ref(reference);
    const Coord ref_star[3] = {ref.Astar(), ref.Bstar(), ref.Cstar()};
    const Coord frame_star[3] = {frame.Astar(), frame.Bstar(), frame.Cstar()};
    const float dist_tol = settings.GetUnitCellDistTolerance();
    const float angle_tol = settings.GetUnitCellAngleTolerance_deg();

    // The frame's reciprocal lattice vectors as long as each reference axis.
    struct Candidate { std::array<int, 3> n; Coord v; };
    std::vector<Candidate> axis[3];
    for (int i = -MAX_COEFFICIENT; i <= MAX_COEFFICIENT; ++i)
        for (int j = -MAX_COEFFICIENT; j <= MAX_COEFFICIENT; ++j)
            for (int k = -MAX_COEFFICIENT; k <= MAX_COEFFICIENT; ++k) {
                const Coord v = frame_star[0] * i + frame_star[1] * j + frame_star[2] * k;
                for (int a = 0; a < 3; ++a)
                    if (std::fabs(v.Length() - ref_star[a].Length()) <= dist_tol * ref_star[a].Length())
                        axis[a].push_back({{i, j, k}, v});
            }

    const float ref_angle[3] = {AngleDeg(ref_star[1], ref_star[2]), AngleDeg(ref_star[0], ref_star[2]),
                                AngleDeg(ref_star[0], ref_star[1])};
    std::vector<std::array<std::array<int, 3>, 3>> bases;
    for (const auto &a: axis[0])
        for (const auto &b: axis[1]) {
            if (std::fabs(AngleDeg(a.v, b.v) - ref_angle[2]) > angle_tol)
                continue;
            for (const auto &c: axis[2]) {
                if (std::fabs(AngleDeg(b.v, c.v) - ref_angle[0]) > angle_tol
                    || std::fabs(AngleDeg(a.v, c.v) - ref_angle[1]) > angle_tol)
                    continue;
                const std::array<std::array<int, 3>, 3> m = {a.n, b.n, c.n};
                if (Det(m) > 0 && Det(m) <= MAX_INDEX)
                    bases.push_back(m);
            }
        }
    return bases;
}

namespace {
    // The frame's reflections at nodes of the lattice M picks out, re-indexed on it: a frame index h'
    // is h = h' M^-1 = h' adj(M) / det(M) there, a node only where all three are integers.
    ReflectionVector NodesOf(const ReflectionVector &reflections, const std::array<std::array<int, 3>, 3> &M,
                             float d_min) {
        const int det = Det(M);
        int adj[3][3];
        for (int i = 0; i < 3; ++i)
            for (int j = 0; j < 3; ++j) {
                const int r0 = (j + 1) % 3, r1 = (j + 2) % 3, c0 = (i + 1) % 3, c1 = (i + 2) % 3;
                adj[i][j] = M[r0][c0] * M[r1][c1] - M[r0][c1] * M[r1][c0];
            }
        ReflectionVector nodes;
        for (const auto &r: reflections) {
            const int hp[3] = {r.h, r.k, r.l};
            int h[3];
            bool node = true;
            for (int j = 0; j < 3; ++j) {
                const int s = hp[0] * adj[0][j] + hp[1] * adj[1][j] + hp[2] * adj[2][j];
                node = node && s % det == 0;
                h[j] = s / det;
            }
            if (!node || !(r.d >= d_min))
                continue;
            Reflection relabelled = r;
            relabelled.h = h[0];
            relabelled.k = h[1];
            relabelled.l = h[2];
            nodes.push_back(relabelled);
        }
        return nodes;
    }

    double MeanISigma(const ReflectionVector &reflections) {
        double sum = 0.0;
        size_t n = 0;
        for (const auto &r: reflections)
            if (std::isfinite(r.I) && r.sigma > 0.0f) {
                sum += r.I / r.sigma;
                ++n;
            }
        return n > 0 ? sum / static_cast<double>(n) : -std::numeric_limits<double>::infinity();
    }
}

SublatticeRelabelResult RelabelOnSublattice(std::vector<IntegrationOutcome> &outcomes,
                                            const UnitCell &reference,
                                            const IndexingSettings &settings) {
    const auto on_reference = [&](const IntegrationOutcome &o) {
        return o.latt.GetUnitCell().is_close(reference, settings.GetUnitCellDistTolerance(),
                                             settings.GetUnitCellAngleTolerance_deg());
    };
    // A frame indexed in the wrong lattice was also predicted to the wrong resolution - its limit read
    // off spots that were mostly noise - so its reflections are kept only as fine as the frames already
    // on the reference cell reach.
    float d_min = std::numeric_limits<float>::infinity();
    for (const auto &o: outcomes)
        if (on_reference(o))
            for (const auto &r: o.reflections)
                if (std::isfinite(r.d) && r.d > 0.0f)
                    d_min = std::min(d_min, r.d);

    SublatticeRelabelResult result;
    for (auto &o: outcomes) {
        if (o.reflections.empty() || on_reference(o))
            continue;
        // A lattice can hold the reference in more than one way - a supercell whose metric is nearly
        // symmetric fits it along either of two axes - and only the true one has signal at its nodes:
        // the reflections between them are empty. So the bases are told apart by the data.
        std::optional<std::array<std::array<int, 3>, 3>> best;
        ReflectionVector best_nodes;
        double best_signal = 0.0;
        for (const auto &m: FindSublatticeBases(o.latt, reference, settings)) {
            auto nodes = NodesOf(o.reflections, m, d_min);
            const double signal = MeanISigma(nodes);
            if (!best || signal > best_signal) {
                best = m;
                best_nodes = std::move(nodes);
                best_signal = signal;
            }
        }
        if (!best)
            continue;

        result.reflections_dropped += o.reflections.size() - best_nodes.size();
        result.reflections_kept += best_nodes.size();
        o.reflections = std::move(best_nodes);

        // The reference basis in this frame's orientation: the reciprocal axes are rows of M times the
        // frame's, and the real axes their reciprocal.
        const auto &M = *best;
        const Coord fs[3] = {o.latt.Astar(), o.latt.Bstar(), o.latt.Cstar()};
        Coord rs[3];
        for (int i = 0; i < 3; ++i)
            rs[i] = fs[0] * M[i][0] + fs[1] * M[i][1] + fs[2] * M[i][2];
        const float v_star = rs[0] * (rs[1] % rs[2]);
        o.latt = CrystalLattice((rs[1] % rs[2]) / v_star, (rs[2] % rs[0]) / v_star, (rs[0] % rs[1]) / v_star);
        ++result.frames;
    }
    return result;
}