// 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.GetBackgroundPixelGain() == 1.0f);
    s.BackgroundPixelGain(4.5f);
    CHECK(s.GetBackgroundPixelGain() == 4.5f);
    CHECK_THROWS_AS(s.BackgroundPixelGain(0.5f), JFJochException);
}

namespace {

// A background of mean photons per pixel, each worth gain units, with Gaussian noise of sigma units, the
// pedestal subtracted.
std::vector<uint64_t> PhotonHistogram(double gain, double sigma, uint32_t seed, double mean = 0.2) {
    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);
    for (auto &v: img)
        v = static_cast<int32_t>(std::lround(gain * photons(rng) + noise(rng)));
    std::vector<uint64_t> histogram(background_gain::PEAK_HISTOGRAM_BINS, 0);
    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 / 16.0, Catch::Matchers::WithinAbs(4.57, 0.05));

    // 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, 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());
}