// 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 "../common/CUDAWrapper.h"

#ifdef JFJOCH_USE_CUDA

#include <algorithm>
#include <cmath>
#include <numeric>
#include <random>
#include <vector>

#include "../common/Reflection.h"
#include "../image_analysis/scale_merge/RotationScaleMergeGPU.h"

namespace {

// Frames of reflections with every awkward case the host ingest orders: repeated hkl on one frame (both
// Ewald intersections), the same frame position twice, unusable resolutions and backgrounds, indices
// that fill the sixteen bits of the sort key.
std::vector<std::vector<Reflection>> MakeFrames(int n_frames, uint32_t seed) {
    std::mt19937 rng(seed);
    std::uniform_int_distribution<int> idx(-5, 5), count(0, 900), pick(0, 19);
    std::uniform_real_distribution<float> u(0.0f, 1.0f);
    std::vector<std::vector<Reflection>> frames(n_frames);
    for (int f = 0; f < n_frames; ++f) {
        const int n = (f == 7) ? 0 : count(rng);
        for (int j = 0; j < n; ++j) {
            Reflection r{};
            r.h = idx(rng); r.k = idx(rng); r.l = idx(rng);
            if (pick(rng) == 0) r.h = (j % 2) ? 32767 : -32768;
            r.image_number = static_cast<float>(f) + 0.25f * static_cast<float>(pick(rng) % 3);
            r.d = 1.0f + 0.1f * static_cast<float>(std::abs(r.h % 7) + std::abs(r.k) + std::abs(r.l));
            if (pick(rng) == 1) r.d = NAN;
            if (pick(rng) == 2) r.d = 0.0f;
            r.I = 100.0f * u(rng);
            r.sigma = 1.0f + u(rng);
            r.bkg = (pick(rng) == 3) ? NAN : 10.0f * u(rng) - 1.0f;
            r.var_bkg = u(rng);
            r.prescaling_corr = 0.5f + u(rng);
            r.qe_corr = 0.9f + 0.1f * u(rng);
            r.flight_corr = 1.0f + 0.01f * u(rng);
            r.partiality = u(rng);
            r.zeta = u(rng);
            r.image_scale_corr = u(rng);
            r.delta_phi_deg = u(rng) - 0.5f;
            r.predicted_x = 1000.0f * u(rng);
            r.predicted_y = 1000.0f * u(rng);
            r.on_ice_ring = pick(rng) == 4;
            r.clipped = pick(rng) == 5;
            r.overloaded = pick(rng) == 6;
            frames[f].push_back(r);
        }
    }
    return frames;
}

struct HostKey { int16_t h, k, l; float img; int32_t idx; float d; };

}  // namespace

TEST_CASE("RotationIngestGPU_SortAndCutMatchHost", "[RotationScale][gpu]") {
    if (get_gpu_count() == 0)
        SKIP("no GPU");
    const int n_frames = 60;
    const auto frames = MakeFrames(n_frames, 17);
    std::vector<const Reflection *> ptr(n_frames);
    std::vector<int32_t> count(n_frames), src_start(n_frames + 1, 0);
    for (int f = 0; f < n_frames; ++f) {
        ptr[f] = frames[f].data();
        count[f] = static_cast<int32_t>(frames[f].size());
        src_start[f + 1] = src_start[f] + count[f];
    }
    const int n = src_start[n_frames];
    std::vector<const Reflection *> all(n);
    for (int f = 0; f < n_frames; ++f)
        for (int j = 0; j < count[f]; ++j) all[src_start[f] + j] = &frames[f][j];

    for (const auto &[d_min, d_max] : {std::pair<std::optional<double>, std::optional<double>>{std::nullopt, std::nullopt},
                                       {1.4, 2.1}, {std::nullopt, 1.8}}) {
        RotationScaleMergeGPU gpu(0);
        gpu.IngestUpload(n_frames, ptr.data(), count.data(), 4);

        // The flux meter, summed per frame in reflection order.
        std::vector<double> mean(n_frames);
        gpu.IngestMeanBackground(mean.data());
        for (int f = 0; f < n_frames; ++f) {
            double sum = 0.0;
            int c = 0;
            for (const auto &r : frames[f])
                if (std::isfinite(r.bkg) && r.bkg > 0.0f) { sum += r.bkg; ++c; }
            if (c > 0) CHECK(mean[f] == sum / c);
            else CHECK(std::isnan(mean[f]));
        }

        std::vector<double> flux(n_frames);
        for (int f = 0; f < n_frames; ++f) flux[f] = 0.9 + 0.01 * f;
        const int n_keep = gpu.IngestSort(flux.data(), d_min, d_max);

        // The host ingest: one total order, its raw-hkl runs, each with its first usable resolution, and
        // the runs the range keeps with their observations.
        std::vector<HostKey> key(n);
        for (int i = 0; i < n; ++i)
            key[i] = {static_cast<int16_t>(all[i]->h), static_cast<int16_t>(all[i]->k), static_cast<int16_t>(all[i]->l),
                      all[i]->image_number, i, all[i]->d};
        std::sort(key.begin(), key.end(), [](const HostKey &a, const HostKey &b) {
            if (a.h != b.h) return a.h < b.h;
            if (a.k != b.k) return a.k < b.k;
            if (a.l != b.l) return a.l < b.l;
            if (a.img != b.img) return a.img < b.img;
            return a.idx < b.idx;
        });
        const bool limited = d_min || d_max;
        std::vector<uint8_t> keep(n, limited ? 0 : 1);
        struct Run { int start, count; int32_t h, k, l; float d; };
        std::vector<Run> runs;
        int n_runs_all = 0;
        for (int i = 0; i < n; ) {
            int j = i;
            float d = NAN;
            while (j < n && key[j].h == key[i].h && key[j].k == key[i].k && key[j].l == key[i].l) {
                if (!std::isfinite(d) && std::isfinite(key[j].d) && key[j].d > 0.0f) d = key[j].d;
                ++j;
            }
            ++n_runs_all;
            const bool in = !limited || (std::isfinite(d) && d > 0.0f && !(d_min && d < *d_min) && !(d_max && d > *d_max));
            if (in) {
                runs.push_back({i, j - i, key[i].h, key[i].k, key[i].l, d});
                for (int p = i; p < j; ++p) keep[key[p].idx] = 1;
            }
            i = j;
        }
        std::vector<int32_t> new_idx(n, -1);
        int at = 0;
        for (int i = 0; i < n; ++i) if (keep[i]) new_idx[i] = at++;
        REQUIRE(n_keep == at);

        RotationScaleMergeGPU::IngestLayout L;
        gpu.IngestGetLayout(L);
        CHECK(L.n_reflections == n);
        CHECK(L.n_runs_all == n_runs_all);
        REQUIRE(L.run_start.size() == runs.size());
        std::vector<int32_t> h(n_keep), k(n_keep), l(n_keep), perm(n_keep);
        std::vector<float> dphi(n_keep), part(n_keep), img(n_keep), zeta(n_keep), d(n_keep), corr(n_keep),
                I(n_keep), sigma(n_keep);
        gpu.IngestGetObservations(h.data(), k.data(), l.data(), dphi.data(), part.data(), img.data(), zeta.data(),
                                  d.data(), corr.data(), I.data(), sigma.data(), perm.data());
        int at_run = 0;
        for (size_t r = 0; r < runs.size(); ++r) {
            CHECK(L.run_start[r] == at_run);
            CHECK(L.run_count[r] == runs[r].count);
            CHECK(L.run_h[r] == runs[r].h);
            CHECK(L.run_k[r] == runs[r].k);
            CHECK(L.run_l[r] == runs[r].l);
            CHECK((L.run_d[r] == runs[r].d || (std::isnan(L.run_d[r]) && std::isnan(runs[r].d))));
            for (int p = 0; p < runs[r].count; ++p)
                CHECK(perm[at_run + p] == new_idx[key[runs[r].start + p].idx]);
            at_run += runs[r].count;
        }
        for (int f = 0, s = 0; f < n_frames; ++f) {
            int c = 0;
            for (int j = 0; j < count[f]; ++j) c += keep[src_start[f] + j];
            CHECK(L.frame_start[f] == s);
            CHECK(L.frame_count[f] == c);
            s += c;
        }
        for (int i = 0; i < n; ++i) {
            if (!keep[i]) continue;
            const int o = new_idx[i];
            const Reflection &r = *all[i];
            CHECK(h[o] == static_cast<int16_t>(r.h));
            CHECK(I[o] == r.I);
            CHECK(corr[o] == r.image_scale_corr);
            CHECK(dphi[o] == r.delta_phi_deg);
            CHECK(img[o] == r.image_number);
        }
    }
}

#endif