// 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 <vector>

#include "../common/CUDAWrapper.h"
#include "../common/DiffractionExperiment.h"
#include "../common/PixelMask.h"
#include "../writer/HDF5Objects.h"
#include "../image_analysis/bragg_integration/SpotWidth.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorCPU.h"
#include "../image_analysis/spot_finding/ImageSpotFinderCPU.h"
#ifdef JFJOCH_USE_CUDA
#include "../image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h"
#endif

// The pre-scan measures the spot width and the footprint on squares cut out about each spot instead
// of the whole frame (MeasureSpotsOnTiles). That is only allowed because it gives the whole frame's
// answer to the bit, so this checks exactly that on real diffraction frames - including the spots
// whose footprint window walks off its square and is measured again, on a larger square or on the
// frame, and the frame whose squares would be more than the frame itself.
namespace {

struct Frame {
    DiffractionExperiment x;
    ImagePreprocessorBuffer image;
    std::vector<DiffractionSpot> spots;
};

std::vector<Frame> ReadFrames() {
    HDF5ReadOnlyFile data("../../tests/test_data/compression_benchmark.h5");
    HDF5DataSet dataset(data, "/entry/data/data");
    HDF5DataSpace file_space(dataset);
    const auto dim = file_space.GetDimensions();

    DiffractionExperiment x(DetDECTRIS(dim[2], dim[1], "Test", {}));
    x.PixelSigned(true).DetectorDistance_mm(75).BeamY_pxl(1136).BeamX_pxl(1090).IncidentEnergy_keV(12.4);
    const PixelMask mask(x);
    ImagePreprocessorCPU preprocessor(x, mask);
    ImageSpotFinderCPU finder(static_cast<int32_t>(x.GetXPixelsNumConv()), static_cast<int32_t>(x.GetYPixelsNumConv()));

    std::vector<Frame> frames;
    std::vector<int16_t> raw(dim[1] * dim[2]);
    for (hsize_t i = 0; i < dim[0]; i += 4) {
        dataset.ReadVector(raw, {i, 0, 0}, {1, dim[1], dim[2]});
        Frame f{x, ImagePreprocessorBuffer(x.GetPixelsNum()), {}};
        preprocessor.Analyze(f.image, reinterpret_cast<const uint8_t *>(raw.data()), CompressedImageMode::Int16);
        f.spots = finder.Run(f.image, SpotFindingSettings{});
        frames.push_back(std::move(f));
    }
    return frames;
}

bool Same(const std::vector<spot_width::FluxCurve> &a, const std::vector<spot_width::FluxCurve> &b) {
    if (a.size() != b.size()) return false;
    for (size_t k = 0; k < a.size(); k++) {
        const auto &p = a[k], &q = b[k];
        if (p.d_A != q.d_A || p.c != q.c || p.shape.valid != q.shape.valid || p.shape.two_theta != q.shape.two_theta
            || p.shape.jr != q.shape.jr || p.shape.jt != q.shape.jt || p.shape.m_rad != q.shape.m_rad
            || p.shape.m_tan != q.shape.m_tan || p.shape.cos_psi != q.shape.cos_psi
            || p.shape.par_u != q.shape.par_u || p.shape.par_v != q.shape.par_v)
            return false;
    }
    return true;
}

bool Same(const std::vector<FootprintSpot> &a, const std::vector<FootprintSpot> &b) {
    if (a.size() != b.size()) return false;
    for (size_t k = 0; k < a.size(); k++)
        if (a[k].r_px != b[k].r_px || a[k].sigma_rad != b[k].sigma_rad || a[k].sigma_tan != b[k].sigma_tan)
            return false;
    return true;
}

// Every frame through MeasureSpotsOnTiles on `image(f)`, against the whole frame; the misses summed.
template <class Image>
SpotTileMisses CheckFrames(std::vector<Frame> &frames, Image &&image, int half, int footprint_half) {
    SpotTileMisses total;
    size_t curves = 0, footprints = 0;
    for (auto &f : frames) {
        const int width = static_cast<int>(f.x.GetXPixelsNumConv()), height = static_cast<int>(f.x.GetYPixelsNumConv());
        const auto &g = f.x.GetDiffractionGeometry();
        std::vector<spot_width::FluxCurve> c_frame, c_tiles;
        std::vector<FootprintSpot> f_frame, f_tiles;
        MeasureSpotFluxCurves(f.image, width, height, g, f.spots, c_frame);
        std::vector<float> sx, sy;
        for (const auto &s : f.spots) {
            sx.push_back(s.RawCoord().x);
            sy.push_back(s.RawCoord().y);
        }
        MeasureFootprintSpots(f.image.data(), width, height, g.GetBeamX_pxl(), g.GetBeamY_pxl(), sx, sy, f_frame);

        const auto misses = MeasureSpotsOnTiles(image(f), width, height, g, f.spots, c_tiles, f_tiles,
                                                half, footprint_half);
        CHECK(Same(c_frame, c_tiles));
        CHECK(Same(f_frame, f_tiles));
        total.larger += misses.larger;
        total.whole_frame += misses.whole_frame;
        curves += c_frame.size();
        footprints += f_frame.size();
    }
    // Something was measured, or the comparison says nothing.
    CHECK(curves > 0);
    CHECK(footprints > 100);
    return total;
}

}   // namespace

TEST_CASE("SpotTiles_MatchWholeFrame", "[SpotWidth]") {
    auto frames = ReadFrames();
    REQUIRE(!frames.empty());
    const auto host = [](Frame &f) -> ImagePreprocessorBuffer & { return f.image; };

    // As the pre-scan reads them.
    CheckFrames(frames, host, SPOT_TILE_HALF, FOOTPRINT_TILE_HALF);
    // Squares too small for the footprint: its spots are measured again on the larger squares...
    CHECK(CheckFrames(frames, host, spot_width::R_BKG_OUT, spot_width::R_BKG_OUT + 1).larger > 0);
    // ... and on the frame, where the larger squares would be more than it.
    CHECK(CheckFrames(frames, host, spot_width::R_BKG_OUT, 4000).whole_frame > 0);
    // Squares so large that the frame is read instead from the start.
    CHECK(CheckFrames(frames, host, 1000, FOOTPRINT_TILE_HALF).whole_frame == 0);
}

#ifdef JFJOCH_USE_CUDA
// The same with the squares cut out on the card, from a device image the host never holds.
TEST_CASE("SpotTiles_MatchWholeFrame_GPU", "[SpotWidth][GPU]") {
    if (get_gpu_count() == 0)
        SKIP("No GPU");
    auto frames = ReadFrames();
    REQUIRE(!frames.empty());
    ImagePreprocessorBufferGPU device(frames.front().x.GetPixelsNum(), /*host_mirror=*/false);
    const auto upload = [&](Frame &f) -> ImagePreprocessorBuffer & {
        REQUIRE(cudaMemcpy(device.getGPUBuffer(), f.image.data(), f.image.size() * sizeof(int32_t),
                           cudaMemcpyHostToDevice) == cudaSuccess);
        return device;
    };
    CheckFrames(frames, upload, SPOT_TILE_HALF, FOOTPRINT_TILE_HALF);
    CHECK(CheckFrames(frames, upload, spot_width::R_BKG_OUT, spot_width::R_BKG_OUT + 1).larger > 0);
    CHECK(CheckFrames(frames, upload, spot_width::R_BKG_OUT, 4000).whole_frame > 0);
    CheckFrames(frames, upload, 1000, FOOTPRINT_TILE_HALF);
}
#endif