// Copyright (2019-2023) Paul Scherrer Institute

#include <iostream>
#include <future>

#include "bitshuffle/bitshuffle_core.h"
#include "zstd/lib/zstd.h"
#include "../writer/HDF5Objects.h"
#include "../common/FrameTransformation.h"
#include "../common/RawToConvertedGeometry.h"
#include "JFJochDecompress.h"

double CheckCompressionThread(const DiffractionExperiment &x,
                              int32_t nimages,
                              int32_t image0,
                              int32_t stride,
                              const std::vector<int16_t> &image,
                              std::vector<uint16_t> &output) {
    double ret = 0;

    FrameTransformation transformation(x);
    for (int32_t i = image0; i < nimages; i += stride) {
        for (int j = 0; j < x.GetModulesNum(); j++ ) {
            transformation.ProcessModule(image.data() + (j + i * x.GetModulesNum()) * RAW_MODULE_SIZE, j, 0);
        }
        ret += transformation.SaveCompressedImage(output.data() + i * x.GetMaxCompressedSize());
    }
    return ret;
}

std::string CheckCompression(const DiffractionExperiment &x, int32_t nimages, const std::vector<int16_t> &image, uint16_t nthreads) {

    double original_size = nimages * x.GetModulesNum() * RAW_MODULE_SIZE * x.GetPixelDepth();

    std::vector<uint16_t> output(nimages * x.GetMaxCompressedSize());

    double compressed_size = 0;
    auto start_time = std::chrono::system_clock::now();

    std::vector<std::future<double>> futures;

    for (int i = 0; i < nthreads; i++)
        futures.emplace_back(std::async(std::launch::async, &CheckCompressionThread, std::ref(x),
                                        nimages, i, nthreads,
                                        std::ref(image), std::ref(output)));

    for (int i = 0; i < nthreads; i++)
        compressed_size += futures[i].get();

    auto end_time = std::chrono::system_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);
    char buf[256];

    snprintf(buf, 255, "%.2fx compression (%5.2f bits/pxl)   Throughput: %5.2f GB/s\n", original_size / compressed_size, compressed_size * 8 / double(nimages * x.GetModulesNum() * RAW_MODULE_SIZE), (original_size / elapsed.count()) / 1000.0);
    return std::string(buf);
}

std::string CheckDecompression(const DiffractionExperiment &x, size_t nimages, const std::vector<int16_t> &image) {
    double original_size = nimages * x.GetModulesNum() * RAW_MODULE_SIZE * x.GetPixelDepth();

    FrameTransformation transformation(x);
    std::vector<char> output[nimages];
    for (auto &v: output)
        v.resize(x.GetMaxCompressedSize());

    std::vector<size_t> compressed_size(nimages);

    for (int i = 0; i < nimages; i++) {
        for (int j = 0; j < x.GetModulesNum(); j++ ) {
            transformation.ProcessModule(image.data() + (j + i * x.GetModulesNum()) * RAW_MODULE_SIZE, j, 0);
        }
        compressed_size[i] = transformation.SaveCompressedImage(output[i].data());
        output[i].resize(compressed_size[i]);
    }

    std::vector<uint16_t> decompress_v;
    auto start_time = std::chrono::system_clock::now();
    for (int i = 0; i < nimages; i++) {
        JFJochDecompress(decompress_v, x.GetCompressionAlgorithm(), output[i], x.GetPixelsNum());
    }

    auto end_time = std::chrono::system_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);
    char buf[256];

    snprintf(buf, 255, "    decompression                    Throughput: %5.2f GB/s\n", (original_size / elapsed.count()) / 1000.0);
    return {buf};
}

int main(int argc, char **argv) {

    RegisterHDF5Filter();

    if ((argc != 2) && (argc != 3) && (argc != 4)) {
        std::cout << "Usage: ./CompressionBenchmark <JF4M hdf5 file> {<nthreads>}" << std::endl;
        exit(EXIT_FAILURE);
    }

    uint32_t nthreads = 1;
    if (argc >= 3)
        nthreads = atol(argv[2]);

    uint64_t nimages = 25;
    if (argc == 4)
        nimages = atol(argv[3]);

    if ((nthreads <= 0) || (nimages <= 0)) {
        std::cerr << "Error in input parameters" << std::endl;
        exit(EXIT_FAILURE);
    }

    HDF5ReadOnlyFile data(argv[1]);

    HDF5DataSet dataset(data, "/entry/data/data");

    HDF5DataSpace file_space(dataset);

    if (file_space.GetNumOfDimensions() != 3) {
        std::cout << "/entry/data/data must be 3D" << std::endl;
        exit(EXIT_FAILURE);
    }

    DiffractionExperiment x(DetectorGeometry(8, 2, 8, 36));

    if ((file_space.GetDimensions()[1] == 2164) && (file_space.GetDimensions()[2] == 2068)) {
        std::cout << "JF4M with gaps detected (2068 x 2164)" << std::endl;
    } else {
        std::cout << "Unknown geometry - exiting" << std::endl;
        exit(EXIT_FAILURE);
    }

    uint64_t nimages_dataset = file_space.GetDimensions()[0];
    std::cout << "Number of images in the dataset: " << nimages_dataset << std::endl;
    if (nimages_dataset > 200) {
        nimages_dataset = 200;
        std::cout << "Using only " << nimages_dataset << " images" << std::endl;
    }
    x.Mode(DetectorMode::Conversion);

    std::vector<int16_t> image_conv ( nimages_dataset * file_space.GetDimensions()[1] * file_space.GetDimensions()[2]);


    std::vector<hsize_t> start = {0,0,0};
    std::vector<hsize_t> dim = {nimages_dataset, file_space.GetDimensions()[1], file_space.GetDimensions()[2]};

    auto start_time = std::chrono::system_clock::now();
    dataset.ReadVector(image_conv, start, dim);
    auto end_time = std::chrono::system_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time);

    std::cout << "Images loaded      " << elapsed.count() / nimages_dataset << " ms/image      "
              << (image_conv.size() * sizeof(uint16_t)) / (1000.0 * elapsed.count())  <<  " MB/s" << std::endl;


    std::cout << "Images to benchmark " << nimages << std::endl << "Threads: " << nthreads << std::endl << std::endl;
    std::vector<int16_t> image( nimages * x.GetModulesNum() * RAW_MODULE_SIZE);

    for (int i = 0; i < nimages; i++) {
        ConvertedToRawGeometry(x,
                               image.data() + i * RAW_MODULE_SIZE * x.GetModulesNum(),
                               image_conv.data() + (i % nimages_dataset) * file_space.GetDimensions()[1] * file_space.GetDimensions()[2]);
    }

    auto image_shuffled = bitshuffle(image, 4096);

    x.MaskChipEdges(false);

    x.Compression(CompressionAlgorithm::NO_COMPRESSION);
    for (int i = 0; i < 3; i++)
        CheckCompression(x, nimages, image, nthreads);

    for (int i = 0; i <3; i++) {
        x.Compression(CompressionAlgorithm::NO_COMPRESSION).Mode(DetectorMode::Raw);
        std::cout << "None (memcpy)               " << CheckCompression(x, nimages, image, nthreads);
        x.Compression(CompressionAlgorithm::NO_COMPRESSION).Mode(DetectorMode::Conversion);

        std::cout << "None (geom transform)       " << CheckCompression(x, nimages, image, nthreads);

        x.Compression(CompressionAlgorithm::BSHUF_LZ4);
        std::cout << "BSHUF/LZ4                   " << CheckCompression(x, nimages, image, nthreads);

        x.Compression(CompressionAlgorithm::BSHUF_ZSTD);
        std::cout << "BSHUF/ZSTD (0)              " << CheckCompression(x, nimages, image, nthreads);

        x.Compression(CompressionAlgorithm::BSHUF_ZSTD_RLE);
        std::cout << "BSHUF/ZSTD (RLE)            " << CheckCompression(x, nimages, image, nthreads);
        std::cout << std::endl;
    }

    std::cout << std::endl << std::endl << "Decompression" << std::endl << std::endl;

    x.Compression(CompressionAlgorithm::BSHUF_LZ4);
    std::cout << "BSHUF/LZ4                   " << CheckDecompression(x, nimages, image);

    x.Compression(CompressionAlgorithm::BSHUF_ZSTD);
    std::cout << "BSHUF/ZSTD                  " << CheckDecompression(x, nimages, image);

    x.Compression(CompressionAlgorithm::BSHUF_ZSTD_RLE);
    std::cout << "BSHUF/ZSTD (RLE)            " << CheckDecompression(x, nimages, image);
}