// 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 <cmath>
#include <algorithm>
#include <random>
#include <tuple>
#include <vector>

#include "../common/CUDAWrapper.h"
#include "../common/DetectorSetup.h"
#include "../common/DiffractionExperiment.h"
#include "../common/PixelMask.h"
#include "../image_analysis/hot_pixels/HotPixels.h"

namespace {
    constexpr int W = 257, H = 257, C = 128;
    constexpr int NFRAMES = 60;
    constexpr double OSC_DEG = 0.1, SPACING_DEG = 6.0;   // 60 frames spread over a full turn

    constexpr size_t I(int x, int y) { return static_cast<size_t>(y) * W + x; }
}

// A Poisson background with a powder ring, and on it four kinds of pixel: one that reads 60 counts
// high on every frame, one only 20 high - persistent, but too weak to make an outlier - one that holds
// the error value on every frame, and one crossed by a genuine reflection. The last sits close to the rotation axis, where one reflection stays on a pixel for a
// long stretch of rotation - here nine consecutive sampled frames, more than the chance bound allows
// but fewer than the one-reflection bound at that zeta. Only the first and the third are masked - and
// a 3x3 patch reading high on every frame, which is stationary in the lab and so no reflection either.
// Three dark pixels besides: one dead on the powder ring, which is masked, one on the ring reading 70%
// of it, and one dead on the 5-count background, whose deficit is too small to make an outlier.
TEST_CASE("HotPixelFinder_PersistentPixelNotBragg", "[HotPixelFinder]") {
    DiffractionExperiment x(DetDECTRIS(W, H, "Test detector", ""));
    x.IncidentEnergy_keV(WVL_1A_IN_KEV).DetectorDistance_mm(10.0f);
    x.BeamX_pxl(static_cast<float>(C)).BeamY_pxl(static_cast<float>(C));
    x.Goniometer(GoniometerAxis("omega", 0.0f, static_cast<float>(OSC_DEG), Coord(1, 0, 0), std::nullopt));
    const PixelMask pixel_mask(x);
    HotPixelFinder finder(x, pixel_mask, 4);

    constexpr int HOT_X = 200, HOT_Y = 200, WARM_X = 190, WARM_Y = 60, ERR_X = 60, ERR_Y = 190;
    constexpr int BRAGG_X = 40, BRAGG_Y = 115;
    constexpr int PATCH_X = 150, PATCH_Y = 225;
    constexpr int DEAD_X = C + 82, DEAD_Y = C, DIM_X = C, DIM_Y = C + 82, DEAD_BKG_X = 20, DEAD_BKG_Y = 20;
    std::mt19937 rng(1);
    // The powder ring is a smooth radial profile, as a real one is: 45 counts over the background at
    // 82 px, 4 px sigma.
    std::vector<int32_t> frame(static_cast<size_t>(W) * H), scratch;
    for (int f = 0; f < NFRAMES; f++) {
        for (int y = 0; y < H; y++)
            for (int x_ = 0; x_ < W; x_++) {
                const double r = std::hypot(x_ - C, y - C);
                const double mean = 5.0 + 45.0 * std::exp(-0.5 * (r - 82.0) * (r - 82.0) / 16.0);
                frame[I(x_, y)] = std::poisson_distribution<int>(mean)(rng);
            }
        frame[I(HOT_X, HOT_Y)] += 60;
        frame[I(WARM_X, WARM_Y)] += 20;
        frame[I(ERR_X, ERR_Y)] = INT32_MIN;
        for (int dy = -1; dy <= 1; dy++)
            for (int dx = -1; dx <= 1; dx++)
                frame[I(PATCH_X + dx, PATCH_Y + dy)] += 60;
        if (f >= 20 && f < 29)
            frame[I(BRAGG_X, BRAGG_Y)] += 500;
        frame[I(DEAD_X, DEAD_Y)] = 0;
        frame[I(DIM_X, DIM_Y)] = static_cast<int32_t>(0.7 * frame[I(DIM_X, DIM_Y)]);
        frame[I(DEAD_BKG_X, DEAD_BKG_Y)] = 0;
        finder.AddImage(frame.data(), scratch);
    }

    const auto result = finder.GetMask(OSC_DEG, SPACING_DEG, 4);
    CHECK(result.frames == NFRAMES);
    CHECK(result.mask[I(HOT_X, HOT_Y)] == 1);   // 1 = hot, 2 = error value, 3 = dead
    CHECK(result.mask[I(ERR_X, ERR_Y)] == 2);
    CHECK(result.mask[I(WARM_X, WARM_Y)] == 0);
    CHECK(result.mask[I(BRAGG_X, BRAGG_Y)] == 0);
    for (int dy = -1; dy <= 1; dy++)
        for (int dx = -1; dx <= 1; dx++)
            CHECK(result.mask[I(PATCH_X + dx, PATCH_Y + dy)] == 1);
    CHECK(result.mask[I(DEAD_X, DEAD_Y)] == 3);
    CHECK(result.mask[I(DIM_X, DIM_Y)] == 0);
    CHECK(result.mask[I(DEAD_BKG_X, DEAD_BKG_Y)] == 0);
    CHECK(result.hot == 10);
    CHECK(result.error == 1);
    CHECK(result.dead == 1);
}

#ifdef JFJOCH_USE_CUDA
// The device path against the host one, on frames that take every branch of both: a background bright
// enough on one annulus that the host reads its rings' statistics off the exact selection rather than
// its histogram, negative counts, saturated and error pixels scattered at random, masked pixels, and a
// few hundred planted pixels whose excess and persistence straddle every threshold, and as many read at
// fractions of their counts that straddle the dead rule's, so that a level or a threshold one count off
// would move some of them across. The masks must be identical.
TEST_CASE("HotPixelFinder_DeviceMatchesHost", "[HotPixelFinder]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping HotPixelFinder_DeviceMatchesHost");
        return;
    }
    DiffractionExperiment x(DetDECTRIS(W, H, "Test detector", ""));
    x.IncidentEnergy_keV(WVL_1A_IN_KEV).DetectorDistance_mm(10.0f);
    x.BeamX_pxl(static_cast<float>(C) + 0.3f).BeamY_pxl(static_cast<float>(C) - 0.6f);
    x.Goniometer(GoniometerAxis("omega", 0.0f, static_cast<float>(OSC_DEG), Coord(1, 0, 0), std::nullopt));
    PixelMask pixel_mask(x);
    std::vector<uint32_t> masked(static_cast<size_t>(W) * H, 0);
    for (int y = 100; y < 110; y++)
        masked[I(30, y)] = 1;
    pixel_mask.LoadUserMask(x, masked);

    HotPixelFinder host(x, pixel_mask, 4), device(x, pixel_mask, 4);
    HotPixelFinderGPU::Frame frame(std::make_shared<CudaStream>());
    CudaDevicePtr<int32_t> device_image(static_cast<size_t>(W) * H);

    std::mt19937 rng(7);
    std::uniform_int_distribution<int> coord(0, W - 1);
    struct Planted { size_t i; int excess; double rate; };
    std::vector<Planted> planted;
    for (int p = 0; p < 300; p++)
        planted.push_back({I(coord(rng), coord(rng)), 5 + p / 3, 0.2 + 0.8 * (p % 7) / 6.0});
    std::vector<std::pair<size_t, double>> dark;
    for (int p = 0; p < 300; p++)
        dark.push_back({I(coord(rng), coord(rng)), 0.3 + 0.3 * (p % 13) / 12.0});

    std::vector<int32_t> image(static_cast<size_t>(W) * H), scratch;
    std::uniform_real_distribution<double> u(0.0, 1.0);
    for (int f = 0; f < NFRAMES; f++) {
        for (int y = 0; y < H; y++)
            for (int x_ = 0; x_ < W; x_++) {
                const double r = std::hypot(x_ - C, y - C);
                const double mean = 5.0 + 45.0 * std::exp(-0.5 * (r - 82.0) * (r - 82.0) / 16.0)
                                    + (r > 40.0 && r < 50.0 ? 3000.0 : 0.0);
                int32_t v = std::poisson_distribution<int>(mean)(rng) - (r > 110.0 ? 3 : 0);
                const double d = u(rng);
                if (d < 0.002) v = INT32_MIN;
                else if (d < 0.003) v = INT32_MAX;
                image[I(x_, y)] = v;
            }
        for (const auto &p : planted)
            if (u(rng) < p.rate && image[p.i] != INT32_MIN && image[p.i] != INT32_MAX)
                image[p.i] += p.excess;
        for (const auto &[i, fraction] : dark)
            if (image[i] != INT32_MIN && image[i] != INT32_MAX)
                image[i] = static_cast<int32_t>(fraction * image[i]);
        for (size_t i = 0; i < image.size(); i++)
            if (pixel_mask.GetMask()[i] != 0)
                image[i] = INT32_MIN;   // as the preprocessor leaves a masked pixel

        host.AddImage(image.data(), scratch);
        REQUIRE(cudaMemcpy(device_image, image.data(), image.size() * sizeof(int32_t), cudaMemcpyHostToDevice)
                == cudaSuccess);
        device.AddDeviceImage(device_image, frame);
        // The frame is queued, not processed: the next copy must not overwrite it before its kernels ran.
        REQUIRE(cudaStreamSynchronize(*frame.stream) == cudaSuccess);
    }

    const auto expected = host.GetMask(OSC_DEG, SPACING_DEG, 4);
    const auto result = device.GetMask(OSC_DEG, SPACING_DEG, 4);
    CHECK(expected.hot > 10);
    CHECK(expected.dead > 3);
    CHECK(result.dead == expected.dead);
    CHECK(result.frames == expected.frames);
    CHECK(result.hot == expected.hot);
    CHECK(result.error == expected.error);
    CHECK(result.mask == expected.mask);
}

// The pre-scan keeps one device half per card and folds the cards together when the mask is read. Two
// halves, each given part of the frames, folded into one must give exactly the chance counts and the
// candidates of one half given every frame.
TEST_CASE("HotPixelFinderGPU_FoldedCardsMatchOneCard", "[HotPixelFinder]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping HotPixelFinderGPU_FoldedCardsMatchOneCard");
        return;
    }
    // Four rings of two sectors: a key by band of rows and half of the detector, one column masked.
    constexpr int NRINGS = 4, SECTORS = 2;
    const size_t npixels = static_cast<size_t>(W) * H;
    std::vector<int32_t> key(npixels);
    std::vector<uint32_t> key_begin(NRINGS * SECTORS + 1, 0);
    for (int y = 0; y < H; y++)
        for (int x_ = 0; x_ < W; x_++) {
            const int k = x_ == 7 ? -1 : (y * NRINGS / H) * SECTORS + x_ * SECTORS / W;
            key[I(x_, y)] = k;
            if (k >= 0) key_begin[k + 1]++;
        }
    for (int k = 0; k < NRINGS * SECTORS; k++)
        key_begin[k + 1] += key_begin[k];
    const HotPixelLevelRules rules{HotPixelFinder::MIN_SECTOR_PIXELS, HotPixelFinder::MIN_RING_PIXELS,
                                   HotPixelFinder::LIT_NSIGMA, HotPixelFinder::LIT_OFFSET};
    HotPixelFinderGPU one(key.data(), npixels, key_begin, NRINGS, SECTORS, rules);
    HotPixelFinderGPU first(key.data(), npixels, key_begin, NRINGS, SECTORS, rules);
    HotPixelFinderGPU second(key.data(), npixels, key_begin, NRINGS, SECTORS, rules);
    HotPixelFinderGPU::Frame frame(std::make_shared<CudaStream>());
    CudaDevicePtr<int32_t> device_image(npixels);

    std::mt19937 rng(11);
    std::uniform_real_distribution<double> u(0.0, 1.0);
    std::vector<int32_t> image(npixels);
    for (int f = 0; f < NFRAMES; f++) {
        for (size_t i = 0; i < npixels; i++) {
            const double d = u(rng);
            image[i] = d < 0.01 ? INT32_MIN : d < 0.015 ? INT32_MAX : std::poisson_distribution<int>(8.0)(rng);
        }
        for (int p = 0; p < 40; p++)
            image[I(20 + p, 30 + p)] += f % 3 == 0 ? 0 : 40;   // lit on most frames
        REQUIRE(cudaMemcpy(device_image, image.data(), npixels * sizeof(int32_t), cudaMemcpyHostToDevice)
                == cudaSuccess);
        one.Add(device_image, frame);
        (f % 3 == 0 ? second : first).Add(device_image, frame);
        REQUIRE(cudaStreamSynchronize(*frame.stream) == cudaSuccess);
    }
    first.Fold(second);

    std::vector<int64_t> lit_one, seen_one, lit_folded, seen_folded;
    one.ChanceCounts(lit_one, seen_one);
    first.ChanceCounts(lit_folded, seen_folded);
    CHECK(lit_folded == lit_one);
    CHECK(seen_folded == seen_one);
    const std::vector<int> k_chance(NRINGS, 5);
    const auto expected = one.GetCandidates(NFRAMES, 10, true, k_chance);
    const auto result = first.GetCandidates(NFRAMES, 10, true, k_chance);
    CHECK(expected.index.size() >= 40);
    // The candidates are written by atomic slot, so their order is the device's; compare them by pixel.
    const auto by_pixel = [](const HotPixelFinderGPU::Candidates &c) {
        std::vector<std::tuple<uint32_t, uint16_t, uint16_t, uint16_t, int64_t, int64_t>> v;
        for (size_t j = 0; j < c.index.size(); j++)
            v.emplace_back(c.index[j], c.n_lit[j], c.n_error[j], c.n_error_ring_ok[j], c.sum_value[j],
                           c.error_level_sum[j]);
        std::sort(v.begin(), v.end());
        return v;
    };
    CHECK(by_pixel(result) == by_pixel(expected));
    CHECK(result.key_frames == expected.key_frames);
    CHECK(result.key_level_sum == expected.key_level_sum);
    CHECK(result.key_local_level_sum == expected.key_local_level_sum);
}
#endif