// 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 <bitshuffle/bitshuffle.h>
#include <zstd.h>
#include <iostream>
#include <random>

#include "../compression/JFJochCompressor.h"
#include "../compression/JFJochDecompress.h"
#include "../common/DiffractionExperiment.h"
#include "../writer/HDF5Objects.h"

TEST_CASE("JFjochZstdCompressor_Raw_Block","[ZSTD]") {
    uint8_t src[128*8];
    uint8_t dst[128*8+3];

    for (int i = 0; i < 128*8; i++)
        src[i] = i % 256;

    size_t dst_size = JFJochZstdCompressor::RawBlock(dst, src, 128 * 8, false);

    REQUIRE(dst_size == 128*8+3);
    int diff = 0;
    for (int i = 0; i < 128*8; i++) {
        if (dst[i+3] != i % 256) diff++;
    }
    REQUIRE(diff == 0);
}

TEST_CASE("JFjochZstdCompressor_RLE_Block","[ZSTD]") {
    uint8_t dst[4];
    size_t dst_size = JFJochZstdCompressor::RLEBlock(dst, 0xAF, 128 * 8, false);

    REQUIRE(dst_size == 4);
    REQUIRE(dst[3] == 0xAF);
}

TEST_CASE("JFjochZstdCompressor_Frame_onlyRLE","[ZSTD]") {
    JFJochZstdCompressor compressor;
    uint8_t src[128 * 8];
    uint8_t dst[128 * 8 * 2];
    uint8_t decomp_buffer[128 * 8 * 2];

    for (int i = 0; i < 128 * 8; i++) {
        src[i] = 0;
        dst[i] = 0xFF;
    }

    size_t dst_size = compressor.Compress(dst, (uint64_t *) src, 128*8, 128*8);
    REQUIRE(ZSTD_decompress(decomp_buffer, 128 * 8 * 2, dst, dst_size) == 8 * 128);
    size_t diff = 0;
    for (int i = 0; i < 128 * 8; i++) if (src[i] != decomp_buffer[i]) diff++;
    REQUIRE(diff == 0);

    REQUIRE(dst_size == 4 + 9);
    REQUIRE(dst[9 + 3] == 0);
}

TEST_CASE("JFjochZstdCompressor_Frame_onlyRLE_1block","[ZSTD]") {
    JFJochZstdCompressor compressor;
    uint8_t src[ZSTD_BLOCKSIZE_MAX * 2];
    uint8_t dst[ZSTD_BLOCKSIZE_MAX * 4];
    uint8_t decomp_buffer[ZSTD_BLOCKSIZE_MAX * 4];

    for (int i = 0; i < ZSTD_BLOCKSIZE_MAX * 2; i++) {
        src[i] = 0;
        dst[i] = 0xFF;
    }

    size_t dst_size = compressor.Compress(dst, (uint64_t *) src, ZSTD_BLOCKSIZE_MAX, ZSTD_BLOCKSIZE_MAX);
    REQUIRE(dst_size == 4 + 9);
    REQUIRE(dst[9 + 3] == 0);

    REQUIRE(ZSTD_decompress(decomp_buffer, ZSTD_BLOCKSIZE_MAX * 4, dst, dst_size) == ZSTD_BLOCKSIZE_MAX);
    size_t diff = 0;
    for (int i = 0; i < ZSTD_BLOCKSIZE_MAX; i++) if (src[i] != decomp_buffer[i]) diff++;
    REQUIRE(diff == 0);
}

TEST_CASE("JFjochZstdCompressor_Frame_onlyRLE_2blocks","[ZSTD]") {
    JFJochZstdCompressor compressor;
    uint8_t src[ZSTD_BLOCKSIZE_MAX * 2];
    uint8_t dst[ZSTD_BLOCKSIZE_MAX * 4];
    uint8_t decomp_buffer[ZSTD_BLOCKSIZE_MAX * 4];

    for (int i = 0; i < ZSTD_BLOCKSIZE_MAX * 2; i++) {
        src[i] = 0;
        dst[i] = 0xFF;
    }

    size_t dst_size = compressor.Compress(dst, (uint64_t *) src, ZSTD_BLOCKSIZE_MAX * 2, ZSTD_BLOCKSIZE_MAX * 2);
    REQUIRE(dst_size == 8 + 9);
    REQUIRE(dst[9 + 3] == 0);
    REQUIRE(dst[9 + 7] == 0);

    REQUIRE(ZSTD_decompress(decomp_buffer, ZSTD_BLOCKSIZE_MAX * 4, dst, dst_size) == ZSTD_BLOCKSIZE_MAX * 2);
    size_t diff = 0;
    for (int i = 0; i < ZSTD_BLOCKSIZE_MAX * 2; i++) if (src[i] != decomp_buffer[i]) diff++;
    REQUIRE(diff == 0);
}

TEST_CASE("JFjochZstdCompressor_Frame_onlyRLE_0xFF_2blocks","[ZSTD]") {
    JFJochZstdCompressor compressor;
    uint8_t src[ZSTD_BLOCKSIZE_MAX * 2];
    uint8_t dst[ZSTD_BLOCKSIZE_MAX * 4];
    uint8_t decomp_buffer[ZSTD_BLOCKSIZE_MAX * 4];

    for (int i = 0; i < ZSTD_BLOCKSIZE_MAX * 2; i++) {
        src[i] = 0xFF;
        dst[i] = 0x0;
    }

    size_t dst_size = compressor.Compress(dst, (uint64_t *) src, ZSTD_BLOCKSIZE_MAX * 2, ZSTD_BLOCKSIZE_MAX * 2);
    REQUIRE(dst_size == 8 + 9);
    REQUIRE(dst[9 + 3] == 0xFF);
    REQUIRE(dst[9 + 7] == 0xFF);

    REQUIRE(ZSTD_decompress(decomp_buffer, ZSTD_BLOCKSIZE_MAX * 4, dst, dst_size) == ZSTD_BLOCKSIZE_MAX * 2);
    size_t diff = 0;
    for (int i = 0; i < ZSTD_BLOCKSIZE_MAX * 2; i++) if (src[i] != decomp_buffer[i]) diff++;
    REQUIRE(diff == 0);
}


TEST_CASE("JFjochZstdCompressor_Frame_onlyRAW","[ZSTD]") {
    JFJochZstdCompressor compressor;

    uint8_t src[128 * 8];
    uint8_t dst[128 * 8 * 2];
    uint8_t decomp_buffer[128 * 8 * 2];
    for (int i = 0; i < 128 * 8; i++)
        src[i] = 1 + i % 128;
    size_t dst_size = compressor.Compress(dst, (uint64_t *) src, 128 * 8, 128*8);
    REQUIRE(ZSTD_decompress(decomp_buffer, 128 * 8 * 2, dst, dst_size) == 8 * 128);
    size_t diff = 0;
    for (int i = 0; i < 128 * 8; i++) if (src[i] != decomp_buffer[i]) diff++;
    REQUIRE(diff == 0);

    REQUIRE(dst_size == 3 + 9 + 128 * 8);

    diff = 0;
    for (int i = 0; i < 128 * 8; i++) {
        if (dst[i + 3 + 9] != 1 + i % 128)
            diff++;
    }
    REQUIRE(diff == 0);
}

TEST_CASE("JFjochZstdCompressor_Frame_mixed","[ZSTD]") {
    JFJochZstdCompressor compressor;

    uint8_t src[128 * 8];
    uint8_t dst[128 * 8 * 2];
    uint8_t decomp_buffer[128 * 8 * 2];

    for (int i = 0; i < 128*8; i++)
        src[i] = (i / 128 + 1) % 2 ;
    size_t dst_size = compressor.Compress(dst, (uint64_t *) src, 128 * 8, 128 * 8);
    REQUIRE(dst_size == 9 + 4 * 4 + 4 * (3+128));

    REQUIRE(ZSTD_decompress(decomp_buffer, 128*8*2, dst, dst_size) == 8 * 128);
    size_t diff = 0;
    for (int i = 0; i < 128*8; i++) if (src[i] != decomp_buffer[i]) diff++;
    REQUIRE(diff == 0);
}

TEST_CASE("JFjochZstdCompressor_Frame_simulated","[ZSTD]") {
    for(int seed = 0; seed < 200; seed++) {
        JFJochZstdCompressor compressor;

        // Predictable random number generator
        std::mt19937 g1(seed);
        std::uniform_int_distribution<int16_t> distribution;

        std::vector<int16_t> input(RAW_MODULE_SIZE);
        for (auto &i: input) i = distribution(g1);

        std::vector<char> input_shuffled(RAW_MODULE_SIZE * sizeof(uint16_t));
        bshuf_bitshuffle(input.data(), input_shuffled.data(), RAW_MODULE_SIZE, 2, 4096);
        std::vector<uint8_t> input_compressed(ZSTD_compressBound(RAW_MODULE_SIZE * sizeof(uint16_t)));
        std::vector<uint8_t> output(RAW_MODULE_SIZE * sizeof(uint16_t) * 2);

        size_t dst_size = compressor.Compress(input_compressed.data(), (uint64_t *) input_shuffled.data(),
                                              RAW_MODULE_SIZE * sizeof(uint16_t), RAW_MODULE_SIZE * sizeof(uint16_t));

        REQUIRE(ZSTD_decompress(output.data(), RAW_MODULE_SIZE * sizeof(uint16_t) * 2,
                                input_compressed.data(), dst_size) == RAW_MODULE_SIZE * sizeof(uint16_t));

        REQUIRE(memcmp(input_shuffled.data(), output.data(), RAW_MODULE_SIZE * sizeof(uint16_t)) == 0);
    }
}

TEST_CASE("JFjochZstdCompressor_Frame_zeroes","[ZSTD]") {
    JFJochZstdCompressor compressor;

    std::vector<char> input_shuffled(RAW_MODULE_SIZE * sizeof(uint16_t));
    for (auto &i: input_shuffled) i = 0x0u;

    std::vector<uint8_t> input_compressed(ZSTD_compressBound(RAW_MODULE_SIZE * sizeof(uint16_t)));
    std::vector<uint8_t> output(RAW_MODULE_SIZE * sizeof(uint16_t) * 2);

    size_t dst_size = compressor.Compress(input_compressed.data(), (uint64_t *) input_shuffled.data(),
                                          RAW_MODULE_SIZE * sizeof(uint16_t), RAW_MODULE_SIZE * sizeof(uint16_t));
    REQUIRE(dst_size == 9 + 4 * 8);
    REQUIRE(ZSTD_decompress(output.data(), RAW_MODULE_SIZE * sizeof(uint16_t) * 2,
                            input_compressed.data(), dst_size) == RAW_MODULE_SIZE * sizeof(uint16_t));

    REQUIRE(memcmp(input_shuffled.data(), output.data(), RAW_MODULE_SIZE * sizeof(uint16_t)) == 0);
}

TEST_CASE("JFjochZstdCompressor_Frame_ones","[ZSTD]") {
    JFJochZstdCompressor compressor;

    std::vector<char> input_shuffled(RAW_MODULE_SIZE * sizeof(uint16_t));
    for (auto &i: input_shuffled) i = 0xFFu;

    std::vector<uint8_t> input_compressed(ZSTD_compressBound(RAW_MODULE_SIZE * sizeof(uint16_t)));
    std::vector<uint8_t> output(RAW_MODULE_SIZE * sizeof(uint16_t) * 2);

    size_t dst_size = compressor.Compress(input_compressed.data(), (uint64_t *) input_shuffled.data(),
                                          RAW_MODULE_SIZE * sizeof(uint16_t), RAW_MODULE_SIZE * sizeof(uint16_t));
    REQUIRE(dst_size == 9 + 4 * 8);
    REQUIRE(ZSTD_decompress(output.data(), RAW_MODULE_SIZE * sizeof(uint16_t) * 2,
                            input_compressed.data(), dst_size) == RAW_MODULE_SIZE * sizeof(uint16_t));

    REQUIRE(memcmp(input_shuffled.data(), output.data(), RAW_MODULE_SIZE * sizeof(uint16_t)) == 0);
}

TEST_CASE("JFJochCompressor_JFJochDecompressor_ZSTD","[ZSTD]") {
    DiffractionExperiment x(DetJF4M());
    x.Compression(CompressionAlgorithm::BSHUF_ZSTD).BitDepthImage(32).PixelSigned(true);

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

    for (auto &i: image)
        i = 345;

    JFJochBitShuffleCompressor compressor(x.GetCompressionAlgorithm());

    std::vector<char> tmp(x.GetPixelsNum() * sizeof(int32_t) * 4 + 12);

    auto tmp_size = compressor.Compress(tmp.data(), tmp.size(), image);
    tmp.resize(tmp_size);

    std::vector<int32_t> output;
    REQUIRE_NOTHROW(JFJochDecompress(output, x.GetCompressionAlgorithm(), tmp, x.GetPixelsNum()));
    REQUIRE(output.size() == x.GetPixelsNum());
    REQUIRE(memcmp(image.data(), output.data(), x.GetPixelsNum() * sizeof(int32_t)) == 0);
}

TEST_CASE("JFJochCompressor_DestTooSmall_Throws","[ZSTD]") {
    DiffractionExperiment x(DetJF4M());
    x.Compression(CompressionAlgorithm::BSHUF_ZSTD).BitDepthImage(32).PixelSigned(true);

    std::vector<int32_t> image(x.GetPixelsNum(), 345);

    JFJochBitShuffleCompressor compressor(x.GetCompressionAlgorithm());

    // A destination far too small for the compressed output must throw, not overflow it.
    std::vector<char> tiny(64);
    REQUIRE_THROWS_AS(compressor.Compress(tiny.data(), tiny.size(), image),
                      CompressionBufferTooSmallException);

    // A buffer sized to the worst case never throws.
    std::vector<char> big(MaxCompressedSize(x.GetCompressionAlgorithm(), x.GetPixelsNum(), sizeof(int32_t)));
    REQUIRE_NOTHROW(compressor.Compress(big.data(), big.size(), image));
}

TEST_CASE("JFJochCompressor_JFJochDecompressor_LZ4","[ZSTD]") {
    DiffractionExperiment x(DetJF4M());
    x.Compression(CompressionAlgorithm::BSHUF_LZ4).BitDepthImage(32).PixelSigned(true);

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

    for (auto &i: image)
        i = 5678;

    JFJochBitShuffleCompressor compressor(x.GetCompressionAlgorithm());

    std::vector<char> tmp(x.GetPixelsNum() * sizeof(int32_t) * 4 + 12);

    auto tmp_size = compressor.Compress(tmp.data(), tmp.size(), image);
    tmp.resize(tmp_size);

    std::vector<int32_t> output;
    REQUIRE_NOTHROW(JFJochDecompress(output, x.GetCompressionAlgorithm(), tmp, x.GetPixelsNum()));
    REQUIRE(output.size() == x.GetPixelsNum());
    REQUIRE(memcmp(image.data(), output.data(), x.GetPixelsNum() * sizeof(int32_t)) == 0);
}

TEST_CASE("JFJochDecompressor_None","[ZSTD]") {
    DiffractionExperiment x(DetJF4M());
    x.Compression(CompressionAlgorithm::NO_COMPRESSION).BitDepthImage(32).PixelSigned(true);

    std::vector<int32_t> image(x.GetPixelsNum());
    for (auto &i: image)
        i = 578;

    std::vector<int32_t> output;
    REQUIRE_NOTHROW(JFJochDecompress(output, x.GetCompressionAlgorithm(), image, x.GetPixelsNum()));
    REQUIRE(output.size() == x.GetPixelsNum());
    REQUIRE(memcmp(image.data(), output.data(), x.GetPixelsNum() * sizeof(int32_t)) == 0);
}

TEST_CASE("Bitshuffle_ZSTD","[ZSTD]") {
    std::vector<int32_t> image(RAW_MODULE_SIZE * sizeof(int32_t));
    std::vector<char> compressed(bshuf_compress_zstd_bound(RAW_MODULE_SIZE, 4, 0));
    std::vector<int32_t> decompressed(RAW_MODULE_SIZE * sizeof(int32_t));
    for (int i = 0; i < RAW_MODULE_SIZE; i++)
        image[i] = i;

    auto out_size = bshuf_compress_zstd(image.data(), compressed.data(), RAW_MODULE_SIZE, 4, 0,0);
    REQUIRE(out_size > 0);
    REQUIRE(bshuf_decompress_zstd(compressed.data(), decompressed.data(), RAW_MODULE_SIZE, 4, 0) == out_size);

    REQUIRE(memcmp(image.data(), decompressed.data(), RAW_MODULE_SIZE*sizeof(uint32_t)) == 0);
}
// ---- BSHUF_ZSTD_RLE_HUFF (RLE runs + adaptive Huffman literals, standard zstd frame) ----
// Round-trips through JFJochDecompress (which calls stock ZSTD_decompress per block) and checks
// the output never exceeds the advertised worst-case size. Run on good and adversarial data to
// confirm the algorithm always produces valid output, even when the ratio is poor.
template<class T>
static void RequireHuffRoundTrip(const std::vector<T> &image) {
    JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_ZSTD_RLE_HUFF);
    auto compressed = compressor.Compress(image);

    REQUIRE(compressed.size() <= (size_t) MaxCompressedSize(CompressionAlgorithm::BSHUF_ZSTD_RLE_HUFF,
                                                            image.size(), sizeof(T)));
    std::vector<T> output;
    REQUIRE_NOTHROW(JFJochDecompress(output, CompressionAlgorithm::BSHUF_ZSTD_RLE_HUFF,
                                     compressed, image.size()));
    REQUIRE(output.size() == image.size());
    REQUIRE(memcmp(image.data(), output.data(), image.size() * sizeof(T)) == 0);
}

TEST_CASE("ZstdHuff_PureZeros", "[ZSTD]") {
    std::vector<int16_t> image(200000, 0);            // all runs, no literals
    RequireHuffRoundTrip(image);
}

TEST_CASE("ZstdHuff_Poisson10", "[ZSTD]") {
    std::vector<uint16_t> image(200000);
    std::mt19937 gen(12345);
    std::poisson_distribution<int> dist(10);
    for (auto &v : image) v = (uint16_t) dist(gen);
    RequireHuffRoundTrip(image);
}

TEST_CASE("ZstdHuff_RandomIncompressible", "[ZSTD]") {
    std::vector<uint16_t> image(200000);              // Mersenne-Twister noise: ~zero compression
    std::mt19937 gen(98765);
    for (auto &v : image) v = (uint16_t) gen();
    RequireHuffRoundTrip(image);                      // must still be valid and within bound
}

TEST_CASE("ZstdHuff_MaskLike", "[ZSTD]") {
    std::vector<uint32_t> image(200000, 0);           // extreme sparsity like a pixel_mask
    std::mt19937 gen(555);
    for (auto &v : image) {
        uint32_t r = gen() % 100;
        if (r < 5) v = 1;
        else if (r == 5) v = 0x80000000u;             // module-gap style flag
    }
    RequireHuffRoundTrip(image);
}

TEST_CASE("ZstdHuff_LysoImage", "[ZSTD]") {
    RegisterHDF5Filter();   // bitshuffle filter, needed to read the compressed benchmark dataset
    HDF5ReadOnlyFile data("../../tests/test_data/compression_benchmark.h5");
    HDF5DataSet dataset(data, "/entry/data/data");
    HDF5DataSpace file_space(dataset);
    auto dims = file_space.GetDimensions();
    std::vector<int16_t> image(dims[1] * dims[2]);
    std::vector<hsize_t> start = {0, 0, 0}, size = {1, dims[1], dims[2]};
    dataset.ReadVector(image, start, size);
    RequireHuffRoundTrip(image);
}