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

// Ported from the CUDA kernels of the fast-feedback indexer,
// https://github.com/paulscherrerinstitute/fast-feedback-indexer (indexer/src/indexer_gpu.cu)
// (c) 2022 Paul Scherrer Institute, BSD-3-Clause
//
// The samples and the scores are the library's; the loops are reorganised for a CPU. Each block of
// samples is scored against one spot at a time, so the inner loop runs over samples and vectorises,
// and in the cell search the rotation about the candidate vector is separated out of the spot terms.

#include "FFBIDXIndexerCPU.h"

#include <algorithm>
#include <array>
#include <bit>
#include <cmath>
#include <cstdint>
#include <limits>
#include <numbers>
#include <numeric>

#include "PostIndexingRefinement.h"

namespace {
    // The GPU library's runtime defaults (config_runtime in ffbidx/indexer.h) and rugnux's persistent
    // configuration of it (FFBIDXIndexer.h). Keep them in step: the CPU vs GPU test compares the two.
    constexpr int NUM_HALFSPHERE_POINTS = 32 * 1024;
    constexpr int NUM_ANGLE_POINTS = 1024;  // the library's own choice for 32k half-sphere points
    constexpr int NUM_CANDIDATE_VECTORS = 32;
    constexpr int NUM_OUTPUT_CELLS = 32;
    constexpr int MIN_SPOTS = 6;            // vector refinement
    constexpr float TRIML = 0.001f;
    constexpr float TRIMH = 0.3f;
    constexpr float DELTA = 0.1f;
    constexpr float DIST1 = 0.2f;
    constexpr float DIST3 = 0.15f;
    constexpr float LENGTH_THRESHOLD = 1e-9f;

    // Half-sphere directions are scored in blocks of this many, so a block's accumulators stay in L1
    constexpr int BLOCK = 1024;
    static_assert(NUM_HALFSPHERE_POINTS % BLOCK == 0);

    // Every factor of a score's product is in [TRIML + DELTA, TRIMH + DELTA]; after 8 of them the
    // product is still a normal float
    constexpr int RENORMALISE_EVERY = 8;

    struct Tables {
        std::array<std::vector<float>, 3> dir;      // Fibonacci points on the half sphere z >= 0
        std::vector<float> cos_alpha, sin_alpha;    // rotation angles about a cell vector
    };

    // (a, rest) = a + x * y + rest, a Kahan sum, with the product fused into the first addition as
    // nvcc compiles the library's
    inline void KahanAdd(float &a, float &rest, float x, float y = 1.0f) {
        const float s = std::fma(x, y, rest);
        const float t = a;
        a = t + s;
        rest = s - (a - t);
    }

    Tables MakeTables() {
        constexpr double pi2 = 2.0 * std::numbers::pi;
        const double dl = 3.0 - std::sqrt(5.0);

        Tables t;
        for (auto &v: t.dir)
            v.resize(NUM_HALFSPHERE_POINTS);
        std::array<float, 32> dl2n;
        for (int n = 0; n < 32; n++)
            dl2n[n] = static_cast<float>(std::fmod(dl * std::ldexp(1.0, n), pi2));
        for (int i = 0; i < NUM_HALFSPHERE_POINTS; i++) {
            const float z = 1.0f - static_cast<float>(i) / static_cast<float>(NUM_HALFSPHERE_POINTS - 1);
            const float r = std::sqrt(std::fma(-z, z, 1.0f));
            float l = 0, rest = 0;
            for (int n = 0; n < 32; n++)
                if (i & (1 << n))
                    KahanAdd(l, rest, dl2n[n]);
            t.dir[0][i] = std::cos(l) * r;
            t.dir[1][i] = std::sin(l) * r;
            t.dir[2][i] = z;
        }
        for (int q = 0; q < NUM_ANGLE_POINTS; q++) {
            const double alpha = q * pi2 / NUM_ANGLE_POINTS;
            t.cos_alpha.push_back(static_cast<float>(std::cos(alpha)));
            t.sin_alpha.push_back(static_cast<float>(std::sin(alpha)));
        }
        return t;
    }

    const Tables &GetTables() {
        static const Tables t = MakeTables();
        return t;
    }

    inline float Dist2Int(float x) {
        return std::abs(x - std::rint(x));
    }

    // One spot's share of a sample's score: a factor of the product, and whether it counts as good
    inline void AddSpot(float d, float good_limit, float &product, int &good) {
        good += static_cast<int>(d < good_limit);
        product *= std::clamp(d, TRIML, TRIMH) + DELTA;
    }

    // Moves the power of two of the product into the exponent, leaving the product in [1, 2). The
    // library takes a log2 every 8 spots instead; this is the same sum, and it vectorises.
    inline void Renormalise(float &product, int &exponent) {
        const auto bits = std::bit_cast<uint32_t>(product);
        exponent += static_cast<int>(bits >> 23) - 127;
        product = std::bit_cast<float>((bits & 0x007fffffu) | 0x3f800000u);
    }

    // exp2(mean log2(trim(d) + delta)) - delta - (good + 1): lower is better. A NaN, from a degenerate
    // candidate, goes last instead of breaking the sort.
    inline float Score(float log2_sum, int good, int nspots) {
        const float s = std::exp2(log2_sum / static_cast<float>(nspots)) - DELTA - static_cast<float>(good + 1);
        return std::isnan(s) ? std::numeric_limits<float>::infinity() : s;
    }

    // Indices of the k lowest scores, best first
    std::vector<int> Best(const std::vector<float> &score, int k) {
        std::vector<int> idx(score.size());
        std::iota(idx.begin(), idx.end(), 0);
        auto lower = [&score](int a, int b) { return score[a] < score[b] || (score[a] == score[b] && a < b); };
        std::nth_element(idx.begin(), idx.begin() + k, idx.end(), lower);
        idx.resize(k);
        std::sort(idx.begin(), idx.end(), lower);
        return idx;
    }

    // Stage 1: every half-sphere direction as the direction of a cell vector of this length, scored on
    // how close the spots' projections on it come to integers
    std::vector<int> FindCandidates(const Tables &t, const std::vector<Eigen::Vector3f> &spots, float length) {
        const int nspots = static_cast<int>(spots.size());
        std::vector<float> score(NUM_HALFSPHERE_POINTS);
        std::array<float, BLOCK> product;
        std::array<int, BLOCK> exponent, good;

        for (int b = 0; b < NUM_HALFSPHERE_POINTS; b += BLOCK) {
            const float *vx = t.dir[0].data() + b;
            const float *vy = t.dir[1].data() + b;
            const float *vz = t.dir[2].data() + b;
            product.fill(1.0f);
            exponent.fill(0);
            good.fill(0);
            for (int i = 0; i < nspots; i++) {
                const Eigen::Vector3f s = length * spots[i];
                for (int k = 0; k < BLOCK; k++)
                    AddSpot(Dist2Int(s.x() * vx[k] + s.y() * vy[k] + s.z() * vz[k]), DIST1, product[k], good[k]);
                if ((i + 1) % RENORMALISE_EVERY == 0)
                    for (int k = 0; k < BLOCK; k++)
                        Renormalise(product[k], exponent[k]);
            }
            for (int k = 0; k < BLOCK; k++)
                score[b + k] = Score(static_cast<float>(exponent[k]) + std::log2(product[k]), good[k], nspots);
        }
        return Best(score, NUM_CANDIDATE_VECTORS);
    }

    // The library's dot product: rounding matters here, as a spot on the edge of a tolerance decides
    // which spots the candidate is refined against
    inline float Dot(const Eigen::Vector3f &a, const Eigen::Vector3f &b) {
        return std::fma(a[0], b[0], std::fma(a[1], b[1], a[2] * b[2]));
    }

    // A sum as the library forms it on a warp of 32 lanes: spot i is added on lane i % 32, and the
    // lanes are reduced by shuffles that leave partial sums behind in every lane but the first. Those
    // stay in the lanes and enter the next round's sum. Reproduced as it is, because the refined
    // candidates - and so the cells searched - depend on it.
    struct WarpSum {
        std::array<float, 32> value{}, rest{};

        void Add(size_t spot, float x, float y) {
            KahanAdd(value[spot % 32], rest[spot % 32], x, y);
        }

        float Reduce() {
            std::array<float, 32> r{};
            for (int d = 1; d < 32; d *= 2) {
                const auto old_value = value;
                const auto old_r = r;
                for (int l = 0; l < 32; l++) {
                    const int from = (l % (2 * d) < d) ? l + d : l;
                    r[l] = old_r[from];
                    KahanAdd(value[l], r[l], old_value[from]);
                }
            }
            return value[0];
        }
    };

    // x = A^-1 v for symmetric A given by its upper triangle (00 01 02 11 12 22), by cofactors, as
    // the library solves it (products fused as nvcc compiles them)
    Eigen::Vector3f SymSolve(const std::array<float, 6> &A, const std::array<float, 3> &v) {
        auto cofactor = [](float a, float b, float c, float d) { return std::fma(a, b, -(c * d)); };
        const float co00 = cofactor(A[3], A[5], A[4], A[4]);
        const float co01 = cofactor(A[1], A[5], A[2], A[4]);
        const float co02 = cofactor(A[1], A[4], A[2], A[3]);
        const float co11 = cofactor(A[0], A[5], A[2], A[2]);
        const float co12 = cofactor(A[0], A[4], A[1], A[2]);
        const float co22 = cofactor(A[0], A[3], A[1], A[1]);
        // a0 * b0 + a1 * b1 + a2 * b2
        auto sum3 = [](float a0, float b0, float a1, float b1, float a2, float b2) {
            float t = a0 * b0;
            float rest = 0;
            KahanAdd(t, rest, a1, b1);
            KahanAdd(t, rest, a2, b2);
            return t;
        };
        const float r_det = 1.0f / sum3(A[0], co00, -A[1], co01, A[2], co02);
        return {sum3(co00, v[0], -co01, v[1], co02, v[2]) * r_det,
                sum3(-co01, v[0], co11, v[1], -co12, v[2]) * r_det,
                sum3(co02, v[0], -co12, v[1], co22, v[2]) * r_det};
    }

    // Stage 2: a candidate refined by least squares against the spots it nearly indexes, in four rounds
    // of tightening tolerance; the Miller indices stay those of the unrefined vector
    Eigen::Vector3f RefineCandidate(const std::vector<Eigen::Vector3f> &spots, const Eigen::Vector3f &z0) {
        std::array<WarpSum, 6> STS;
        std::array<WarpSum, 3> STm;
        Eigen::Vector3f z = z0;
        for (const float beta: {0.1f, 0.05f, 0.025f, 0.01f}) {
            int good = 0;
            for (size_t i = 0; i < spots.size(); i++) {
                const Eigen::Vector3f &s = spots[i];
                const float m = std::rint(Dot(s, z0));
                if (std::abs(Dot(s, z) - m) >= beta)
                    continue;
                good++;
                int e = 0;
                for (int j = 0; j < 3; j++) {
                    STm[j].Add(i, s[j], m);
                    for (int k = j; k < 3; k++)
                        STS[e++].Add(i, s[j], s[k]);
                }
            }
            if (good < MIN_SPOTS)
                break;
            std::array<float, 6> A;
            std::array<float, 3> v;
            for (int e = 0; e < 3; e++)
                v[e] = STm[e].Reduce();
            for (int e = 0; e < 6; e++)
                A[e] = STS[e].Reduce();
            z = SymSolve(A, v);
        }
        return z;
    }

    // The input cell turned so its vector cell_vec lies along c, in the frame (x, y, z = c/|c|) the
    // rotation about c is sampled in: a(alpha) = a_z z + a_xy (cos alpha x + sin alpha y), and b the
    // same at alpha + delta.
    struct CellFrame {
        int cell_vec;
        Eigen::Vector3f x, y, z;
        float c_length;
        float a_z, a_xy, b_z, b_xy;
        float cos_delta, sin_delta;

        Eigen::Vector3f A(float cos_a, float sin_a) const {
            return a_z * z + a_xy * (cos_a * x + sin_a * y);
        }

        Eigen::Vector3f B(float cos_a, float sin_a) const {
            const float c = cos_a * cos_delta - sin_a * sin_delta;
            const float s = sin_a * cos_delta + cos_a * sin_delta;
            return b_z * z + b_xy * (c * x + s * y);
        }
    };

    CellFrame MakeFrame(const Eigen::Matrix3f &cell, int cell_vec, const Eigen::Vector3f &c) {
        CellFrame f;
        f.cell_vec = cell_vec;
        f.c_length = c.norm();
        f.z = c / f.c_length;

        // Mirroring on the bisector of the input vector and z takes one onto the other
        const Eigen::Vector3f v = cell.row(cell_vec);
        Eigen::Vector3f a = cell.row((cell_vec + 1) % 3);
        Eigen::Vector3f b = cell.row((cell_vec + 2) % 3);
        Eigen::Vector3f t = v + v.norm() * f.z;
        if (t.norm() >= std::numeric_limits<float>::epsilon())
            t.normalize();
        else
            t = f.z.cross(a).normalized();
        a = 2.0f * a.dot(t) * t - a;
        b = 2.0f * b.dot(t) * t - b;

        f.a_z = a.dot(f.z);
        f.x = (a - f.a_z * f.z).normalized();
        f.a_xy = a.dot(f.x);
        f.y = f.z.cross(f.x);

        f.b_z = b.dot(f.z);
        const Eigen::Vector3f b_dir = (b - f.b_z * f.z).normalized();
        f.b_xy = b.dot(b_dir);
        f.cos_delta = b_dir.dot(f.x);
        f.sin_delta = b_dir.dot(f.y);
        return f;
    }

    // Stage 3: the cell of one frame at every rotation angle, scored on how close each spot comes to
    // a lattice point. The spot's projections on z, x and y are taken once; c does not turn at all.
    void ScoreAngles(const Tables &t, const CellFrame &f, const std::vector<Eigen::Vector3f> &spots, float *score) {
        const int nspots = static_cast<int>(spots.size());
        const float *cos_a = t.cos_alpha.data();
        const float *sin_a = t.sin_alpha.data();
        std::array<float, NUM_ANGLE_POINTS> product;
        std::array<int, NUM_ANGLE_POINTS> exponent, good;
        product.fill(1.0f);
        exponent.fill(0);
        good.fill(0);

        for (int i = 0; i < nspots; i++) {
            const float zs = f.z.dot(spots[i]);
            const float xs = f.x.dot(spots[i]);
            const float ys = f.y.dot(spots[i]);
            const float dc = Dist2Int(f.c_length * zs);
            const float a0 = f.a_z * zs;
            const float a1 = f.a_xy * xs;
            const float a2 = f.a_xy * ys;
            const float b0 = f.b_z * zs;
            const float b1 = f.b_xy * (f.cos_delta * xs + f.sin_delta * ys);
            const float b2 = f.b_xy * (f.cos_delta * ys - f.sin_delta * xs);
            for (int q = 0; q < NUM_ANGLE_POINTS; q++) {
                const float da = Dist2Int(a0 + a1 * cos_a[q] + a2 * sin_a[q]);
                const float db = Dist2Int(b0 + b1 * cos_a[q] + b2 * sin_a[q]);
                AddSpot(std::sqrt(dc * dc + da * da + db * db), DIST3, product[q], good[q]);
            }
            if ((i + 1) % RENORMALISE_EVERY == 0)
                for (int q = 0; q < NUM_ANGLE_POINTS; q++)
                    Renormalise(product[q], exponent[q]);
        }
        // The library starts this sum at 1, not 0; kept, so the scores agree with the GPU
        for (int q = 0; q < NUM_ANGLE_POINTS; q++)
            score[q] = Score(1.0f + static_cast<float>(exponent[q]) + std::log2(product[q]), good[q], nspots);
    }
}

// Following the fast-feedback indexer (H.-C. Stadler, PSI) and TORO, Gasparotto et al. (2024)
// J. Appl. Cryst. 57, 931-944
FFBIDXCells FFBIDXFindCells(const Eigen::Matrix3f &cell, const std::vector<Coord> &coord, size_t nspots) {
    const Tables &t = GetTables();

    std::vector<Eigen::Vector3f> spots(nspots);
    for (size_t i = 0; i < nspots; i++)
        spots[i] = Eigen::Vector3f(coord[i].x, coord[i].y, coord[i].z);

    // Cell vectors of the same length share one set of candidates
    std::array<int, 3> group{};
    std::vector<float> group_length;
    for (int v = 0; v < 3; v++) {
        const float length = cell.row(v).norm();
        auto it = std::find_if(group_length.begin(), group_length.end(),
                               [length](float l) { return std::abs(l - length) < LENGTH_THRESHOLD; });
        group[v] = static_cast<int>(it - group_length.begin());
        if (it == group_length.end())
            group_length.push_back(length);
    }

    std::vector<std::vector<Eigen::Vector3f>> candidates(group_length.size());
    for (size_t g = 0; g < group_length.size(); g++) {
        for (const int i: FindCandidates(t, spots, group_length[g])) {
            const Eigen::Vector3f dir(t.dir[0][i], t.dir[1][i], t.dir[2][i]);
            candidates[g].push_back(RefineCandidate(spots, group_length[g] * dir));
        }
    }

    // Every cell vector along every candidate of its length, both ways round
    std::vector<CellFrame> frames;
    for (int v = 0; v < 3; v++) {
        for (const auto &c: candidates[group[v]]) {
            frames.push_back(MakeFrame(cell, v, c));
            frames.push_back(MakeFrame(cell, v, -c));
        }
    }

    std::vector<float> score(frames.size() * NUM_ANGLE_POINTS);
    for (size_t i = 0; i < frames.size(); i++)
        ScoreAngles(t, frames[i], spots, score.data() + i * NUM_ANGLE_POINTS);

    FFBIDXCells ret;
    ret.cells.resize(3 * NUM_OUTPUT_CELLS, 3);
    ret.scores.resize(NUM_OUTPUT_CELLS);
    const auto best = Best(score, NUM_OUTPUT_CELLS);
    for (int k = 0; k < NUM_OUTPUT_CELLS; k++) {
        const CellFrame &f = frames[best[k] / NUM_ANGLE_POINTS];
        const int q = best[k] % NUM_ANGLE_POINTS;
        ret.scores(k) = score[best[k]];
        ret.cells.row(3 * k + f.cell_vec) = f.c_length * f.z;
        ret.cells.row(3 * k + (f.cell_vec + 1) % 3) = f.A(t.cos_alpha[q], t.sin_alpha[q]);
        ret.cells.row(3 * k + (f.cell_vec + 2) % 3) = f.B(t.cos_alpha[q], t.sin_alpha[q]);
    }
    return ret;
}

void FFBIDXIndexerCPU::SetupUnitCell(const std::optional<UnitCell> &cell) {
    if (!cell.has_value())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "FFBIDX requires unit cell");
    reference_unit_cell = cell;

    CrystalLattice l(cell.value());
    input_cell << l.Vec0().x, l.Vec0().y, l.Vec0().z,
            l.Vec1().x, l.Vec1().y, l.Vec1().z,
            l.Vec2().x, l.Vec2().y, l.Vec2().z;
}

std::vector<CrystalLattice> FFBIDXIndexerCPU::RunInternal(const std::vector<Coord> &coord, size_t nspots) {
    if (nspots > coord.size())
        nspots = coord.size();

    if (nspots < viable_cell_min_spots)
        return {};

    auto found = FFBIDXFindCells(input_cell, coord, nspots);

    RefineParameters parameters{
        .viable_cell_min_spots = viable_cell_min_spots,
        .dist_tolerance_vs_reference = dist_tolerance_vs_reference,
        .reference_unit_cell = reference_unit_cell,
        .min_length_A = 1, // doesn't matter
        .max_length_A = 1000, // doesn't matter
        .min_angle_deg = 30,
        .max_angle_deg = 150,
        .indexing_tolerance = indexing_tolerance
    };

    return Refine(coord, nspots, found.cells, found.scores, parameters);
}