// 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 "../common/PixelMask.h"

TEST_CASE("PixelMask_MaskModuleEdges","[PixelMask]") {
    DiffractionExperiment experiment(DetJF(4, 1));
    experiment.MaskModuleEdges(true).MaskChipEdges(false).Raw();

    PixelMask mask(experiment);

    auto mask_out = mask.GetMaskRaw(experiment);

    REQUIRE(mask_out.size() == experiment.GetModulesNum() * RAW_MODULE_SIZE);

    CHECK(mask_out[0] == (1u<<30));
    CHECK(mask_out[1022] == (1u<<30));
    CHECK(mask_out[1024*5] == (1u<<30));
    CHECK(mask_out[1024*5+1023] == (1u<<30));
    CHECK(mask_out[1024*512-1] == (1u<<30));
    CHECK(mask_out[1024*512] == (1u<<30));
    CHECK(mask_out[1024*512*3+123] == (1u<<30));
    CHECK(mask_out[1024*512*3+1024] == (1u<<30));
    CHECK(mask_out[1024*800+256] == 0);

    CHECK(mask.GetMask()[0] == (1u<<30));
    CHECK(mask.GetMask()[1022] == (1u<<30));
}

TEST_CASE("PixelMask_MaskChipEdges","[PixelMask]") {
    DiffractionExperiment experiment(DetJF(4,1));
    experiment.MaskChipEdges(true).MaskModuleEdges(false).Raw();

    PixelMask mask(experiment);
    auto mask_out = mask.GetMaskRaw(experiment);
    REQUIRE(mask_out.size() == experiment.GetModulesNum() * RAW_MODULE_SIZE);

    CHECK(mask_out[255] == (1u<<31));
    CHECK(mask_out[1024*800+256] == (1u<<31));

    CHECK(mask_out[1024*233+511] == (1u<<31));
    CHECK(mask_out[1024*512+512] == (1u<<31));
    CHECK(mask_out[1024*1000+767] == (1u<<31));
    CHECK(mask_out[1024*(512*3+12)+768] == (1u<<31));
    CHECK(mask_out[1024*(512+255)+345] == (1u<<31));
    CHECK(mask_out[1024*(1024+256)+876] == (1u<<31));

    CHECK(mask_out[1022] == 0);
    CHECK(mask_out[1024*5+1023] == 0);
    CHECK(mask_out[1024*18] == 0);
}

TEST_CASE("PixelMask_MaskWrongGain","[PixelMask]") {
    ImageFormatSettings image_format;
    image_format.MaskModuleEdges(false).MaskChipEdges(false).MaskPixelsWithoutG0(true).GeometryTransformed(false);

    DiffractionExperiment experiment(DetJF(4,1));
    experiment.ImportImageFormatSettings(image_format);

    PixelMask mask(experiment);

    JFCalibration calibration(experiment);
    for (int m = 0; m < experiment.GetModulesNum(); m++) {
        calibration.Pedestal(m, 0) = JFModulePedestal(1500);
        calibration.Pedestal(m, 1) = JFModulePedestal(12000);
        calibration.Pedestal(m, 2) = JFModulePedestal(13000);
    }

    calibration.Pedestal(2,0).GetPedestal()[345] = UINT16_MAX;
    calibration.Pedestal(2,1).GetPedestal()[346] = UINT16_MAX;
    calibration.Pedestal(2,2).GetPedestal()[347] = UINT16_MAX;

    mask.LoadDetectorBadPixelMask(experiment, &calibration);

    auto mask_out = mask.GetMaskRaw(experiment);
    REQUIRE(mask_out.size() == experiment.GetModulesNum() * RAW_MODULE_SIZE);

    CHECK(mask_out[2*RAW_MODULE_SIZE + 345] != 0);
    CHECK(mask_out[2*RAW_MODULE_SIZE + 346] != 0);
    CHECK(mask_out[2*RAW_MODULE_SIZE + 347] != 0);

    CHECK(mask.GetStatistics().noisy_pixel == 0);
    CHECK(mask.GetStatistics().error_pixel == 3);
    CHECK(mask.GetStatistics().user_mask == 0);

    image_format.MaskPixelsWithoutG0(false);
    experiment.ImportImageFormatSettings(image_format);

    REQUIRE(!experiment.IsMaskPixelsWithoutG0());
    mask.LoadDetectorBadPixelMask(experiment, &calibration);
    mask_out = mask.GetMaskRaw(experiment);

    CHECK(mask_out[2*RAW_MODULE_SIZE + 345] == 0);
    CHECK(mask_out[2*RAW_MODULE_SIZE + 346] != 0);
    CHECK(mask_out[2*RAW_MODULE_SIZE + 347] != 0);

    CHECK(mask.GetStatistics().noisy_pixel == 0);
    CHECK(mask.GetStatistics().error_pixel == 2);
    CHECK(mask.GetStatistics().user_mask == 0);
}

TEST_CASE("PixelMask_MaskG0RMS","[PixelMask]") {
    ImageFormatSettings image_format;
    image_format.MaskModuleEdges(false).MaskChipEdges(false).MaskPixelsWithoutG0(true).GeometryTransformed(false);
    image_format.PedestalG0RMSLimit(100);

    DiffractionExperiment experiment(DetJF(4,1));
    experiment.ImportImageFormatSettings(image_format);

    PixelMask mask(experiment);

    JFCalibration calibration(experiment);
    for (int m = 0; m < experiment.GetModulesNum(); m++) {
        calibration.Pedestal(m, 0) = JFModulePedestal(1500);
        calibration.Pedestal(m, 1) = JFModulePedestal(12000);
        calibration.Pedestal(m, 2) = JFModulePedestal(13000);
    }

    calibration.Pedestal(2,0).GetPedestal()[345] = UINT16_MAX;
    calibration.Pedestal(2,1).GetPedestal()[346] = UINT16_MAX;
    calibration.Pedestal(2,2).GetPedestal()[347] = UINT16_MAX;
    calibration.Pedestal(2, 0).GetPedestalRMS()[245] = 150;
    calibration.Pedestal(2, 0).GetPedestalRMS()[345] = UINT16_MAX;
    calibration.Pedestal(2, 0).GetPedestalRMS()[346] = UINT16_MAX;

    mask.LoadDetectorBadPixelMask(experiment, &calibration);

    auto mask_out = mask.GetMaskRaw(experiment);
    REQUIRE(mask_out.size() == experiment.GetModulesNum() * RAW_MODULE_SIZE);

    CHECK(mask_out[2*RAW_MODULE_SIZE + 245] == (1 << PixelMask::NoisyPixelBit));

    CHECK(mask_out[2*RAW_MODULE_SIZE + 345] == (1 << PixelMask::ErrorPixelBit));
    CHECK(mask_out[2*RAW_MODULE_SIZE + 346] == ((1 << PixelMask::ErrorPixelBit) | (1 << PixelMask::NoisyPixelBit)));
    CHECK(mask_out[2*RAW_MODULE_SIZE + 347] == (1 << PixelMask::ErrorPixelBit));

    CHECK(mask.GetStatistics().noisy_pixel == 2);
    CHECK(mask.GetStatistics().error_pixel == 3);
    CHECK(mask.GetStatistics().user_mask == 0);
}


TEST_CASE("PixelMask_SCs","[PixelMask]") {
    ImageFormatSettings image_format;
    image_format.MaskModuleEdges(false).MaskChipEdges(false).MaskPixelsWithoutG0(true).GeometryTransformed(false);
    image_format.PedestalG0RMSLimit(100);

    DiffractionExperiment experiment(DetJF(4,1));
    experiment.ImportImageFormatSettings(image_format);
    experiment.StorageCells(16);

    PixelMask mask(experiment);

    JFCalibration calibration(experiment);
    calibration.Pedestal(0,0, 12).GetPedestal()[456] = UINT16_MAX;
    calibration.Pedestal(1,2, 15).GetPedestal()[324] = UINT16_MAX;
    calibration.Pedestal(2,0, 1).GetPedestalRMS()[324] = 500;

    calibration.Pedestal(3,0, 5).GetPedestalRMS()[0] = 500;
    calibration.Pedestal(3,2, 3).GetPedestal()[0] = UINT16_MAX;
    calibration.Pedestal(3,2, 11).GetPedestal()[0] = UINT16_MAX;
    calibration.Pedestal(3,2, 14).GetPedestal()[0] = UINT16_MAX;

    mask.LoadDetectorBadPixelMask(experiment, &calibration);

    auto mask_out = mask.GetMaskRaw(experiment);
    REQUIRE(mask_out.size() == experiment.GetModulesNum() * RAW_MODULE_SIZE);

    CHECK(mask_out[456] == (1 << PixelMask::ErrorPixelBit));
    CHECK(mask_out[RAW_MODULE_SIZE * 1 + 324] == (1 << PixelMask::ErrorPixelBit));
    CHECK(mask_out[RAW_MODULE_SIZE * 2 + 324] == (1 << PixelMask::NoisyPixelBit));
    CHECK(mask_out[RAW_MODULE_SIZE * 3 + 0] == ((1 << PixelMask::ErrorPixelBit) | (1 << PixelMask::NoisyPixelBit)));

    CHECK(mask.GetStatistics().noisy_pixel == 2);
    CHECK(mask.GetStatistics().error_pixel == 3);
}

TEST_CASE("PixelMask_CalculateNexusMask_UserMaskConv","[PixelMask]") {
    DiffractionExperiment experiment(DetJF(4, 1, 8, 36, false));
    experiment.MaskModuleEdges(false).MaskChipEdges(false);

    PixelMask mask(experiment);

    std::vector<uint32_t> v(experiment.GetPixelsNum(), 0);
    v[1030 * 700 + 300] = 1;

    REQUIRE_NOTHROW(mask.LoadUserMask(experiment, v));

    auto mask_v = mask.GetMask(experiment);

    REQUIRE(mask_v.size() == experiment.GetPixelsNum() );
    REQUIRE(mask_v[1030 * 700 + 300] == (1 << PixelMask::UserMaskedPixelBit));
    REQUIRE(mask_v[(1030+8)*514] == (1 << PixelMask::ModuleGapPixelBit));

    auto user_mask_v = mask.GetUserMask(experiment);

    REQUIRE(user_mask_v.size() == experiment.GetPixelsNum() );
    REQUIRE(user_mask_v[1030 * 700 + 300] == 1);

    REQUIRE(user_mask_v[(1030+8)*514] == 0);
}

TEST_CASE("PixelMask_MaskDetectorGaps","[PixelMask]") {
    DiffractionExperiment experiment(DetJF(8, 2, 8, 36, true));
    experiment.MaskChipEdges(false).MaskModuleEdges(false);
    experiment.IncidentEnergy_keV(WVL_1A_IN_KEV);

    PixelMask mask(experiment);
    auto mask_export = mask.GetMask();
    REQUIRE(mask_export.size() == experiment.GetPixelsNum());

    REQUIRE(mask_export[0] == 0);
    REQUIRE(mask_export[1029] == 0);
    REQUIRE(mask_export[1030] == 1);
    REQUIRE(mask_export[1031] == 1);
    REQUIRE(mask_export[1037] == 1);
    REQUIRE(mask_export[1038] == 0);

    REQUIRE(mask_export[(1030*2+8)*514] == 1);
    REQUIRE(mask_export[(1030*2+8)*514 + 566] == 1);
    REQUIRE(mask_export[(1030*2+8)*(514*3+36*2+12) + 566] == 1);
}

TEST_CASE("PixelMask_LoadDECTRISBadPixelMask","[PixelMask]") {
    DiffractionExperiment experiment(DetJF(4,1));
    PixelMask mask(experiment);

    std::vector<uint32_t> det_mask(experiment.GetPixelsNum(), 0);
    det_mask[0] = (1 << PixelMask::UserMaskedPixelBit);
    det_mask[534] = (1 << PixelMask::UserMaskedPixelBit) | (1 << PixelMask::ModuleGapPixelBit);
    det_mask[535] = (1 << PixelMask::ModuleGapPixelBit);
    det_mask[2 * 1000 * 5] = (1 << 3);
    det_mask[2 * 1000 * 5 + 3] = (1 << 2);
    det_mask[2 * 1000 * 5 + 4] = (1 << PixelMask::NoisyPixelBit);
    det_mask[23 * 1000 * 5] = (1 << 5) | (1 << PixelMask::UserMaskedPixelBit);

    mask.LoadDECTRISBadPixelMask(det_mask);

    CHECK(mask.GetStatistics().user_mask == 2);
    CHECK(mask.GetStatistics().error_pixel == 4);
    CHECK(mask.GetStatistics().noisy_pixel == 0);

    CHECK(mask.GetMask()[0] == (1 << PixelMask::UserMaskedPixelBit));
    CHECK(mask.GetMask()[534] == (1 << PixelMask::ModuleGapPixelBit));
    CHECK(mask.GetMask()[535] == (1 << PixelMask::ModuleGapPixelBit));
    CHECK(mask.GetMask()[2 * 1000 * 5] == (1 << PixelMask::ErrorPixelBit));
    CHECK(mask.GetMask()[2 * 1000 * 5 + 3] == (1 << PixelMask::ErrorPixelBit));
    CHECK(mask.GetMask()[2 * 1000 * 5 + 4] == (1 << PixelMask::ErrorPixelBit));
    CHECK(mask.GetMask()[23 * 1000 * 5] == ((1 << PixelMask::ErrorPixelBit) | (1 << PixelMask::UserMaskedPixelBit)));
}


TEST_CASE("PixelMask_LoadDarkBadPixelMask","[PixelMask]") {
    DiffractionExperiment experiment(DetJF(4,1));
    PixelMask mask(experiment);

    std::vector<uint32_t> det_mask(experiment.GetPixelsNum(), 0);
    det_mask[0] = (1 << PixelMask::UserMaskedPixelBit);
    det_mask[534] = (1 << PixelMask::UserMaskedPixelBit) | (1 << PixelMask::ModuleGapPixelBit);
    det_mask[535] = (1 << PixelMask::ModuleGapPixelBit);
    det_mask[2 * 1000 * 5] = (1 << 3);
    det_mask[2 * 1000 * 5 + 3] = (1 << 2);
    det_mask[2 * 1000 * 5 + 4] = (1 << PixelMask::NoisyPixelBit);
    det_mask[23 * 1000 * 5] = (1 << 5) | (1 << PixelMask::UserMaskedPixelBit);

    mask.LoadDECTRISBadPixelMask(det_mask);

    std::vector<uint32_t> dark_mask(experiment.GetPixelsNum(), 0);
    dark_mask[0] = 1; // user masked, should be included
    dark_mask[534] = 1; // On gap, should be ignored
    dark_mask[2 * 1000 * 5] = 1; // bad pixel, should be ignored
    dark_mask[67867] = 1; // Nothing else, should be included
    mask.LoadDarkBadPixelMask(dark_mask);

    CHECK(mask.GetStatistics().user_mask == 2);
    CHECK(mask.GetStatistics().error_pixel == 4);
    CHECK(mask.GetStatistics().noisy_pixel == 2);

    CHECK(mask.GetMask()[0] == (1 << PixelMask::UserMaskedPixelBit | 1 << PixelMask::NoisyPixelBit));

    CHECK(mask.GetMask()[534] == (1 << PixelMask::ModuleGapPixelBit));
    CHECK(mask.GetMask()[2 * 1000 * 5] == (1 << PixelMask::ErrorPixelBit));
    CHECK(mask.GetMask()[67867] == (1 << PixelMask::NoisyPixelBit));
}