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

#include "../common/CUDAWrapper.h"
#include "../common/DetectorSetup.h"
#include "../common/DiffractionExperiment.h"
#include "../common/PixelMask.h"
#include "../rugnux/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.
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;
    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;
        if (f >= 20 && f < 29)
            frame[I(BRAGG_X, BRAGG_Y)] += 500;
        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
    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);
    CHECK(result.hot == 1);
    CHECK(result.error == 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, 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<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 (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);
    }

    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(result.frames == expected.frames);
    CHECK(result.hot == expected.hot);
    CHECK(result.error == expected.error);
    CHECK(result.mask == expected.mask);
}
#endif