// Copyright (2019-2022) Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-or-later

#include <catch2/catch.hpp>
#include <nlohmann/json.hpp>

#include "../frame_serialize/JFJochFrameSerializer.h"
#include "../frame_serialize/JFJochFrameDeserializer.h"

TEST_CASE("CBORSerialize_Start", "[CBOR]") {
    JFJochFrameSerializer serializer(8*1024*1024);

    StartMessage message {
            .data_file_count = 3,
            .detector_distance = 0.0005,
            .beam_center_x = 456.6,
            .beam_center_y = 124.3,
            .number_of_images = 34567,
            .image_size_x = 456,
            .image_size_y = 457,
            .pixel_bit_depth = 32,
            .pixel_signed = true,
            .incident_energy = 12400,
            .incident_wavelength = 0.988,
            .frame_time = 0.0001,
            .count_time = 0.000098,
            .saturation_value = 65535,
            .min_value = -32767,
            .pixel_size_x = 0.000075,
            .pixel_size_y = 0.000075,
            .sensor_thickness = 0.0005,
            .sensor_material = "Si",
            .compression_algorithm = CompressionAlgorithm::BSHUF_LZ4,
            .compression_block_size = 8192,
            .unit_cell = {45,37,45,90,108,120},
            .space_group_number = 154,
            .max_spot_count = 250,
            .storage_cell_number = 16,
            .pixel_mask_enabled = true,
            .arm_date = "abc",
            .sample_name = "lyso",
            .file_prefix = "lyso1/dir/file",
            .channels = {"default", "sc2"},
            .detector_description = "EIGER 16M",
            .detector_serial_number = "123",
            .series_unique_id = "bla",
            .series_id = 4567,
            .goniometer = {{"omega", {
                    .increment = 0.1f,
                    .start = 10.0f
            }}},
            .detector_translation = {0.5f, 0.0f, 0.5f},
            .source_name = "Swiss Light Source",
            .source_name_short = "SLS",
            .instrument_name = "X06SA",
            .instrument_name_short = "PXIII",
            .rad_int_bin_number = 35
    };
    message.pixel_mask["sc0"] = std::vector<uint32_t>(456*457, 15);

    REQUIRE_NOTHROW(serializer.SerializeSequenceStart(message));

    auto image = serializer.GetBuffer();

    JFJochFrameDeserializer deserializer;
    REQUIRE_NOTHROW(deserializer.Process(image));
    REQUIRE(deserializer.GetType() == JFJochFrameDeserializer::Type::START);

    StartMessage output_message;
    REQUIRE_NOTHROW(output_message = deserializer.GetStartMessage());
    REQUIRE(output_message.pixel_mask.contains("sc0"));
    CHECK(output_message.data_file_count == message.data_file_count);
    CHECK(output_message.detector_distance == Approx(message.detector_distance));
    CHECK(output_message.beam_center_x == Approx(message.beam_center_x));
    CHECK(output_message.beam_center_y == Approx(message.beam_center_y));
    CHECK(output_message.number_of_images == message.number_of_images);
    CHECK(output_message.image_size_x == message.image_size_x);
    CHECK(output_message.image_size_y == message.image_size_y);
    CHECK(output_message.pixel_bit_depth == message.pixel_bit_depth);
    CHECK(output_message.incident_energy == Approx(message.incident_energy));
    CHECK(output_message.incident_wavelength == Approx(message.incident_wavelength));
    CHECK(output_message.frame_time == Approx(message.frame_time));
    CHECK(output_message.count_time == Approx(message.count_time));
    CHECK(output_message.saturation_value == message.saturation_value);
    CHECK(output_message.min_value == message.min_value);
    CHECK(output_message.pixel_size_x == Approx(message.pixel_size_x));
    CHECK(output_message.pixel_size_y == Approx(message.pixel_size_y));
    CHECK(output_message.sensor_thickness == Approx(message.sensor_thickness));
    CHECK(output_message.sensor_material == message.sensor_material);
    CHECK(output_message.pixel_mask == message.pixel_mask);
    CHECK(output_message.pixel_mask_enabled == message.pixel_mask_enabled);
    CHECK(output_message.compression_algorithm == message.compression_algorithm);
    CHECK(output_message.compression_block_size == message.compression_block_size);
    CHECK(output_message.space_group_number == message.space_group_number);
    CHECK(output_message.arm_date == message.arm_date);
    CHECK(output_message.storage_cell_number == message.storage_cell_number);
    CHECK(output_message.pixel_signed == message.pixel_signed);
    CHECK(output_message.sample_name == message.sample_name);
    CHECK(output_message.file_prefix == message.file_prefix);
    CHECK(output_message.max_spot_count == message.max_spot_count);
    CHECK(output_message.channels == message.channels);
    CHECK(output_message.detector_description == message.detector_description);
    CHECK(output_message.detector_serial_number == message.detector_serial_number);
    CHECK(output_message.series_unique_id == message.series_unique_id);
    CHECK(output_message.series_id == message.series_id);
    CHECK(output_message.source_name == message.source_name);
    CHECK(output_message.source_name_short == message.source_name_short);
    CHECK(output_message.instrument_name == message.instrument_name);
    CHECK(output_message.instrument_name_short == message.instrument_name_short);
    CHECK(output_message.rad_int_bin_number == message.rad_int_bin_number);

    for (int i = 0; i < 3; i++)
        CHECK(output_message.detector_translation[i] == message.detector_translation[i]);
    for (int i = 0; i < 6; i++)
        CHECK(output_message.unit_cell[i] == message.unit_cell[i]);

    REQUIRE(output_message.goniometer.contains("omega"));
    CHECK(output_message.goniometer["omega"].increment == 0.1f);
    CHECK(output_message.goniometer["omega"].start == 10.0f);
}

TEST_CASE("CBORSerialize_End", "[CBOR]") {
    JFJochFrameSerializer serializer(8*1024*1024);

    EndMessage message {
            .number_of_images = 57789,
            .max_receiver_delay = 3456,
            .efficiency = 0.99,
            .write_master_file = true,
            .end_date = "ccc"
    };

    REQUIRE_NOTHROW(serializer.SerializeSequenceEnd(message));

    auto image = serializer.GetBuffer();

    JFJochFrameDeserializer deserializer;
    REQUIRE_NOTHROW(deserializer.Process(image));
    REQUIRE(deserializer.GetType() == JFJochFrameDeserializer::Type::END);

    EndMessage output_message{};
    REQUIRE_NOTHROW(output_message = deserializer.GetEndMessage());
    REQUIRE(output_message.max_receiver_delay == message.max_receiver_delay);
    REQUIRE(output_message.number_of_images == message.number_of_images);
    REQUIRE(output_message.efficiency == Approx(message.efficiency));
    REQUIRE(output_message.end_date == message.end_date);
    REQUIRE(output_message.write_master_file == message.write_master_file);
}

TEST_CASE("CBORSerialize_Image", "[CBOR]") {
    JFJochFrameSerializer serializer(8*1024*1024);

    std::vector<SpotToSave> spots;

    std::vector<uint8_t> test(1024);
    for (int i = 0; i < test.size(); i++)
        test[i] = (i * 253 + 56) % 256;

    CBORImage image {
            .data = test.data(),
            .size = 1024,
            .xpixel = 256,
            .ypixel = 2,
            .pixel_depth_bytes = 2,
            .pixel_is_signed = true,
            .algorithm = CompressionAlgorithm::NO_COMPRESSION,
            .channel = "default"
    };

    DataMessage message {
            .number = 456,
            .image = image,
            .spots = spots,
            .indexing_result = 1
    };

    REQUIRE_NOTHROW(serializer.SerializeImage(message));
    auto serialized = serializer.GetBuffer();

    JFJochFrameDeserializer deserializer;
    REQUIRE_NOTHROW(deserializer.Process(serialized));
    REQUIRE(deserializer.GetType() == JFJochFrameDeserializer::Type::IMAGE);

    auto image_array = deserializer.GetDataMessage();
    REQUIRE(image_array.image.algorithm == CompressionAlgorithm::NO_COMPRESSION);
    REQUIRE(image_array.image.xpixel == 256);
    REQUIRE(image_array.image.ypixel == 2);
    REQUIRE(image_array.image.pixel_depth_bytes == 2);
    REQUIRE(image_array.image.pixel_is_signed);
    REQUIRE(image_array.image.channel == "default");
    REQUIRE(image_array.image.size == test.size());
    REQUIRE(image_array.indexing_result == message.indexing_result);
    REQUIRE(image_array.number == 456);
    REQUIRE(memcmp(image_array.image.data, test.data(), test.size()) == 0);
}

TEST_CASE("CBORSerialize_Image_2", "[CBOR]") {
    JFJochFrameSerializer serializer(8*1024*1024);

    std::vector<SpotToSave> spots;

    std::vector<uint8_t> test(512*1024);
    for (int i = 0; i < test.size(); i++)
        test[i] = (i * 253 + 56) % 256;

    CBORImage image {
            .data = test.data(),
            .size = 1024 * 512,
            .xpixel = 1024,
            .ypixel = 512,
            .pixel_depth_bytes = 1,
            .pixel_is_signed = false,
            .algorithm = CompressionAlgorithm::NO_COMPRESSION,
            .channel = "default"
    };

    DataMessage message {
            .number = 480,
            .image = image,
            .spots = spots,
            .indexing_result = 3
    };

    REQUIRE_NOTHROW(serializer.SerializeImage(message));
    auto serialized = serializer.GetBuffer();

    JFJochFrameDeserializer deserializer;
    REQUIRE_NOTHROW(deserializer.Process(serialized));
    REQUIRE(deserializer.GetType() == JFJochFrameDeserializer::Type::IMAGE);


    auto image_array = deserializer.GetDataMessage();
    REQUIRE(image_array.image.algorithm == CompressionAlgorithm::NO_COMPRESSION);
    REQUIRE(image_array.image.xpixel == 1024);
    REQUIRE(image_array.image.ypixel == 512);
    REQUIRE(image_array.image.pixel_depth_bytes == 1);
    REQUIRE(!image_array.image.pixel_is_signed);
    REQUIRE(image_array.image.channel == "default");
    REQUIRE(image_array.image.size == test.size());
    REQUIRE(image_array.indexing_result == message.indexing_result);
    REQUIRE(image_array.number == 480);
    REQUIRE(memcmp(image_array.image.data, test.data(), test.size()) == 0);
}

TEST_CASE("CBORSerialize_Image_Compressed", "[CBOR]") {
    JFJochFrameSerializer serializer(8*1024*1024);

    std::vector<SpotToSave> spots;

    std::vector<uint8_t> test(512);
    for (int i = 0; i < test.size(); i++)
        test[i] = (i * 253 + 56) % 256;

    CBORImage image {
            .data = test.data(),
            .size = 512,
            .xpixel = 256,
            .ypixel = 2,
            .pixel_depth_bytes = 4,
            .pixel_is_signed = true,
            .algorithm = CompressionAlgorithm::BSHUF_LZ4,
            .channel = "default"
    };

    DataMessage message {
            .number = 456,
            .image = image,
            .spots = spots
    };

    REQUIRE_NOTHROW(serializer.SerializeImage(message));
    auto serialized = serializer.GetBuffer();

    JFJochFrameDeserializer deserializer;
    REQUIRE_NOTHROW(deserializer.Process(serialized));
    REQUIRE(deserializer.GetType() == JFJochFrameDeserializer::Type::IMAGE);

    auto image_array = deserializer.GetDataMessage();
    REQUIRE(image_array.image.algorithm == CompressionAlgorithm::BSHUF_LZ4);
    REQUIRE(image_array.image.xpixel == 256);
    REQUIRE(image_array.image.ypixel == 2);
    REQUIRE(image_array.image.pixel_depth_bytes == 4);
    REQUIRE(image_array.image.channel == "default");
    REQUIRE(image_array.image.size == test.size());
    REQUIRE(image_array.image.pixel_is_signed == false);
    REQUIRE(image_array.number == 456);
    REQUIRE(memcmp(image_array.image.data, test.data(), test.size()) == 0);
}

TEST_CASE("CBORSerialize_Image_Rad_Int_Profile", "[CBOR]") {
    JFJochFrameSerializer serializer(8 * 1024 * 1024);

    std::vector<uint8_t> test(1024);
    for (int i = 0; i < test.size(); i++)
        test[i] = (i * 253 + 56) % 256;

    CBORImage image{
            .data = test.data(),
            .size = 1024,
            .xpixel = 256,
            .ypixel = 2,
            .pixel_depth_bytes = 2,
            .algorithm = CompressionAlgorithm::NO_COMPRESSION,
            .channel = "default"
    };

    DataMessage message{
            .number = 789,
            .image = image,
            .rad_int_profile = {4.0, 5.0, 7.0, 12.0, 13.25, 0.125}
    };

    REQUIRE_NOTHROW(serializer.SerializeImage(message));

    auto serialized = serializer.GetBuffer();

    JFJochFrameDeserializer deserializer;
    REQUIRE_NOTHROW(deserializer.Process(serialized));
    REQUIRE(deserializer.GetType() == JFJochFrameDeserializer::Type::IMAGE);

    auto image_array = deserializer.GetDataMessage();
    REQUIRE(image_array.number == 789);

    REQUIRE(image_array.image.size == test.size());
    REQUIRE(memcmp(image_array.image.data, test.data(), test.size()) == 0);
    REQUIRE(image_array.rad_int_profile == message.rad_int_profile);
}

TEST_CASE("CBORSerialize_Image_Spots", "[CBOR]") {
    JFJochFrameSerializer serializer(8*1024*1024);

    std::vector<SpotToSave> spots;
    spots.push_back(SpotToSave{.x = 7, .y = 8, .intensity = 34});
    spots.push_back(SpotToSave{.x = 37, .y = 48, .intensity = 123});

    std::vector<uint8_t> test(1024);
    for (int i = 0; i < test.size(); i++)
        test[i] = (i * 253 + 56) % 256;

    CBORImage image {
            .data = test.data(),
            .size = 1024,
            .xpixel = 256,
            .ypixel = 2,
            .pixel_depth_bytes = 2,
            .algorithm = CompressionAlgorithm::NO_COMPRESSION,
            .channel = "default"
    };

    DataMessage message {
            .number = 789,
            .image = image,
            .spots = spots
    };

    REQUIRE_NOTHROW(serializer.SerializeImage(message));

    auto serialized = serializer.GetBuffer();

    JFJochFrameDeserializer deserializer;
    REQUIRE_NOTHROW(deserializer.Process(serialized));
    REQUIRE(deserializer.GetType() == JFJochFrameDeserializer::Type::IMAGE);

    auto image_array = deserializer.GetDataMessage();
    REQUIRE(image_array.number == 789);

    REQUIRE(image_array.image.size == test.size());
    REQUIRE(memcmp(image_array.image.data, test.data(), test.size()) == 0);
    REQUIRE(image_array.spots.size() == 2);

    REQUIRE(image_array.spots[0].intensity == 34);

    REQUIRE(image_array.spots[1].x == 37);
    REQUIRE(image_array.spots[1].y == 48);
    REQUIRE(image_array.spots[1].intensity == 123);
}

TEST_CASE("CBORSerialize_End_NLohmann", "[CBOR]") {
    JFJochFrameSerializer serializer(8*1024*1024);
    EndMessage message {
            .number_of_images = 57789,
            .max_receiver_delay = 3456,
            .efficiency = 0.99,
            .write_master_file = true,
            .end_date = "ccc"
    };

    REQUIRE_NOTHROW(serializer.SerializeSequenceEnd(message));

    auto serialized = serializer.GetBuffer();

    nlohmann::json j ;
    REQUIRE_NOTHROW(j = nlohmann::json::from_cbor(serialized));
    REQUIRE(j["type"] == "end");
    REQUIRE(j["write_master_file"] == true);
    REQUIRE(j["end_date"] == "ccc");
}