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

#include <catch2/catch_all.hpp>
#include <chrono>
#include <cmath>
#include <fstream>
#include <random>
#include <sstream>
#include <thread>

#include <Eigen/Geometry>

#include "../common/CUDAWrapper.h"
#include "../common/DiffractionExperiment.h"
#include "../image_analysis/indexing/FFBIDXIndexerCPU.h"

#ifdef JFJOCH_USE_CUDA
#include "../image_analysis/indexing/FFBIDXIndexer.h"
#endif

namespace {
    // A still of a test cell in a fixed random orientation: the reciprocal lattice points to 2 A that lie
    // within 0.001 1/A of the Ewald sphere of a 1 A beam along z, and then random spots in the same
    // resolution range. The seed is fixed, and the engine scaled by hand, so the frame is the same
    // with every standard library.
    std::vector<Coord> SyntheticStill(const UnitCell &uc, uint32_t seed, size_t noise_spots) {
        std::mt19937 rng(seed);
        auto draw = [&rng]() {
            return 2.0f * static_cast<float>(rng()) / static_cast<float>(std::mt19937::max()) - 1.0f;
        };
        const Eigen::Matrix3f rot = Eigen::Quaternionf(draw(), draw(), draw(), draw()).normalized().toRotationMatrix();

        const CrystalLattice l(uc);
        auto rotate = [&rot](const Coord &v) {
            const Eigen::Vector3f x = rot * Eigen::Vector3f(v.x, v.y, v.z);
            return Coord(x.x(), x.y(), x.z());
        };
        const Coord astar = rotate(l.Astar()), bstar = rotate(l.Bstar()), cstar = rotate(l.Cstar());

        const Coord k0(0, 0, 1.0f);
        std::vector<Coord> spots;
        for (int h = -40; h <= 40; h++)
            for (int k = -40; k <= 40; k++)
                for (int m = -40; m <= 40; m++) {
                    const Coord r = h * astar + k * bstar + m * cstar;
                    if (r.Length() > 0.5f || r.Length() < 0.02f)
                        continue;
                    if (std::abs((r + k0).Length() - 1.0f) < 0.001f)
                        spots.push_back(r);
                }
        while (noise_spots > 0) {
            const Coord r(0.5f * draw(), 0.5f * draw(), 0.5f * draw());
            if (r.Length() > 0.5f || r.Length() < 0.02f)
                continue;
            spots.push_back(r);
            noise_spots--;
        }
        return spots;
    }

    // Fraction of spots within 0.1 of a lattice point (in Miller index units)
    float IndexedFraction(const std::vector<Coord> &spots, const CrystalLattice &l) {
        size_t n = 0;
        for (const auto &s: spots) {
            const float h = s * l.Vec0(), k = s * l.Vec1(), m = s * l.Vec2();
            const float dh = h - std::round(h), dk = k - std::round(k), dm = m - std::round(m);
            if (dh * dh + dk * dk + dm * dm < 0.01f)
                n++;
        }
        return static_cast<float>(n) / static_cast<float>(spots.size());
    }

    bool SameLengths(const UnitCell &a, const UnitCell &b, float rel_tol) {
        std::array<double, 3> x = {a.a, a.b, a.c}, y = {b.a, b.b, b.c};
        std::sort(x.begin(), x.end());
        std::sort(y.begin(), y.end());
        for (int i = 0; i < 3; i++)
            if (std::abs(x[i] - y[i]) > rel_tol * y[i])
                return false;
        return true;
    }

    // The cell vectors as rows
    Eigen::Matrix3f CellMatrix(const UnitCell &uc) {
        const CrystalLattice l(uc);
        Eigen::Matrix3f m;
        m << 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;
        return m;
    }

    const std::vector<UnitCell> &TestCells() {
        static const std::vector<UnitCell> cells = {
            {30, 40, 50, 90, 90, 90},
            {80, 80, 90, 90, 90, 120},
            {40, 45, 80, 90, 82.5, 90},
            {60, 60, 60, 90, 90, 90}
        };
        return cells;
    }
}

TEST_CASE("Portable_FFBIDXIndexerCPU_SyntheticStill", "[portable][Indexing]") {
    for (size_t i = 0; i < TestCells().size(); i++) {
        const UnitCell &uc = TestCells()[i];
        const auto lattice_spots = SyntheticStill(uc, 1234 + i, 0);
        auto spots = SyntheticStill(uc, 1234 + i, lattice_spots.size() / 5);
        INFO("cell " << i << ", " << lattice_spots.size() << " lattice spots");
        REQUIRE(lattice_spots.size() > 30);

        DiffractionExperiment experiment;
        experiment.SetUnitCell(uc);
        FFBIDXIndexerCPU indexer;
        indexer.Setup(experiment);
        const auto result = indexer.Run(spots);
        REQUIRE(!result.lattice.empty());
        CHECK(SameLengths(result.lattice[0].GetUnitCell(), uc, 0.01f));
        CHECK(IndexedFraction(lattice_spots, result.lattice[0]) > 0.95f);
    }
}

// Time per call of the CPU indexer, alone and with Refine(), on one thread and on 16 at once
TEST_CASE("FFBIDXIndexerCPU_Benchmark", "[.][ffbidx_bench]") {
    const UnitCell uc(40, 45, 80, 90, 82.5, 90);
    DiffractionExperiment experiment;
    experiment.SetUnitCell(uc);

    for (const size_t n: {30, 80, 250}) {
        auto spots = SyntheticStill(uc, 99, 0);
        const auto lattice_spots = spots.size();
        spots = SyntheticStill(uc, 99, n > lattice_spots ? n - lattice_spots : 0);
        spots.resize(n);

        const int calls = 20;
        auto t0 = std::chrono::steady_clock::now();
        for (int i = 0; i < calls; i++)
            FFBIDXFindCells(CellMatrix(uc), spots, n);
        const double find_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - t0).count() / calls;

        FFBIDXIndexerCPU indexer;
        indexer.Setup(experiment);
        t0 = std::chrono::steady_clock::now();
        for (int i = 0; i < calls; i++)
            indexer.Run(spots);
        const double run_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - t0).count() / calls;

        const int threads = 16;
        t0 = std::chrono::steady_clock::now();
        std::vector<std::thread> workers;
        for (int t = 0; t < threads; t++)
            workers.emplace_back([&] {
                FFBIDXIndexerCPU w;
                w.Setup(experiment);
                for (int i = 0; i < calls; i++)
                    w.Run(spots);
            });
        for (auto &w: workers)
            w.join();
        const double wall_s = std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count();

        WARN(n << " spots: FindCells " << find_ms << " ms, with Refine " << run_ms << " ms; "
             << threads << " threads " << threads * calls / wall_s << " calls/s");
    }
}

#ifdef JFJOCH_USE_CUDA

namespace {
    // The GPU library's raw output, before Refine(), with the configuration FFBIDXIndexer uses
    FFBIDXCells RunGPU(const Eigen::Matrix3f &cell, const std::vector<Coord> &spots) {
        fast_feedback::config_runtime<float> crt{};
        fast_feedback::config_persistent<float> cpers{
            .max_output_cells = 32,
            .max_input_cells = 1,
            .max_spots = static_cast<unsigned>(spots.size()),
            .num_candidate_vectors = 32,
            .redundant_computations = true,
        };
        fast_feedback::refine::indexer<float> indexer{cpers, crt};
        indexer.iCellM() = cell;
        for (size_t i = 0; i < spots.size(); i++) {
            indexer.spotX(i) = spots[i].x;
            indexer.spotY(i) = spots[i].y;
            indexer.spotZ(i) = spots[i].z;
        }
        indexer.index(1, spots.size());
        return {indexer.oCellM(), indexer.oScoreV()};
    }

    // Cell (rows) and spots of one of the library's test files: the first line is the cell, then a
    // spot per line; '#' starts a comment
    std::pair<Eigen::Matrix3f, std::vector<Coord>> ReadSimpleData(const std::string &path) {
        std::ifstream f(path);
        REQUIRE(f.good());
        std::vector<float> v;
        std::string line;
        while (std::getline(f, line)) {
            if (line.empty() || line[0] == '#')
                continue;
            std::istringstream s(line);
            float x;
            while (s >> x)
                v.push_back(x);
        }
        REQUIRE(v.size() >= 12);
        Eigen::Matrix3f cell;
        cell << v[0], v[1], v[2], v[3], v[4], v[5], v[6], v[7], v[8];
        std::vector<Coord> spots;
        for (size_t i = 9; i + 2 < v.size(); i += 3)
            spots.emplace_back(v[i], v[i + 1], v[i + 2]);
        return {cell, spots};
    }

    float RelativeDifference(const Eigen::Matrix3f &a, const Eigen::Matrix3f &b) {
        return (a - b).norm() / b.norm();
    }

    // CPU and GPU agree when the best scores agree and the CPU's best cell is one of the GPU's best
    // cells - the samples are the same, so a matching cell agrees to rounding, and two cells within
    // rounding of each other in score may come out in either order
    void CompareWithGPU(const std::string &name, const Eigen::Matrix3f &cell, const std::vector<Coord> &spots) {
        const auto gpu = RunGPU(cell, spots);
        const auto cpu = FFBIDXFindCells(cell, spots, spots.size());

        float best_match = INFINITY;
        for (int k = 0; k < gpu.scores.size(); k++)
            best_match = std::min(best_match, RelativeDifference(cpu.cells.block<3, 3>(0, 0), gpu.cells.block<3, 3>(3 * k, 0)));

        int top_matched = 0;
        for (int j = 0; j < cpu.scores.size(); j++)
            for (int k = 0; k < gpu.scores.size(); k++)
                if (RelativeDifference(cpu.cells.block<3, 3>(3 * j, 0), gpu.cells.block<3, 3>(3 * k, 0)) < 1e-3f) {
                    top_matched++;
                    break;
                }

        UNSCOPED_INFO(name << ": " << spots.size() << " spots, score CPU " << cpu.scores(0) << " GPU " << gpu.scores(0)
                      << ", best cell to GPU " << best_match << ", " << top_matched << "/32 cells matched");
        CHECK(std::abs(cpu.scores(0) - gpu.scores(0)) <= 1e-3f * std::abs(gpu.scores(0)));
        CHECK(best_match < 1e-3f);
        CHECK(top_matched >= 24);
    }
}

TEST_CASE("FFBIDXIndexerCPU_MatchesGPU", "[Indexing][GPU]") {
    if (get_gpu_count() == 0)
        SKIP("no GPU");

    for (int i = 0; i < 9; i++) {
        for (const std::string kind: {"local", "peakfinder8", "radial"}) {
            const std::string name = "image" + std::to_string(i) + "_" + kind + ".txt";
            const auto [cell, spots] = ReadSimpleData(std::string(FFBIDX_SIMPLE_DATA_DIR) + "/" + name);
            CompareWithGPU(name, cell, spots);
        }
    }

    for (size_t i = 0; i < TestCells().size(); i++) {
        const auto spots = SyntheticStill(TestCells()[i], 1234 + i, 10);
        CompareWithGPU("synthetic " + std::to_string(i), CellMatrix(TestCells()[i]), spots);
    }
}

#endif