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

#include "catch2/catch_all.hpp"
#include <random>

#include "../jungfrau/JFCalibration.h"

TEST_CASE("JFCalibration_Constructor","[JFCalibration]") {
    DiffractionExperiment experiment(DetJF(8, 2));
    experiment.StorageCells(16);

    JFCalibration c(experiment);
    REQUIRE(c.GetStorageCellNum() == 16);
    REQUIRE(c.GetModulesNum() == 8);
}

TEST_CASE("JFCalibration_GainStatistics","[JFCalibration]") {
    JFCalibration calibration(4);
    std::vector<double> tmp(3 * RAW_MODULE_SIZE);
    for (int m = 0; m < calibration.GetModulesNum(); m++) {
        for (int i = 0; i < RAW_MODULE_SIZE; i++) {
            tmp[i] = 35 + m;
            tmp[i + RAW_MODULE_SIZE] = -56 + m;
            tmp[i + 2 * RAW_MODULE_SIZE] = 1/32. * m;
        }
        calibration.GainCalibration(m) = JFModuleGainCalibration(tmp);
    }

    auto s = calibration.GetModuleStatistics();

    REQUIRE(s.size() == 4);

    REQUIRE(s[0].gain_g0_mean  == Catch::Approx(35.0));
    REQUIRE(s[1].gain_g0_mean  == Catch::Approx(36.0));
    REQUIRE(s[2].gain_g0_mean  == Catch::Approx(37.0));
    REQUIRE(s[3].gain_g0_mean  == Catch::Approx(38.0));

    REQUIRE(s[0].gain_g1_mean  == Catch::Approx(-56));
    REQUIRE(s[1].gain_g1_mean  == Catch::Approx(-55));
    REQUIRE(s[2].gain_g1_mean  == Catch::Approx(-54));
    REQUIRE(s[3].gain_g1_mean  == Catch::Approx(-53));

    REQUIRE(s[0].gain_g2_mean  == Catch::Approx(0/32.));
    REQUIRE(s[1].gain_g2_mean  == Catch::Approx(1/32.));
    REQUIRE(s[2].gain_g2_mean  == Catch::Approx(2/32.));
    REQUIRE(s[3].gain_g2_mean  == Catch::Approx(3/32.));
}

TEST_CASE("JFCalibration_PedestalAndMaskStatistics","[JFCalibration]") {
    JFCalibration calibration(4, 2);

    for (int module = 0; module < 4; module++) {
        for (int i = 0; i < RAW_MODULE_SIZE; i++) {
            calibration.Pedestal(module, 0, 1).GetPedestal()[i] = 1000 * (module+1);
            calibration.Pedestal(module, 1, 1).GetPedestal()[i] = 2000 * (module+1);
            calibration.Pedestal(module, 2, 1).GetPedestal()[i] = 3000 * (module+1);
        }
    }

    calibration.Pedestal(3, 0, 1).GetPedestal()[0] = UINT16_MAX;
    calibration.Pedestal(3, 1, 1).GetPedestal()[1] = UINT16_MAX;
    calibration.Pedestal(3, 2, 1).GetPedestal()[1] = UINT16_MAX;
    calibration.Pedestal(3, 2, 1).GetPedestal()[2] = UINT16_MAX;

    auto s1 = calibration.GetModuleStatistics(1);

    REQUIRE(s1[0].pedestal_g0_mean == Catch::Approx(1000));
    REQUIRE(s1[0].pedestal_g1_mean == Catch::Approx(2000));
    REQUIRE(s1[0].pedestal_g2_mean == Catch::Approx(3000));

    REQUIRE(s1[1].pedestal_g0_mean == Catch::Approx(2000));
    REQUIRE(s1[1].pedestal_g1_mean == Catch::Approx(4000));
    REQUIRE(s1[1].pedestal_g2_mean == Catch::Approx(6000));

    REQUIRE(s1[2].pedestal_g0_mean == Catch::Approx(3000));
    REQUIRE(s1[2].pedestal_g1_mean == Catch::Approx(6000));
    REQUIRE(s1[2].pedestal_g2_mean == Catch::Approx(9000));

    REQUIRE(s1[0].bad_pixels == 0);
    REQUIRE(s1[1].bad_pixels == 0);
    REQUIRE(s1[2].bad_pixels == 0);
    REQUIRE(s1[3].bad_pixels == 3);

    auto s0 = calibration.GetModuleStatistics(0);

    REQUIRE(s0[0].pedestal_g0_mean == 0);
    REQUIRE(s0[0].pedestal_g1_mean == 0);
    REQUIRE(s0[0].pedestal_g2_mean == 0);

    REQUIRE(s0[0].bad_pixels == 0);
    REQUIRE(s0[1].bad_pixels == 0);
    REQUIRE(s0[2].bad_pixels == 0);
    REQUIRE(s0[3].bad_pixels == 0);
}

TEST_CASE("JFCalibration_Statistics_All","[JFCalibration]") {
    JFCalibration calibration(4, 2);

    for (int module = 0; module < 4; module++) {
        for (int i = 0; i < RAW_MODULE_SIZE; i++) {
            calibration.Pedestal(module, 0, 1).GetPedestal()[i] = 1000 * (module+1);
            calibration.Pedestal(module, 1, 1).GetPedestal()[i] = 2000 * (module+1);
            calibration.Pedestal(module, 2, 1).GetPedestal()[i] = 3000 * (module+1);
        }
    }

    calibration.Pedestal(3, 0, 1).GetPedestal()[0] = UINT16_MAX;
    calibration.Pedestal(3, 1, 1).GetPedestal()[1] = UINT16_MAX;
    calibration.Pedestal(3, 2, 1).GetPedestal()[2] = UINT16_MAX;
    calibration.Pedestal(3, 2, 1).GetPedestal()[2] = UINT16_MAX;

    auto s0 = calibration.GetModuleStatistics();

    REQUIRE(s0[0].module_number == 0);
    REQUIRE(s0[0].storage_cell_number == 0);
    REQUIRE(s0[5].module_number == 1);
    REQUIRE(s0[5].storage_cell_number == 1);
    REQUIRE(s0[7].module_number == 3);
    REQUIRE(s0[7].storage_cell_number == 1);

    REQUIRE(s0[0].pedestal_g0_mean == 0);
    REQUIRE(s0[0].pedestal_g1_mean == 0);
    REQUIRE(s0[0].pedestal_g2_mean == 0);

    REQUIRE(s0[0].bad_pixels == 0);
    REQUIRE(s0[1].bad_pixels == 0);
    REQUIRE(s0[2].bad_pixels == 0);
    REQUIRE(s0[3].bad_pixels == 0);

    REQUIRE(s0[4].pedestal_g0_mean == Catch::Approx(1000));
    REQUIRE(s0[4].pedestal_g1_mean == Catch::Approx(2000));
    REQUIRE(s0[4].pedestal_g2_mean == Catch::Approx(3000));

    REQUIRE(s0[5].pedestal_g0_mean == Catch::Approx(2000));
    REQUIRE(s0[5].pedestal_g1_mean == Catch::Approx(4000));
    REQUIRE(s0[5].pedestal_g2_mean == Catch::Approx(6000));

    REQUIRE(s0[6].pedestal_g0_mean == Catch::Approx(3000));
    REQUIRE(s0[6].pedestal_g1_mean == Catch::Approx(6000));
    REQUIRE(s0[6].pedestal_g2_mean == Catch::Approx(9000));

    REQUIRE(s0[4].bad_pixels == 0);
    REQUIRE(s0[5].bad_pixels == 0);
    REQUIRE(s0[6].bad_pixels == 0);
    REQUIRE(s0[7].bad_pixels == 3);
}

TEST_CASE("JFCalibration_GetPedestal","[JFCalibration]") {
    JFCalibration calibration(4, 2);

    for (int module = 0; module < 4; module++) {
        for (int i = 0; i < RAW_MODULE_SIZE; i++) {
            calibration.Pedestal(module, 1, 1).GetPedestal()[i] = 2000 * (module + 1) + (i % 30000);
        }
    }

    auto x = calibration.GetPedestal(1, 1);
    REQUIRE(x[0] == 2000);
    REQUIRE(x[2*RAW_MODULE_SIZE] == 3 * 2000);
    REQUIRE(x[3*RAW_MODULE_SIZE+500] == 4 * 2000 + 500);
}