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

#include "../common/BraggIntegrationSettings.h"
#include "../common/CompressedImage.h"
#include "../common/DetectorSetup.h"
#include "../common/DiffractionExperiment.h"
#include "../common/Reflection.h"
#include "../image_analysis/bragg_integration/BraggIntegrationEngineCPU.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorBuffer.h"

// The FPGA workflow integrates straight off the assembled CompressedImage (any pixel type) instead of
// a preprocessed int32 buffer, reading only the reflection disks. These tests pin that the typed
// sampler produces exactly the same intensities as the equivalent int32 preprocessed buffer - i.e. the
// zero-copy route is numerically identical to the buffer route, not merely close.

namespace {

struct Scene {
    std::vector<int32_t> image;   // INT32_MIN = masked, INT32_MAX = saturated
    std::vector<Reflection> predicted;
};

Reflection MakeReflection(float x, float y, float d, int hkl) {
    Reflection r{};
    r.h = hkl; r.k = hkl; r.l = hkl;
    r.predicted_x = x;
    r.predicted_y = y;
    r.d = d;
    r.rlp = 1.0f;
    r.partiality = 1.0f;
    return r;
}

// A grid of clean Gaussian spots on a flat background, values kept well inside int16 range so the same
// scene can be represented as Int16 and Int32 without clipping. A few masked/saturated pixels exercise
// the sentinel handling.
Scene BuildScene(size_t width, size_t height) {
    Scene s;
    s.image.assign(width * height, 12);
    const int margin = 45, spacing = 60;
    int hkl = 1;
    for (int gy = 0; margin + gy * spacing < static_cast<int>(height) - margin; ++gy)
        for (int gx = 0; margin + gx * spacing < static_cast<int>(width) - margin; ++gx) {
            const float cx = static_cast<float>(margin + gx * spacing) + 0.3f;
            const float cy = static_cast<float>(margin + gy * spacing) - 0.2f;
            const double amp = 150.0 + 40.0 * ((gx * 7 + gy * 13) % 20);   // <= ~950, safe for int16
            const double sigma = 1.3;
            for (int dy = -6; dy <= 6; ++dy)
                for (int dx = -6; dx <= 6; ++dx) {
                    const int x = static_cast<int>(std::lround(cx)) + dx;
                    const int y = static_cast<int>(std::lround(cy)) + dy;
                    if (x < 0 || y < 0 || x >= static_cast<int>(width) || y >= static_cast<int>(height)) continue;
                    const double ex = x - cx, ey = y - cy;
                    s.image[y * width + x] += static_cast<int32_t>(std::lround(amp * std::exp(-(ex * ex + ey * ey) / (2.0 * sigma * sigma))));
                }
            const float d = 1.4f + 0.12f * static_cast<float>((gx + gy) % 12);
            s.predicted.push_back(MakeReflection(cx, cy, d, hkl++));
        }
    for (int k = 0; k < 20; ++k) {
        const size_t idx = (static_cast<size_t>(k) * 2654435761u) % s.image.size();
        s.image[idx] = (k % 2) ? INT32_MIN : INT32_MAX;
    }
    return s;
}

DiffractionExperiment MakeExperiment(IntegratorMode mode) {
    DiffractionExperiment experiment(DetJF(2));
    experiment.DetectorDistance_mm(100.0f).IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(400.0f).BeamY_pxl(400.0f);
    BraggIntegrationSettings settings;
    settings.Integrator(mode);
    experiment.ImportBraggIntegrationSettings(settings);
    return experiment;
}

void RequireIdentical(const std::vector<Reflection> &a, const std::vector<Reflection> &b) {
    REQUIRE(a.size() == b.size());
    REQUIRE(a.size() > 40);
    for (size_t i = 0; i < a.size(); ++i) {
        INFO("reflection " << i << " hkl " << a[i].h);
        CHECK(a[i].h == b[i].h);
        CHECK(a[i].I == b[i].I);
        CHECK(a[i].sigma == b[i].sigma);
        CHECK(a[i].bkg == b[i].bkg);
    }
}

}   // namespace

TEST_CASE("BraggIntegrationEngineCPU_CompressedImageMatchesBuffer", "[Integration]") {
    for (const auto mode : {IntegratorMode::BoxSum, IntegratorMode::ProfileGaussian, IntegratorMode::ProfileEmpirical}) {
        const DiffractionExperiment experiment = MakeExperiment(mode);
        const size_t W = experiment.GetXPixelsNum(), H = experiment.GetYPixelsNum();
        const size_t npixel = experiment.GetPixelsNum();
        const Scene scene = BuildScene(W, H);
        REQUIRE(scene.image.size() == npixel);

        BraggIntegrationEngineCPU engine(experiment);

        // Route A: the preprocessed int32 buffer.
        ImagePreprocessorBuffer buffer(npixel);
        for (size_t i = 0; i < npixel; ++i) buffer[i] = scene.image[i];
        const auto from_buffer = engine.Run(buffer, scene.predicted, scene.predicted.size(), 7);

        SECTION("Int32 CompressedImage") {
            const CompressedImage image(scene.image, W, H);
            const auto from_image = engine.Run(image, scene.predicted, scene.predicted.size(), 7);
            RequireIdentical(from_buffer, from_image);
        }

        SECTION("Int16 CompressedImage") {
            // Same scene as int16 (masked -> INT16_MIN, saturated -> INT16_MAX) with a matching int32
            // buffer built by the exact sampler mapping; the two must integrate identically.
            std::vector<int16_t> img16(npixel);
            std::vector<int32_t> buf32(npixel);
            for (size_t i = 0; i < npixel; ++i) {
                const int32_t v = scene.image[i];
                if (v == INT32_MIN)      { img16[i] = INT16_MIN; buf32[i] = INT32_MIN; }
                else if (v == INT32_MAX) { img16[i] = INT16_MAX; buf32[i] = INT32_MAX; }
                else                     { img16[i] = static_cast<int16_t>(v); buf32[i] = v; }
            }
            ImagePreprocessorBuffer buffer16(npixel);
            for (size_t i = 0; i < npixel; ++i) buffer16[i] = buf32[i];
            const auto ref16 = engine.Run(buffer16, scene.predicted, scene.predicted.size(), 7);

            const CompressedImage image(img16, W, H);
            const auto from_image = engine.Run(image, scene.predicted, scene.predicted.size(), 7);
            RequireIdentical(ref16, from_image);
        }
    }
}