// Copyright (2019-2023) Paul Scherrer Institute

#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;
        uint16_t mask_value;
        DiffractionExperiment x(DetectorGeometry(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
    }
}

#include <iostream>

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(DetectorGeometry(1));

    x.Mode(DetectorMode::Conversion);

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

    // calc.AnalyzeImage(image.data());

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

    std::cout << calib.GetPedestal()[0]  << " " << calib.GetPedestalRMS()[0] << std::endl;
    std::cout << calib.GetPedestal()[1]  << " " << calib.GetPedestalRMS()[1] << std::endl;
    std::cout << calib.GetPedestal()[2]  << " " << calib.GetPedestalRMS()[2] << std::endl;
    std::cout << calib.GetPedestal()[511*1024+34]  << " " << calib.GetPedestalRMS()[511*1024+34] << std::endl;

    REQUIRE(calib.GetPedestalMask()[0] == 0);
    REQUIRE(calib.GetPedestalMask()[1] == 0);
    REQUIRE(calib.GetPedestalMask()[2] == MASK_PEDESTAL_G0_RMS_LIMIT );
    REQUIRE(calib.GetPedestalMask()[511*1024+34] == MASK_PEDESTAL_G0_RMS_LIMIT );

    calc.Export(calib, 1, 20);

    REQUIRE(calib.GetPedestalMask()[0] == MASK_PEDESTAL_G0_RMS_LIMIT );
    REQUIRE(calib.GetPedestalMask()[1] == MASK_PEDESTAL_G0_RMS_LIMIT );
    REQUIRE(calib.GetPedestalMask()[2] == MASK_PEDESTAL_G0_RMS_LIMIT );
    REQUIRE(calib.GetPedestalMask()[511*1024+34] == MASK_PEDESTAL_G0_RMS_LIMIT );

    calc.Export(calib, 1, 80);

    REQUIRE(calib.GetPedestalMask()[0] == 0);
    REQUIRE(calib.GetPedestalMask()[1] == 0);
    REQUIRE(calib.GetPedestalMask()[2] == 0);
    REQUIRE(calib.GetPedestalMask()[511*1024+34] == 0);
}

TEST_CASE("PedestalCalc_ImagesLessThanWindow", "[JFPedestalCalc]") {
    DiffractionExperiment x(DetectorGeometry(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_MIN_IMAGE_COUNT- 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);
}