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

#include <cstdlib>
#include <iostream>
#include "../common/Logger.h"
#include "../writer/FileWriter.h"
#include "../receiver/FrameTransformation.h"
#include "../common/RawToConvertedGeometry.h"
#include "../common/PixelMask.h"
#include "../writer/HDF5NXmx.h"
#include "../common/print_license.h"

int main(int argc, char **argv) {
    print_license("jfjoch_hdf5_test");

    Logger logger("jfjoch_hdf5_test");

    RegisterHDF5Filter();

    if ((argc < 2) || (argc > 4)) {
        std::cout << "Usage: ./jfjoch_hdf5_test <JF4M hdf5 file> {{<#images>} <rate in Hz>}" << std::endl;
        std::cout << std::endl;
        std::cout << "Env. variables:" << std::endl;
        std::cout << "HDF5DATASET_WRITE_TEST_PREFIX" << std::endl;
        std::cout << "HDF5MASTER_NEW_FORMAT" << std::endl;
        exit(EXIT_FAILURE);
    }

    int64_t nimages_out = 100;
    double rate = 2200;

    if (argc >= 3)
        nimages_out = atoi(argv[2]);
    if (argc >= 4)
        rate = atof(argv[3]);

    std::chrono::microseconds period_us((rate == 0) ? 0 : (int64_t) (1.0e6 / rate));

    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(DetJF4M());
    x.Summation(1);

    // Set metadata for the compression_benchmark.h5 dataset
    x.BeamX_pxl(1090).BeamY_pxl(1136).DetectorDistance_mm(75).IncidentEnergy_keV(WVL_1A_IN_KEV);
    x.MaskModuleEdges(true);
    x.MaskChipEdges(true);

    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 = file_space.GetDimensions()[0];
    logger.Info("Number of images in the original dataset: " + std::to_string(nimages));

    // Set file name
    if (std::getenv("HDF5DATASET_WRITE_TEST_PREFIX") == nullptr)
        x.FilePrefix("writing_test");
    else
        x.FilePrefix(std::getenv("HDF5DATASET_WRITE_TEST_PREFIX"));

    if (std::getenv("HDF5MASTER_NEW_FORMAT") != nullptr) {
        std::cout << "Using new format for HDF5 master file" << std::endl;
        x.SetFileWriterFormat(FileWriterFormat::NXmxVDS);
    } else
        x.SetFileWriterFormat(FileWriterFormat::NXmxLegacy);

    x.Mode(DetectorMode::Conversion).ImagesPerTrigger(nimages);

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

    std::vector<hsize_t> start = {0,0,0};
    dataset.ReadVector(image_conv, start, file_space.GetDimensions());

    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 * file_space.GetDimensions()[1] * file_space.GetDimensions()[2]);
    }

    FrameTransformation transformation(x);

    std::vector<size_t> output_size(nimages);
    std::vector<std::vector<uint8_t> > output(nimages);
    for (auto &i: output) i.resize(x.GetMaxCompressedSize());

    for (int i = 0; i < nimages; i++) {
        for (int j = 0; j < 8; j++)
            transformation.ProcessModule(image.data() + (i * x.GetModulesNum() + j) * RAW_MODULE_SIZE, j, 0);

        auto compressed_image = transformation.GetCompressedImage();

        output_size[i] = compressed_image.GetCompressedSize();
        output[i].resize(compressed_image.GetCompressedSize());
        memcpy(output[i].data(), compressed_image.GetCompressed(), compressed_image.GetCompressedSize());
    }

    x.ImagesPerTrigger(nimages_out);
    logger.Info("Number of images to write: " + std::to_string(nimages_out));

    StartMessage start_message;
    x.FillMessage(start_message);

    PixelMask calib(x);

    start_message.pixel_mask["default"] = calib.GetMask(x);

    // Master & calibration files are written outside of timing routine
    auto fileset = std::make_unique<FileWriter>(start_message);

    auto start_time = std::chrono::system_clock::now();
    logger.Info("Writing " + std::to_string(nimages_out) + " images");

    std::vector<SpotToSave> spots;

    size_t total_image_size = 0;
    for (int i = 0; i < nimages_out; i++) {
        std::this_thread::sleep_until(start_time + i * period_us);

        DataMessage message{};
        message.image = CompressedImage(output[i % nimages], x.GetXPixelsNum(), x.GetYPixelsNum(),
                                        x.GetImageMode(), x.GetCompressionAlgorithm());
        message.spots = spots;
        message.number = i;

        fileset->WriteHDF5(message);

        total_image_size += output_size[i % nimages];
    }

    EndMessage end_message;
    end_message.max_image_number = x.GetImageNum();
    fileset->WriteHDF5(end_message);
    fileset.reset(); // Ensure data file is closed here
    auto end_time = std::chrono::system_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);

    int64_t bandwidth_MBs = (double) total_image_size / (double)elapsed.count();
    int64_t frequency_Hz = (nimages_out * 1e6) / (double) (elapsed.count());

    logger.Info("Writing done");

    logger.Info("Write speed " + std::to_string(bandwidth_MBs) + " MB/s");
    logger.Info("Frequency " + std::to_string(frequency_Hz) + " Hz");
}