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

#include <catch2/catch_all.hpp>

#include "../common/DiffractionExperiment.h"
#include "../writer/HDF5Writer.h"
#include "../reader/JFJochHDF5Reader.h"

TEST_CASE("HDF5DataType_Sign","[HDF5]") {
    HDF5DataType type_u8((uint8_t)0), type_fl(0.0f), type_i32((int32_t) 0), type_u32((uint32_t) 0);

    CHECK(!type_u8.IsSigned());
    CHECK(type_fl.IsSigned());
    CHECK(type_i32.IsSigned());
    CHECK(!type_u32.IsSigned());
}

TEST_CASE("HDF5DataType_ElemSize","[HDF5]") {
    HDF5DataType type_u8((uint8_t)0), type_fl(0.0f), type_i32((int32_t) 0), type_u32((uint32_t) 0);

    CHECK(type_u8.GetElemSize() == 1);
    CHECK(type_fl.GetElemSize() == 4);
    CHECK(type_i32.GetElemSize() == 4);
    CHECK(type_u32.GetElemSize() == 4);
}


TEST_CASE("HDF5DataType_ElemType","[HDF5]") {
    HDF5DataType type_u8((uint8_t)0), type_fl(0.0f), type_i32((int32_t) 0), type_u32((uint32_t) 0);

    CHECK(type_u8.IsInteger());
    CHECK(!type_fl.IsInteger());
    CHECK(type_fl.IsFloat());
    CHECK(type_i32.IsInteger());
    CHECK(type_u32.IsInteger());
}

TEST_CASE("JFJochReader_MasterFile", "[HDF5][Full]") {
    DiffractionExperiment x(DetectorGeometry(1));

    x.FilePrefix("test08").ImagesPerTrigger(950).OverwriteExistingFiles(true);
    x.BeamX_pxl(100).BeamY_pxl(200).DetectorDistance_mm(150)
            .IncidentEnergy_keV(WVL_1A_IN_KEV)
            .FrameTime(std::chrono::microseconds(500), std::chrono::microseconds(10));
    RegisterHDF5Filter();

    {
        StartMessage start_message;
        x.FillMessage(start_message);

        EndMessage end_message;
        end_message.max_image_number = 0;
        std::unique_ptr<NXmx> master = std::make_unique<NXmx>(start_message);
        master->Finalize(end_message);
        master.reset();
    }
    {
        JFJochHDF5Reader reader;
        REQUIRE_NOTHROW(reader.ReadFile("test08_master.h5"));
        auto dataset = reader.GetStartMessage();
        CHECK(dataset->geom.GetBeamX_pxl() == Catch::Approx(x.GetBeamX_pxl()));
        CHECK(dataset->geom.GetBeamY_pxl() == Catch::Approx(x.GetBeamY_pxl()));
        CHECK(dataset->geom.GetDetectorDistance_mm() == Catch::Approx(x.GetDetectorDistance_mm()));
        CHECK(dataset->frame_time == Catch::Approx(x.GetFrameTime().count() / 1e6));
        CHECK(dataset->count_time == Catch::Approx(x.GetFrameCountTime().count() / 1e6));
        CHECK(dataset->geom.GetWavelength_A() == Catch::Approx(x.GetWavelength_A()));
        CHECK(dataset->number_of_images == 0);
    }
    remove("test08_master.h5");

    REQUIRE(H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("JFJochReader_PixelMask", "[HDF5][Full]") {
    DiffractionExperiment x(DetectorGeometry(1));

    x.FilePrefix("test16").ImagesPerTrigger(950).OverwriteExistingFiles(true);
    x.BeamX_pxl(100).BeamY_pxl(200).DetectorDistance_mm(150)
            .IncidentEnergy_keV(WVL_1A_IN_KEV).PixelSigned(false).BitDepthImage(16)
            .FrameTime(std::chrono::microseconds(500), std::chrono::microseconds(10));
    RegisterHDF5Filter();

    std::vector<uint32_t> pixel_mask(x.GetPixelsNum(), 0);
    pixel_mask[5767] = 1;
    pixel_mask[x.GetPixelsNum() - 1] = 4;
    pixel_mask[0] = 256;

    std::vector<int16_t> image(x.GetPixelsNum(), 0);

    {

        StartMessage start_message;
        x.FillMessage(start_message);

        start_message.AddPixelMask(CompressedImage{
                .data = (uint8_t *) pixel_mask.data(),
                .size = pixel_mask.size() * sizeof(uint32_t),
                .xpixel = (size_t) x.GetXPixelsNum(),
                .ypixel = (size_t) x.GetYPixelsNum(),
                .pixel_depth_bytes = 4,
                .pixel_is_signed = false,
                .algorithm = CompressionAlgorithm::NO_COMPRESSION,
                .channel = "default"
        });

        HDF5Writer file_set(start_message);

        DataMessage message{};
        message.image.pixel_depth_bytes = x.GetByteDepthImage();
        message.image.pixel_is_signed = x.IsPixelSigned();
        message.image.algorithm = CompressionAlgorithm::NO_COMPRESSION;
        message.image.xpixel = x.GetXPixelsNum();
        message.image.ypixel = x.GetYPixelsNum();
        message.image.data = (uint8_t *) image.data();
        message.image.size = x.GetPixelsNum() * x.GetByteDepthImage();
        message.number = 0;

        REQUIRE_NOTHROW(file_set.Write(message));


        EndMessage end_message;
        end_message.max_image_number = 1;
        file_set.Write(end_message);

        file_set.Finalize();
    }
    {
        JFJochHDF5Reader reader;
        reader.ReadFile("test16_master.h5");
        auto dataset = reader.GetStartMessage();
        REQUIRE(dataset->pixel_mask.size() == x.GetPixelsNum());
        CHECK(dataset->pixel_mask == pixel_mask);

        REQUIRE_NOTHROW(reader.LoadImage(0));
        auto reader_image = reader.CopyImage();

        CHECK(reader_image->image.at(5767) == GAP_PXL_VALUE);
        CHECK(reader_image->image.at(0) == ERROR_PXL_VALUE);
        CHECK(reader_image->image.at(1) == 0);
        CHECK(reader_image->image.at(2) == 0);
        CHECK(reader_image->image.at(x.GetPixelsNum() - 1) == ERROR_PXL_VALUE);
    }
    remove("test16_master.h5");

    REQUIRE(H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("JFJochReader_DataI16", "[HDF5][Full]") {
    DiffractionExperiment x(DetectorGeometry(1));

    x.FilePrefix("test09").ImagesPerTrigger(4).OverwriteExistingFiles(true);
    x.BitDepthImage(16).ImagesPerFile(1).HDF5MasterFormatVersion(2).PixelSigned(true);
    x.Compression(CompressionAlgorithm::NO_COMPRESSION);

    std::vector<int16_t> image(x.GetPixelsNum());
    image[0] = INT16_MAX;
    image[1] = INT16_MIN;
    image[2] = 456;
    image[3] = -3456;

    RegisterHDF5Filter();
    {
        StartMessage start_message;
        x.FillMessage(start_message);
        HDF5Writer file_set(start_message);

        for (int i = 0; i < x.GetImageNum(); i++) {
            std::vector<SpotToSave> spots;
            image[5678] = i;

            DataMessage message{};
            message.image.pixel_depth_bytes = x.GetByteDepthImage();
            message.image.pixel_is_signed = x.IsPixelSigned();
            message.image.algorithm = CompressionAlgorithm::NO_COMPRESSION;
            message.image.xpixel = x.GetXPixelsNum();
            message.image.ypixel = x.GetYPixelsNum();
            message.image.data = (uint8_t *) image.data();
            message.image.size = x.GetPixelsNum() * x.GetByteDepthImage();
            message.spots = spots;
            message.indexing_result = (i % 2 == 0);
            message.bkg_estimate = i * 345.6;
            message.number = i;

            REQUIRE_NOTHROW(file_set.Write(message));

        }
        EndMessage end_message;
        end_message.max_image_number = x.GetImageNum();
        file_set.Write(end_message);
        file_set.Finalize();
    }
    {
        JFJochHDF5Reader reader;
        REQUIRE_NOTHROW(reader.ReadFile("test09_master.h5"));
        auto dataset = reader.GetStartMessage();
        CHECK(dataset->number_of_images == 4);

        REQUIRE_THROWS(reader.LoadImage(4));

        for (int i = 0; i < 4; i++) {
            REQUIRE_NOTHROW(reader.LoadImage(i));

            auto reader_image = reader.CopyImage();
            REQUIRE(reader_image);
            CHECK(reader_image->image[0] == SATURATED_PXL_VALUE);
            CHECK(reader_image->image[1] == ERROR_PXL_VALUE);
            CHECK(reader_image->image[2] == image[2]);
            CHECK(reader_image->image[3] == image[3]);
            CHECK(reader_image->image[5678] == i);
            CHECK(dataset->indexing_result[i] == (i % 2 == 0));
            CHECK(dataset->bkg_estimate[i] == Catch::Approx(i * 345.6));
        }
    }
    remove("test09_master.h5");
    remove("test09_data_000001.h5");
    remove("test09_data_000002.h5");
    remove("test09_data_000003.h5");
    remove("test09_data_000004.h5");

// No leftover HDF5 objects
    REQUIRE(H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}


TEST_CASE("JFJochReader_DataI16_OldMasterFormat", "[HDF5][Full]") {
    DiffractionExperiment x(DetectorGeometry(1));

    x.FilePrefix("test15").ImagesPerTrigger(4).OverwriteExistingFiles(true);
    x.BitDepthImage(16).ImagesPerFile(1).HDF5MasterFormatVersion(1).PixelSigned(true);
    x.Compression(CompressionAlgorithm::NO_COMPRESSION);

    std::vector<int16_t> image(x.GetPixelsNum());
    image[0] = INT16_MAX;
    image[1] = INT16_MIN;
    image[2] = 456;
    image[3] = -3456;

    RegisterHDF5Filter();
    {
        StartMessage start_message;
        x.FillMessage(start_message);
        HDF5Writer file_set(start_message);

        for (int i = 0; i < x.GetImageNum(); i++) {
            std::vector<SpotToSave> spots;
            image[5678] = i;

            DataMessage message{};
            message.image.pixel_depth_bytes = x.GetByteDepthImage();
            message.image.pixel_is_signed = x.IsPixelSigned();
            message.image.algorithm = CompressionAlgorithm::NO_COMPRESSION;
            message.image.xpixel = x.GetXPixelsNum();
            message.image.ypixel = x.GetYPixelsNum();
            message.image.data = (uint8_t *) image.data();
            message.image.size = x.GetPixelsNum() * x.GetByteDepthImage();
            message.spots = spots;
            message.indexing_result = (i % 2 == 0);
            message.bkg_estimate = i * 345.6;
            message.number = i;

            REQUIRE_NOTHROW(file_set.Write(message));

        }
        EndMessage end_message;
        end_message.max_image_number = x.GetImageNum();
        file_set.Write(end_message);
        file_set.Finalize();
    }
    {
        JFJochHDF5Reader reader;
        REQUIRE_NOTHROW(reader.ReadFile("test15_master.h5"));
        auto dataset = reader.GetStartMessage();
        CHECK(dataset->number_of_images == 4);
        REQUIRE(dataset->spot_count.empty());
        REQUIRE(dataset->bkg_estimate.empty());

        REQUIRE_THROWS(reader.LoadImage(4));

        for (int i = 0; i < 4; i++) {
            REQUIRE_NOTHROW(reader.LoadImage(i));

            auto reader_image = reader.CopyImage();
            REQUIRE(reader_image);
            CHECK(reader_image->image[0] == SATURATED_PXL_VALUE);
            CHECK(reader_image->image[1] == ERROR_PXL_VALUE);
            CHECK(reader_image->image[2] == image[2]);
            CHECK(reader_image->image[3] == image[3]);
            CHECK(reader_image->image[5678] == i);
        }
    }
    remove("test15_master.h5");
    remove("test15_data_000001.h5");
    remove("test15_data_000002.h5");
    remove("test15_data_000003.h5");
    remove("test15_data_000004.h5");

// No leftover HDF5 objects
    REQUIRE(H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("JFJochReader_DataU16", "[HDF5][Full]") {
    DiffractionExperiment x(DetectorGeometry(1));

    x.FilePrefix("test10").ImagesPerTrigger(4).OverwriteExistingFiles(true);
    x.BitDepthImage(16).ImagesPerFile(1).HDF5MasterFormatVersion(2).PixelSigned(false);
    x.Compression(CompressionAlgorithm::NO_COMPRESSION);

    std::vector<uint16_t> image(x.GetPixelsNum());
    image[0] = UINT16_MAX;
    image[1] = INT16_MAX;
    image[2] = 456;

    RegisterHDF5Filter();
    {
        StartMessage start_message;
        x.FillMessage(start_message);
        HDF5Writer file_set(start_message);

        for (int i = 0; i < x.GetImageNum(); i++) {
            std::vector<SpotToSave> spots;
            image[5678] = i;

            DataMessage message{};
            message.image.pixel_depth_bytes = x.GetByteDepthImage();
            message.image.pixel_is_signed = x.IsPixelSigned();
            message.image.algorithm = CompressionAlgorithm::NO_COMPRESSION;
            message.image.xpixel = x.GetXPixelsNum();
            message.image.ypixel = x.GetYPixelsNum();
            message.image.data = (uint8_t *) image.data();
            message.image.size = x.GetPixelsNum() * x.GetByteDepthImage();
            message.spots = spots;
            message.indexing_result = (i % 2 == 0);
            message.bkg_estimate = i * 345.6;
            message.number = i;

            REQUIRE_NOTHROW(file_set.Write(message));

        }
        EndMessage end_message;
        end_message.max_image_number = x.GetImageNum();
        file_set.Write(end_message);
        file_set.Finalize();
    }
    {
        JFJochHDF5Reader reader;
        REQUIRE_NOTHROW(reader.ReadFile("test10_master.h5"));
        auto dataset = reader.GetStartMessage();
        CHECK(dataset->number_of_images == 4);

        REQUIRE_THROWS(reader.LoadImage(4));

        for (int i = 0; i < 4; i++) {
            REQUIRE_NOTHROW(reader.LoadImage(i));

            auto reader_image = reader.CopyImage();
            REQUIRE(reader_image);
            CHECK(reader_image->image[0] == SATURATED_PXL_VALUE);
            CHECK(reader_image->image[1] == INT16_MAX);
            CHECK(reader_image->image[2] == 456);
            CHECK(reader_image->image[5678] == i);
            CHECK(dataset->indexing_result[i] == (i % 2 == 0));
            CHECK(dataset->bkg_estimate[i] == Catch::Approx(i * 345.6));
        }
    }
    remove("test10_master.h5");
    remove("test10_data_000001.h5");
    remove("test10_data_000002.h5");
    remove("test10_data_000003.h5");
    remove("test10_data_000004.h5");

// No leftover HDF5 objects
    REQUIRE(H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("JFJochReader_DataI32", "[HDF5][Full]") {
    DiffractionExperiment x(DetectorGeometry(1));

    x.FilePrefix("test11").ImagesPerTrigger(4).OverwriteExistingFiles(true);
    x.BitDepthImage(32).ImagesPerFile(1).HDF5MasterFormatVersion(2).PixelSigned(true);
    x.Compression(CompressionAlgorithm::NO_COMPRESSION);

    std::vector<uint32_t> image(x.GetPixelsNum());
    image[0] = INT32_MAX;
    image[1] = INT32_MIN;
    image[2] = 456;

    RegisterHDF5Filter();
    {
        StartMessage start_message;
        x.FillMessage(start_message);
        HDF5Writer file_set(start_message);

        for (int i = 0; i < x.GetImageNum(); i++) {
            std::vector<SpotToSave> spots;
            image[5678] = i;

            DataMessage message{};
            message.image.pixel_depth_bytes = x.GetByteDepthImage();
            message.image.pixel_is_signed = x.IsPixelSigned();
            message.image.algorithm = CompressionAlgorithm::NO_COMPRESSION;
            message.image.xpixel = x.GetXPixelsNum();
            message.image.ypixel = x.GetYPixelsNum();
            message.image.data = (uint8_t *) image.data();
            message.image.size = x.GetPixelsNum() * x.GetByteDepthImage();
            message.spots = spots;
            message.number = i;

            REQUIRE_NOTHROW(file_set.Write(message));

        }
        EndMessage end_message;
        end_message.max_image_number = x.GetImageNum();
        file_set.Write(end_message);
        file_set.Finalize();
    }
    {
        JFJochHDF5Reader reader;
        REQUIRE_NOTHROW(reader.ReadFile("test11_master.h5"));
        auto dataset = reader.GetStartMessage();
        CHECK(dataset->number_of_images == 4);

        REQUIRE_THROWS(reader.LoadImage(4));

        for (int i = 0; i < 4; i++) {
            REQUIRE_NOTHROW(reader.LoadImage(i));

            auto reader_image = reader.CopyImage();
            REQUIRE(reader_image);
            CHECK(reader_image->image[0] == SATURATED_PXL_VALUE);
            CHECK(reader_image->image[1] == ERROR_PXL_VALUE);
            CHECK(reader_image->image[2] == 456);
            CHECK(reader_image->image[5678] == i);
        }
    }
    remove("test11_master.h5");
    remove("test11_data_000001.h5");
    remove("test11_data_000002.h5");
    remove("test11_data_000003.h5");
    remove("test11_data_000004.h5");

// No leftover HDF5 objects
    REQUIRE(H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("JFJochReader_DataU32", "[HDF5][Full]") {
    DiffractionExperiment x(DetectorGeometry(1));

    x.FilePrefix("test12").ImagesPerTrigger(4).OverwriteExistingFiles(true);
    x.BitDepthImage(32).ImagesPerFile(1).HDF5MasterFormatVersion(2).PixelSigned(false);
    x.Compression(CompressionAlgorithm::NO_COMPRESSION);

    std::vector<uint32_t> image(x.GetPixelsNum());
    image[0] = UINT32_MAX;
    image[1] = static_cast<uint32_t>(INT32_MAX) + 50;
    image[2] = 456;
    image[3] = INT32_MAX;
    image[4] = INT32_MAX - 1;

    RegisterHDF5Filter();
    {
        StartMessage start_message;
        x.FillMessage(start_message);
        HDF5Writer file_set(start_message);

        for (int i = 0; i < x.GetImageNum(); i++) {
            std::vector<SpotToSave> spots;
            image[5678] = i;

            DataMessage message{};
            message.image.pixel_depth_bytes = x.GetByteDepthImage();
            message.image.pixel_is_signed = x.IsPixelSigned();
            message.image.algorithm = CompressionAlgorithm::NO_COMPRESSION;
            message.image.xpixel = x.GetXPixelsNum();
            message.image.ypixel = x.GetYPixelsNum();
            message.image.data = (uint8_t *) image.data();
            message.image.size = x.GetPixelsNum() * x.GetByteDepthImage();
            message.spots = spots;
            message.number = i;

            REQUIRE_NOTHROW(file_set.Write(message));

        }
        EndMessage end_message;
        end_message.max_image_number = x.GetImageNum();
        file_set.Write(end_message);
        file_set.Finalize();
    }
    {
        JFJochHDF5Reader reader;
        REQUIRE_NOTHROW(reader.ReadFile("test12_master.h5"));
        auto dataset = reader.GetStartMessage();
        CHECK(dataset->number_of_images == 4);

        REQUIRE_THROWS(reader.LoadImage(4));

        for (int i = 0; i < 4; i++) {
            REQUIRE_NOTHROW(reader.LoadImage(i));

            auto reader_image = reader.CopyImage();
            REQUIRE(reader_image);
            CHECK(reader_image->image[0] == INT32_MAX);
            CHECK(reader_image->image[1] == INT32_MAX);
            CHECK(reader_image->image[2] == 456);
            CHECK(reader_image->image[3] == INT32_MAX);
            CHECK(reader_image->image[4] == INT32_MAX - 1);
            CHECK(reader_image->image[5678] == i);
        }
    }
    remove("test12_master.h5");
    remove("test12_data_000001.h5");
    remove("test12_data_000002.h5");
    remove("test12_data_000003.h5");
    remove("test12_data_000004.h5");

// No leftover HDF5 objects
    REQUIRE(H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("JFJochReader_ROI", "[HDF5][Full]") {
    DiffractionExperiment x(DetectorGeometry(1));

    x.FilePrefix("test25").ImagesPerTrigger(4).OverwriteExistingFiles(true);
    x.BitDepthImage(16).ImagesPerFile(1).HDF5MasterFormatVersion(2).PixelSigned(false);
    x.Compression(CompressionAlgorithm::NO_COMPRESSION);

    x.ROI().SetROI(ROIDefinition{
            .boxes = {ROIBox("beam", 100, 120, 20, 30)},
            .circles = {ROICircle("roi1", 500, 800, 10)}
    });

    std::vector<uint16_t> image(x.GetPixelsNum());

    RegisterHDF5Filter();
    {
        StartMessage start_message;
        x.FillMessage(start_message);
        HDF5Writer file_set(start_message);

        for (int i = 0; i < x.GetImageNum(); i++) {
            std::vector<SpotToSave> spots;

            DataMessage message{};
            message.image.pixel_depth_bytes = x.GetByteDepthImage();
            message.image.pixel_is_signed = x.IsPixelSigned();
            message.image.algorithm = CompressionAlgorithm::NO_COMPRESSION;
            message.image.xpixel = x.GetXPixelsNum();
            message.image.ypixel = x.GetYPixelsNum();
            message.image.data = (uint8_t *) image.data();
            message.image.size = x.GetPixelsNum() * x.GetByteDepthImage();
            message.spots = spots;
            message.number = i;
            message.roi["beam"] = ROIMessage{
                    .sum = 12 + i,
                    .sum_square = (uint64_t) 123 * i,
                    .max_count = 123 - i,
                    .pixels = (uint64_t) 189 + i
            };

            message.roi["roi1"] = ROIMessage{
                    .sum = 25 + i,
                    .sum_square = (uint64_t) 15 * i,
                    .max_count = 67 - i,
                    .pixels = (uint64_t) 95 + i
            };
            REQUIRE_NOTHROW(file_set.Write(message));

        }
        EndMessage end_message;
        end_message.max_image_number = x.GetImageNum();
        file_set.Write(end_message);
        file_set.Finalize();
    }
    {
        JFJochHDF5Reader reader;
        REQUIRE_NOTHROW(reader.ReadFile("test25_master.h5"));
        auto dataset = reader.GetStartMessage();
        CHECK(dataset->number_of_images == 4);

        CHECK(dataset->roi.size() == 2);
        REQUIRE(dataset->roi_max.size() == 2);
        int index = 0;
        if (dataset->roi[1] == "beam")
            index = 1;

        REQUIRE(dataset->roi_max[index].size() == 4);
        CHECK(dataset->roi_max[index][2] == 123 - 2);
        CHECK(dataset->roi_sum_sq[index][0] == 0);
        CHECK(dataset->roi_sum_sq[index][1] == 123);
        CHECK(dataset->roi_npixel[index][3] == 189 + 3);
        CHECK(dataset->roi_sum[1 - index][3] == 25 + 3);
    }
    remove("test25_master.h5");
    remove("test25_data_000001.h5");
    remove("test25_data_000002.h5");
    remove("test25_data_000003.h5");
    remove("test25_data_000004.h5");
// No leftover HDF5 objects
    REQUIRE(H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}