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

#include <random>
#include <catch2/catch_all.hpp>
#include "../jungfrau/JFPedestalCalc.h"

TEST_CASE("JFPedestalCalc", "[JFPedestalCalc]") {
    for (int gain_level = 0; gain_level < 3; gain_level++) {
        uint16_t base_value;
        uint16_t wrong_value;
        DiffractionExperiment x(DetJF(1));

        switch (gain_level) {
            case 1:
                base_value = 0x4000;
                wrong_value = 0xc000;
                x.Mode(DetectorMode::PedestalG1);
                break;
            case 2:
                base_value = 0xc000;
                wrong_value = 0;
                x.Mode(DetectorMode::PedestalG2);
                break;
            default:
                base_value = 0;
                wrong_value = 0x4000;
                x.Mode(DetectorMode::PedestalG0);
                break;
        }

        std::vector<uint16_t> image(RAW_MODULE_SIZE, base_value);

        for (int i = 0; i < RAW_MODULE_SIZE; i++)
            image[i] += i;

        image[2] = wrong_value;

        JFPedestalCalc calc(x);

        // Predictable random number generator
        std::mt19937 g1(0);
        std::normal_distribution<double> distribution(12000.0, 12.5);

        for (int i = 0; i < 299; i++) {
            image[3] = base_value + distribution(g1);
            calc.AnalyzeImage(image.data());
        }
        image[3] = base_value + distribution(g1);

        image[0] = wrong_value;
        image[1] = wrong_value;
        image[2] = base_value;

        calc.AnalyzeImage(image.data());

        JFModulePedestal calib;

        calc.Export(calib, 1);

        for (int i = 4; i < RAW_MODULE_COLS; i++)
            REQUIRE(calib.GetPedestal()[i] == i);

        REQUIRE(calib.GetPedestal()[3] > 11900.0);
        REQUIRE(calib.GetPedestal()[3] < 12100.0);

        REQUIRE(calib.GetPedestal()[0] == 0);

        calc.Export(calib, 0);
        auto pedestal = calib.GetPedestal();
        for (int i = 4; i < RAW_MODULE_COLS; i++) {
            REQUIRE(pedestal[i] == i);
        }
    }
}

int16_t r(std::mt19937 &g, std::normal_distribution<double> &dist) {
    double tmp = -5;
    while ((tmp < 0) || (tmp > INT16_MAX))
        tmp = dist(g);
    return tmp;
}

TEST_CASE("JFPedestalCalc_MaskRMS", "[JFPedestalCalc]") {
    DiffractionExperiment x(DetJF(1));

    std::vector<uint16_t> image(RAW_MODULE_SIZE, 0);

    JFPedestalCalc calc(x);

    // Predictable random number generator
    std::mt19937 g1(0);
    std::normal_distribution<double> distribution0(12000.0, 30);
    std::normal_distribution<double> distribution1(5000.0, 50);
    std::normal_distribution<double> distribution2(1000.0, 65);
    std::normal_distribution<double> distribution3(1000.0, 70);

    for (int i = 0; i < 1024; i++) {
        image[0] = r(g1, distribution0);
        image[1] = r(g1, distribution1);
        image[2] = r(g1, distribution2);
        image[511 * 1024 + 34] = r(g1, distribution3);
        calc.AnalyzeImage(image.data());
    }

    JFModulePedestal calib;
    calc.Export(calib, 1);

    REQUIRE(fabs(calib.GetPedestalRMS()[0]-30) < 3.0);
    REQUIRE(fabs(calib.GetPedestalRMS()[1]-50) < 3.0);
    REQUIRE(fabs(calib.GetPedestalRMS()[2]-65) < 3.0);
    REQUIRE(fabs(calib.GetPedestalRMS()[511 * 1024 + 34]-70) < 3.0);
}

TEST_CASE("PedestalCalc_ImagesLessThanWindow", "[JFPedestalCalc]") {
    DiffractionExperiment x(DetJF(1));
    x.Mode(DetectorMode::PedestalG2);
    JFModulePedestal calib;
    JFPedestalCalc calc(x);

    // No images at all
    calc.Export(calib);
    REQUIRE(calib.GetPedestal()[0] == PEDESTAL_WRONG);

    std::vector<uint16_t> image(RAW_MODULE_SIZE, 0xc000 + 12000);
    for (int i = 0; i < PEDESTAL_MIN_IMAGE_COUNT- 1; i++)
        calc.AnalyzeImage(image.data());

    // 127 images
    calc.Export(calib);
    REQUIRE(calib.GetPedestal()[0] == PEDESTAL_WRONG);
    REQUIRE(calib.CountMaskedPixels() == RAW_MODULE_SIZE);

    // 128 images
    calc.AnalyzeImage(image.data());
    calc.Export(calib);
    REQUIRE(calib.GetPedestal()[0] == 12000);
    REQUIRE(calib.GetPedestal()[511*1024+33] == 12000);
}