// Copyright (2019-2022) Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-or-later

#include <catch2/catch.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;
        uint16_t mask_value;
        DiffractionExperiment x(1, {1});

        switch (gain_level) {
            case 1:
                base_value = 0x4000;
                wrong_value = 0xc000;
                mask_value = 4;
                x.Mode(DetectorMode::PedestalG1);
                break;
            case 2:
                base_value = 0xc000;
                wrong_value = 0;
                mask_value = 8;
                x.Mode(DetectorMode::PedestalG2);
                break;
            default:
                base_value = 0;
                wrong_value = 0x4000;
                mask_value = 2;
                x.Mode(DetectorMode::Conversion);
                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;
        // Previously these bits were part of the mask - checking if they are cleared properly
        calib.GetPedestalMask()[3] = mask_value;
        calib.GetPedestalMask()[1] = mask_value;

        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);

        REQUIRE(calib.GetPedestalMask()[0] == 0);
        REQUIRE(calib.GetPedestalMask()[1] == 0);
        REQUIRE(calib.GetPedestalMask()[2] == mask_value); // Wrong gain
        REQUIRE(calib.GetPedestalMask()[3] == 0);

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

        REQUIRE(calib.GetPedestalMask()[3] == 0); // RMS condition
        REQUIRE(calib.GetPedestalMask()[2] == mask_value); // Wrong gain
        REQUIRE(calib.GetPedestalMask()[1] == mask_value); // Wrong gain
        REQUIRE(calib.GetPedestalMask()[0] == mask_value); // Wrong gain
    }
}

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

    // No images at all
    calc.Export(calib);
    REQUIRE(calib.GetPedestal()[0] == PEDESTAL_WRONG);
    REQUIRE(calib.GetPedestalMask()[511*1024+33] == 1 << 3);

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

    // 127 images
    calc.Export(calib);
    REQUIRE(calib.GetPedestal()[0] == PEDESTAL_WRONG);
    REQUIRE(calib.GetPedestalMask()[511*1024+33] == 1 << 3);

    size_t cnt = 0;
    for (int i = 0; i < RAW_MODULE_SIZE; i++)
        if (calib.GetPedestalMask()[i] != 0) cnt++;
    REQUIRE(cnt == RAW_MODULE_SIZE);

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

TEST_CASE("JFPedestal_line0","[JFModulePedestal]") {
    DiffractionExperiment x(1, {1});
    x.Mode(DetectorMode::PedestalG0);
    JFModulePedestal calib;

    std::vector<uint16_t> image(RAW_MODULE_SIZE);
    for (int i = 0; i < RAW_MODULE_SIZE; i++)
        image[i] = i / RAW_MODULE_COLS; // per line number

    JFPedestalCalc calc1(x,128, 384);
    for (int i = 0; i < PEDESTAL_WINDOW_SIZE; i++)
        calc1.AnalyzeImage(image.data());

    calc1.Export(calib);
    REQUIRE(calib.GetPedestal()[16 * RAW_MODULE_COLS] == 0);
    REQUIRE(calib.GetPedestalMask()[0] == 0);
    REQUIRE(calib.GetPedestal()[384 * RAW_MODULE_COLS] == 384);
    REQUIRE(calib.GetPedestal()[432 * RAW_MODULE_COLS + 67] == 432);
    REQUIRE(calib.GetPedestal()[RAW_MODULE_SIZE - 1] == 511);

    for (int i = 384 * RAW_MODULE_COLS; i < RAW_MODULE_SIZE; i++)
        image[i] = 0; // clear image for the already processed part - this is to ensure that second pass will not recalculate this part

    JFPedestalCalc calc2(x,128, 0);
    for (int i = 0; i < PEDESTAL_WINDOW_SIZE; i++)
        calc2.AnalyzeImage(image.data());

    calc2.Export(calib);
    REQUIRE(calib.GetPedestal()[0] == 0);
    REQUIRE(calib.GetPedestal()[16 * RAW_MODULE_COLS] == 16);
    REQUIRE(calib.GetPedestal()[256 * RAW_MODULE_COLS] == 0);
    REQUIRE(calib.GetPedestal()[384 * RAW_MODULE_COLS] == 384);
    REQUIRE(calib.GetPedestal()[432 * RAW_MODULE_COLS + 67] == 432);
    REQUIRE(calib.GetPedestal()[RAW_MODULE_SIZE - 1] == 511);

    uint64_t wrong_mask = 0;
    for (int i = 0; i < RAW_MODULE_SIZE; i++) {
        if (calib.GetPedestalMask()[i] != 0) wrong_mask++;
    }
    REQUIRE(wrong_mask == 0);
}

TEST_CASE("JFPedestal_wrong_size","[JFModulePedestal]") {
    DiffractionExperiment x(1, {1});

    std::unique_ptr<JFPedestalCalc> calc;

    REQUIRE_THROWS(calc = std::make_unique<JFPedestalCalc>(x, -1));
    REQUIRE_THROWS(calc = std::make_unique<JFPedestalCalc>(x, 0));
    REQUIRE_THROWS(calc = std::make_unique<JFPedestalCalc>(x, 1, -1));
    REQUIRE_THROWS(calc = std::make_unique<JFPedestalCalc>(x, 512, 1));
    REQUIRE_THROWS(calc = std::make_unique<JFPedestalCalc>(x, 5120, 10));
}