// Copyright (2019-2023) Paul Scherrer Institute

#include <catch2/catch.hpp>
#include <iostream>

#include "../common/DiffractionExperiment.h"
#include "../writer/HDF5Objects.h"
#include "../writer/HDF5Writer.h"
#include "../writer/HDF5NXmx.h"
#include "../compression/JFJochCompressor.h"
#include "../image_analysis/RadialIntegrationProfile.h"

using namespace std::literals::chrono_literals;

TEST_CASE("HDF5DataSet_scalar", "[HDF5][Unit]") {
    uint16_t tmp_scalar = 16788;
    {
        HDF5File file("scratch1.h5", true, true, false);
        file.SaveScalar("scalar", tmp_scalar);
    }
    {
        HDF5ReadOnlyFile file("scratch1.h5");
        HDF5DataSet scalar_dataset(file, "scalar");
        HDF5DataType data_type_scalar(scalar_dataset);
        HDF5DataSpace data_space_scalar(scalar_dataset);
        REQUIRE(data_type_scalar.GetElemSize() == 2);
        REQUIRE(data_space_scalar.GetNumOfDimensions() == 0);
        REQUIRE(scalar_dataset.ReadScalar<uint16_t>() == tmp_scalar);
        REQUIRE(scalar_dataset.ReadScalar<int64_t>() == tmp_scalar);
        std::vector<uint16_t> v;
        REQUIRE_THROWS(scalar_dataset.ReadVector(v));
    }
    remove("scratch1.h5");
    REQUIRE (H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("HDF5DataSet_string", "[HDF5][Unit]") {
    std::string tmp_string = "HDF5Content";

    {
        HDF5File file("scratch2.h5", true, true, false);
        file.SaveScalar("str", tmp_string);
    }
    {
        HDF5ReadOnlyFile file("scratch2.h5");
        HDF5DataSet string_dataset(file, "str");
        HDF5DataType data_type_str(string_dataset);
        HDF5DataSpace data_space_str(string_dataset);
        CHECK(data_type_str.GetElemSize() == tmp_string.size() + 1);
        CHECK(data_space_str.GetNumOfDimensions() == 0);
        CHECK(string_dataset.ReadString() == tmp_string);
    }
    remove("scratch2.h5");
    REQUIRE (H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("HDF5DataSet_vector", "[HDF5][Unit]") {
    std::vector<double> tmp_vector (16384);
    tmp_vector[0] = 599.88;
    tmp_vector[1000] = 800.12;
    tmp_vector[15000] = 3.1415926;

    {
        RegisterHDF5Filter();
        HDF5File file("scratch3.h5", true, true, false);
        file.SaveVector("vec", tmp_vector);
    }

    {
        HDF5ReadOnlyFile file("scratch3.h5");
        HDF5DataSet vector_dataset(file, "vec");
        HDF5DataType data_type_vec(vector_dataset);
        HDF5DataSpace data_space_vec(vector_dataset);
        CHECK(data_type_vec.GetElemSize() == 8);
        CHECK(data_space_vec.GetNumOfDimensions() == 1);
        CHECK(data_space_vec.GetDimensions()[0] == tmp_vector.size());
        std::vector<double> output;
        REQUIRE_NOTHROW(vector_dataset.ReadVector(output));
        CHECK (output[0] == tmp_vector[0]);
        CHECK (output[1000] == tmp_vector[1000]);
        CHECK (output[15000] == tmp_vector[15000]);
        CHECK (output.size() == tmp_vector.size());

        std::vector<double> output2(10);
        REQUIRE_NOTHROW(vector_dataset.ReadVector(output2, {15000},{10}));
        REQUIRE_THROWS(vector_dataset.ReadVector(output2, {15000},{100}));
        CHECK(output2[0] == tmp_vector[15000]);
    }

    remove("scratch3.h5");
    REQUIRE (H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("HDF5File_Delete", "[HDF5][Unit]") {
    uint16_t tmp_scalar = 16788;
    {
        HDF5File file("scratch1.h5", true, true, false);
        file.SaveScalar("scalar", tmp_scalar);
    }
    {
        HDF5File file("scratch1.h5", false, true, false);
        std::unique_ptr<HDF5DataSet> scalar_dataset;
        REQUIRE_NOTHROW(scalar_dataset = std::make_unique<HDF5DataSet>(file, "scalar"));
        file.Delete("/scalar");

    }
    {
        HDF5File file("scratch1.h5", false, true, false);
        std::unique_ptr<HDF5DataSet> scalar_dataset;
        REQUIRE_THROWS(scalar_dataset = std::make_unique<HDF5DataSet>(file, "scalar"));
    }
    remove("scratch1.h5");
    REQUIRE (H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("HDF5MasterFile", "[HDF5][Full]") {
    {
        RegisterHDF5Filter();
        DiffractionExperiment x(DetectorGeometry(8, 2, 8, 36));

        x.FilePrefix("test01").ImagesPerTrigger(950);

        StartMessage start_message;
        x.FillMessage(start_message);
        EndMessage end_message;
        end_message.number_of_images = x.GetImageNum();
        REQUIRE_NOTHROW(HDF5Metadata::NXmx(start_message, end_message));

        x.FilePrefix("test02");
        x.FillMessage(start_message);
        REQUIRE_NOTHROW(HDF5Metadata::NXmx(start_message, end_message));
    }
    remove("test01_master.h5");
    remove("test02_master.h5");
    // No leftover HDF5 objects
    REQUIRE (H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("HDF5MasterFile_RadInt", "[HDF5][Full]") {
    {
        RegisterHDF5Filter();
        DiffractionExperiment x(DetectorGeometry(8, 2, 8, 36));

        x.DetectorDistance_mm(50).BeamX_pxl(1000).BeamY_pxl(1000);
        x.QSpacingForRadialInt_recipA(0.1).LowQForRadialInt_recipA(0.1).HighQForRadialInt_recipA(4);

        x.FilePrefix("test01_rad_int").ImagesPerTrigger(950);

        RadialIntegrationMapping mapping(x);
        RadialIntegrationProfile profile(mapping, x);

        StartMessage start_message;
        x.FillMessage(start_message);
        start_message.rad_int_bin_to_q = mapping.GetBinToQ();
        EndMessage end_message;
        end_message.number_of_images = x.GetImageNum();
        end_message.rad_int_result["avg1"] = profile.GetResult();
        end_message.rad_int_result["avg2"] = profile.GetResult();

        REQUIRE_NOTHROW(HDF5Metadata::NXmx(start_message, end_message));
    }
    remove("test01_rad_int_master.h5");
    // No leftover HDF5 objects
    REQUIRE (H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("HDF5Writer", "[HDF5][Full]") {
    {
        RegisterHDF5Filter();
        DiffractionExperiment x(DetectorGeometry(8, 2, 8, 36));
        std::vector<SpotToSave> spots;

        x.FilePrefix("test02_1p10").ImagesPerTrigger(5).DataFileCount(2).Compression(JFJochProtoBuf::NO_COMPRESSION);
        StartMessage start_message;
        x.FillMessage(start_message);

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

        for (int i = 0; i < x.GetImageNum(); i++) {
            DataMessage message{};
            message.image.data = (uint8_t *) image.data();
            message.image.size = x.GetPixelsNum() * x.GetPixelDepth();
            message.spots = spots;
            message.number = i;

            REQUIRE_NOTHROW(file_set.Write(message));
        }
    }
    // No leftover HDF5 objects
    REQUIRE (H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("HDF5Writer_Spots", "[HDF5][Full]") {
    {
        RegisterHDF5Filter();
        DiffractionExperiment x(DetectorGeometry(8, 2, 8, 36));
        std::vector<SpotToSave> spots;

        spots.push_back({10,10,7});
        spots.push_back({20,50,12});
        spots.push_back({1000,500,3});

        x.FilePrefix("test02_1p10_spots").ImagesPerTrigger(5).DataFileCount(2).Compression(JFJochProtoBuf::NO_COMPRESSION);
        StartMessage start_message;
        x.FillMessage(start_message);

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

        for (int i = 0; i < x.GetImageNum(); i++) {
            DataMessage message{};
            message.image.data = (uint8_t *) image.data();
            message.image.size = x.GetPixelsNum() * x.GetPixelDepth();
            message.spots = spots;
            message.number = i;

            REQUIRE_NOTHROW(file_set.Write(message));
        }
    }
    // No leftover HDF5 objects
    REQUIRE (H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("HDF5Writer_Rad_Int_Profile", "[HDF5][Full]") {
    {
        RegisterHDF5Filter();
        DiffractionExperiment x(DetectorGeometry(8, 2, 8, 36));
        x.DetectorDistance_mm(50).BeamX_pxl(1000).BeamY_pxl(1000);
        x.QSpacingForRadialInt_recipA(0.1).LowQForRadialInt_recipA(0.1).HighQForRadialInt_recipA(4);

        RadialIntegrationMapping mapping(x);
        RadialIntegrationProfile profile(mapping, x);

        std::vector<float> rad_int_profile(mapping.GetBinNumber(), 4.0);
        std::vector<float> rad_int_avg(mapping.GetBinNumber(), 0.33);

        x.FilePrefix("test02_1p10_rad_int").ImagesPerTrigger(5).DataFileCount(2).Compression(JFJochProtoBuf::NO_COMPRESSION);
        StartMessage start_message;
        x.FillMessage(start_message);
        start_message.rad_int_bin_to_q = mapping.GetBinToQ();

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

        for (int i = 0; i < x.GetImageNum(); i++) {
            DataMessage message{};
            message.image.data = (uint8_t *) image.data();
            message.image.size = x.GetPixelsNum() * x.GetPixelDepth();
            message.rad_int_profile = std::vector<float>(mapping.GetBinNumber(), i);
            message.number = i;

            REQUIRE_NOTHROW(file_set.Write(message));
        }
    }
    // No leftover HDF5 objects
    REQUIRE (H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

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

    x.ImagesPerTrigger(5).DataFileCount(5)
     .Compression(JFJochProtoBuf::NO_COMPRESSION).FilePrefix("vds");

    {
        RegisterHDF5Filter();

        StartMessage start_message;
        x.FillMessage(start_message);

        EndMessage end_message;
        end_message.number_of_images = x.GetImageNum();

        HDF5Writer writer(start_message);
        std::vector<uint16_t> image(x.GetPixelsNum());
        std::vector<SpotToSave> spots;
        for (int i = 0; i < x.GetImageNum(); i++) {
            for (auto &j: image)
                j = i;
            DataMessage message{};
            message.image.data = (uint8_t *) image.data();
            message.image.size = x.GetPixelsNum() * x.GetPixelDepth();
            message.spots = spots;
            message.number = i;

            REQUIRE_NOTHROW(writer.Write(message));
        }

        REQUIRE_NOTHROW(HDF5Metadata::NXmx(start_message, end_message));
    }
    {
        HDF5File file("vds_master.h5", false, true, false);
        std::unique_ptr<HDF5DataSet> dataset;
        REQUIRE_NOTHROW(dataset = std::make_unique<HDF5DataSet>(file,"/entry/data/data"));
        HDF5DataSpace file_space(*dataset);
        REQUIRE(file_space.GetNumOfDimensions() == 3);
        REQUIRE(file_space.GetDimensions()[0] == x.GetImageNum());
        REQUIRE(file_space.GetDimensions()[1] == x.GetYPixelsNum());
        REQUIRE(file_space.GetDimensions()[2] == x.GetXPixelsNum());

        std::vector<int16_t> v(x.GetImageNum() * x.GetYPixelsNum() * x.GetXPixelsNum());
        std::vector<hsize_t> start = {0, 0, 0};
        dataset->ReadVector(v, start, file_space.GetDimensions());

        REQUIRE(v[0] == 0);
        REQUIRE(v[x.GetYPixelsNum() * x.GetXPixelsNum()] == 1);
        REQUIRE(v[2 * x.GetYPixelsNum() * x.GetXPixelsNum()] == 2);
        REQUIRE(v[3 * x.GetYPixelsNum() * x.GetXPixelsNum()] == 3);
        REQUIRE(v[4 * x.GetYPixelsNum() * x.GetXPixelsNum()] == 4);
    }
    // No leftover HDF5 objects
    REQUIRE (H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("HDF5Writer_VDS_missing", "[HDF5][Full]") {
    DiffractionExperiment x(DetectorGeometry(1));
    x.ImagesPerTrigger(5).DataFileCount(5).Compression(JFJochProtoBuf::NO_COMPRESSION).FilePrefix("vds_missing");

    {
        RegisterHDF5Filter();

        StartMessage start_message;
        x.FillMessage(start_message);

        EndMessage end_message;
        end_message.number_of_images = x.GetImageNum() - 1;

        HDF5Writer writer(start_message);

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

        for (int i = 0; i < x.GetImageNum() - 1; i++) {
            for (auto &j: image)
                j = i;
            DataMessage message{};
            message.image.data = (uint8_t *) image.data();
            message.image.size = x.GetPixelsNum() * x.GetPixelDepth();
            message.spots = spots;
            message.number = i;

            REQUIRE_NOTHROW(writer.Write(message));
        }

        REQUIRE_NOTHROW(HDF5Metadata::NXmx(start_message, end_message));
    }
    {
        HDF5File file("vds_missing_master.h5", false, true, false);
        std::unique_ptr<HDF5DataSet> dataset;
        REQUIRE_NOTHROW(dataset = std::make_unique<HDF5DataSet>(file,"/entry/data/data"));
        HDF5DataSpace file_space(*dataset);
        REQUIRE(file_space.GetNumOfDimensions() == 3);
        REQUIRE(file_space.GetDimensions()[0] == x.GetImageNum()-1);
        REQUIRE(file_space.GetDimensions()[1] == x.GetYPixelsNum());
        REQUIRE(file_space.GetDimensions()[2] == x.GetXPixelsNum());

        std::vector<int16_t> v((x.GetImageNum() - 1) * x.GetYPixelsNum() * x.GetXPixelsNum());
        std::vector<hsize_t> start = {0, 0, 0};
        dataset->ReadVector(v, start, file_space.GetDimensions());

        REQUIRE(v[0] == 0);
        REQUIRE(v[x.GetYPixelsNum() * x.GetXPixelsNum()] == 1);
        REQUIRE(v[2 * x.GetYPixelsNum() * x.GetXPixelsNum()] == 2);
        REQUIRE(v[3 * x.GetYPixelsNum() * x.GetXPixelsNum()] == 3);
    }
    // No leftover HDF5 objects
    REQUIRE (H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}


TEST_CASE("HDF5Writer_VDS_zero_images", "[HDF5][Full]") {
    DiffractionExperiment x(DetectorGeometry(1));
    x.ImagesPerTrigger(5).DataFileCount(5).Compression(JFJochProtoBuf::NO_COMPRESSION).FilePrefix("vds_zero");

    {
        RegisterHDF5Filter();

        StartMessage start_message;
        x.FillMessage(start_message);

        EndMessage end_message;
        end_message.number_of_images = 0;

        REQUIRE_NOTHROW(HDF5Metadata::NXmx(start_message, end_message));
    }
    {
        std::unique_ptr<HDF5ReadOnlyFile> file;
        std::unique_ptr<HDF5DataSet> dataset;
        REQUIRE_NOTHROW(file = std::make_unique<HDF5ReadOnlyFile>("vds_zero_master.h5"));
        REQUIRE_THROWS(dataset = std::make_unique<HDF5DataSet>(*file,"/entry/data/data"));
    }
    // No leftover HDF5 objects
    REQUIRE (H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("HDF5NXmx_DataFileName", "[HDF5]") {
    REQUIRE(HDF5Metadata::DataFileName("z/x", 34) == "z/x_data_034.h5");
    REQUIRE(HDF5Metadata::DataFileName("z/x", 0) == "z/x_data_000.h5");
    REQUIRE_THROWS(HDF5Metadata::DataFileName("z/x", 1000));
    REQUIRE_THROWS(HDF5Metadata::DataFileName("z/x", -1));
}

TEST_CASE("HDF5Objects_ExtractFilename", "[HDF5]") {
    REQUIRE(ExtractFilename("filename_data_000001.h5") == "filename_data_000001.h5");
    REQUIRE(ExtractFilename("dir1/filename_data_000001.h5") == "filename_data_000001.h5");
    REQUIRE(ExtractFilename("dir1/dir2/filename_data_000001.h5") == "filename_data_000001.h5");
}

TEST_CASE("HDF5DataType", "[HDF5]") {
    HDF5DataType type1(1,true);
    REQUIRE(type1.GetElemSize() == 1);

    HDF5DataType type2(2,true);
    REQUIRE(type2.GetElemSize() == 2);

    HDF5DataType type4(4,true);
    REQUIRE(type4.GetElemSize() == 4);

    HDF5DataType type8(8,true);
    REQUIRE(type8.GetElemSize() == 8);

    REQUIRE_THROWS(HDF5DataType(7,true));
}