// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>

#include <cmath>
#include <random>

#include "../image_analysis/bragg_integration/BackgroundGain.h"
#include "../common/BraggIntegrationSettings.h"
#include "../common/JFJochException.h"

namespace {

constexpr size_t W = 512, H = 512;

// A smooth background of 100..2000 photons per pixel, read out as bias + gain * photons + read noise,
// with a grid of bright 5x5 "Bragg" spots and a masked strip.
std::vector<int32_t> Frame(double gain, uint32_t seed, std::vector<uint32_t> &mask) {
    std::mt19937 rng(seed);
    std::normal_distribution<double> read_noise(0.0, 3.0);
    std::vector<int32_t> img(W * H);
    mask.assign(W * H, 0);
    for (size_t y = 0; y < H; y++)
        for (size_t x = 0; x < W; x++) {
            const double photons = 100.0 + 1900.0 * (x + y) / static_cast<double>(W + H);
            std::poisson_distribution<int> p(photons);
            double v = (gain > 1.0 ? 20.0 : 0.0) + gain * p(rng) + (gain > 1.0 ? read_noise(rng) : 0.0);
            if (x % 32 >= 14 && x % 32 < 19 && y % 32 >= 14 && y % 32 < 19)
                v += 50.0 * gain * photons;
            img[y * W + x] = static_cast<int32_t>(std::lround(v));
            if (y >= 200 && y < 210)
                mask[y * W + x] = 1;
        }
    return img;
}

}

TEST_CASE("BackgroundGain_CCD_gain4") {
    std::vector<background_gain::Levels> frames;
    std::vector<uint32_t> mask;
    for (uint32_t seed = 1; seed <= 3; seed++) {
        const auto img = Frame(4.0, seed, mask);
        frames.push_back(background_gain::MeasureFrame(img, W, H, mask, 1 << 30));
    }
    const auto g = background_gain::Gain(frames);
    REQUIRE(g.has_value());
    CHECK_THAT(*g, Catch::Matchers::WithinAbs(4.0, 0.25));
}

TEST_CASE("BackgroundGain_photon_counting") {
    std::vector<uint32_t> mask;
    const auto img = Frame(1.0, 7, mask);
    const auto g = background_gain::Gain({background_gain::MeasureFrame(img, W, H, mask, 1 << 30)});
    REQUIRE(g.has_value());
    CHECK_THAT(*g, Catch::Matchers::WithinAbs(1.0, 0.1));
}

TEST_CASE("BackgroundGain_saturated_and_empty") {
    std::vector<uint32_t> mask;
    const auto img = Frame(4.0, 11, mask);
    // Everything at or above the saturation limit: nothing to measure.
    CHECK_FALSE(background_gain::Gain({background_gain::MeasureFrame(img, W, H, mask, 1)}).has_value());
    CHECK_FALSE(background_gain::Gain({}).has_value());
}

TEST_CASE("BackgroundGain_settings") {
    BraggIntegrationSettings s;
    CHECK(s.GetPixelGain() == 1.0f);
    s.PixelGain(4.5f);
    CHECK(s.GetPixelGain() == 4.5f);
    CHECK_THROWS_AS(s.PixelGain(0.5f), JFJochException);
}

namespace {

// A background of mean photons per pixel, each worth gain units, with Gaussian noise of sigma units, on a
// baseline of pedestal units (0: subtracted); the value taken in steps of `step` (a float file of whole
// numbers read at a fraction of its unit), over a number of frames.
std::vector<uint64_t> PhotonHistogram(double gain, double sigma, uint32_t seed, double mean = 0.2,
                                      double pedestal = 0.0, int step = 1, int frames = 1) {
    std::mt19937 rng(seed);
    std::poisson_distribution<int> photons(mean);
    std::normal_distribution<double> noise(0.0, sigma);
    std::vector<int32_t> img(W * H);
    std::vector<uint64_t> histogram(background_gain::PEAK_HISTOGRAM_BINS, 0);
    for (int f = 0; f < frames; f++) {
        for (auto &v: img)
            v = step * static_cast<int32_t>(std::lround(pedestal + gain * photons(rng) + noise(rng)));
        background_gain::AddToHistogram(img, std::vector<uint32_t>(W * H, 0), 1 << 30, histogram);
    }
    return histogram;
}

}

TEST_CASE("BackgroundGain_PhotonPeak") {
    // A photon of 4.57 keV in a file in keV, read at a sixteenth of a keV, 0.3 keV of noise.
    const auto g = background_gain::PhotonPeak(PhotonHistogram(4.57 * 16, 0.3 * 16, 1));
    REQUIRE(g.has_value());
    CHECK_THAT(g->gain / 16.0, Catch::Matchers::WithinAbs(4.57, 0.05));
    CHECK(g->pedestal == 0.0f);

    // Whole detector units, 12 per photon.
    const auto g12 = background_gain::PhotonPeak(PhotonHistogram(12.0, 1.0, 2));
    REQUIRE(g12.has_value());
    CHECK_THAT(g12->gain, Catch::Matchers::WithinAbs(12.0, 0.2));

    // Integer photon counts: the single photon is the first bin, no peak past the zero peak.
    CHECK_FALSE(background_gain::PhotonPeak(PhotonHistogram(1.0, 0.0, 3)).has_value());

    // 30 photons per pixel at 4 units each: one hump, with no second peak at twice its value.
    CHECK_FALSE(background_gain::PhotonPeak(PhotonHistogram(4.0, 3.0, 4, 30.0)).has_value());

    // A CCD's background (bias + gain * photons + read noise): a hump, not a pair of photon peaks.
    std::vector<uint32_t> mask;
    std::vector<uint64_t> ccd(background_gain::PEAK_HISTOGRAM_BINS, 0);
    background_gain::AddToHistogram(Frame(4.0, 5, mask), mask, 1 << 30, ccd);
    CHECK_FALSE(background_gain::PhotonPeak(ccd).has_value());

    // A CCD's pixels outside the active area, all at zero, and its background in two humps far above, the
    // second at twice the first: no photon peaks, as the zeros are a single bin with nothing beside it.
    std::vector<uint64_t> gap(background_gain::PEAK_HISTOGRAM_BINS, 0);
    gap[0] = 800000;
    for (size_t b = 1; b < gap.size(); b++)
        gap[b] = std::lround(200000.0 * std::exp(-0.5 * std::pow((b - 2600.0) / 300.0, 2))
                             + 20000.0 * std::exp(-0.5 * std::pow((b - 5200.0) / 300.0, 2)));
    CHECK_FALSE(background_gain::PhotonPeak(gap).has_value());
}

// A frame whose baseline was never subtracted: the zero-photon peak sits at 122 units, the photons 62 apart
// with 14 units of noise - in whole units, and as a float file of whole numbers read at a sixteenth of its unit;
// 16 frames, as the gain is measured on.
TEST_CASE("BackgroundGain_PhotonPeak_Baseline") {
    const auto g = background_gain::PhotonPeak(PhotonHistogram(62.0, 14.0, 6, 0.3, 122.0, 1, 16));
    REQUIRE(g.has_value());
    CHECK_THAT(g->gain, Catch::Matchers::WithinAbs(62.0, 1.0));
    CHECK_THAT(g->pedestal, Catch::Matchers::WithinAbs(122.0, 1.0));

    const auto g16 = background_gain::PhotonPeak(PhotonHistogram(62.0, 14.0, 6, 0.3, 122.0, 16, 16));
    REQUIRE(g16.has_value());
    CHECK_THAT(g16->gain / 16.0, Catch::Matchers::WithinAbs(62.0, 1.0));
    CHECK_THAT(g16->pedestal / 16.0, Catch::Matchers::WithinAbs(122.0, 1.0));
}

// One estimate for every input in detector units: a CCD's frames show no photon peaks and read their variance per
// unit mean, a background that resolves photons reads its peaks.
TEST_CASE("BackgroundGain_Estimate") {
    std::vector<uint32_t> mask;
    const auto img = Frame(4.0, 21, mask);
    std::vector<uint64_t> ccd(background_gain::PEAK_HISTOGRAM_BINS, 0);
    background_gain::AddToHistogram(img, mask, 1 << 30, ccd);
    const auto g = background_gain::Estimate({background_gain::MeasureFrame(img, W, H, mask, 1 << 30)}, ccd);
    REQUIRE(g.has_value());
    CHECK_THAT(g->gain, Catch::Matchers::WithinAbs(4.0, 0.25));
    CHECK(g->pedestal == 0.0f);

    const auto peaks = background_gain::Estimate({}, PhotonHistogram(62.0, 14.0, 6, 0.3, 122.0, 1, 16));
    REQUIRE(peaks.has_value());
    CHECK_THAT(peaks->gain, Catch::Matchers::WithinAbs(62.0, 1.0));

    CHECK_FALSE(background_gain::Estimate({}, std::vector<uint64_t>(background_gain::PEAK_HISTOGRAM_BINS, 0)).has_value());
}