// 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();

    int64_t nimages_out = 100;

    std::string prefix = "writing_test";
    FileWriterFormat format = FileWriterFormat::NXmxLegacy;
    std::optional<int64_t> images_per_file;
    std::optional<float> rotation;

    int opt;
    while ((opt = getopt(argc, argv, "o:n:Vf:R:S")) != -1) {
        switch (opt) {
            case 'o':
                prefix = optarg;
                break;
            case 'n':
                nimages_out = atoll(optarg);
                break;
            case 'V':
                format = FileWriterFormat::NXmxVDS;
                break;
            case 'S':
                format = FileWriterFormat::NXmxIntegrated;
                break;
            case 'R':
                rotation = atof(optarg);
                break;
            case 'f':
                images_per_file = atoll(optarg);
                if (images_per_file.value() <= 0) {
                    std::cerr << "Invalid number of images per file: " << optarg << std::endl;
                    exit(EXIT_FAILURE);
                }
                images_per_file = atoll(optarg);
                break;
            default:
                std::cout << "Usage: ./jfjoch_hdf5_test <JF4M hdf5 file> [-o <prefix>] [-n <num images>] [-V] [-f <num images per file>] [-R<rotation angle>]" << std::endl;
                exit(EXIT_FAILURE);
        }
    }

    if (optind >= argc) {
        std::cout << "Usage: ./jfjoch_hdf5_test <JF4M hdf5 file> [-o <prefix>] [-n <num images>] [-V] [-f <num images per file>] [-R<rotation angle>]" << std::endl;
        exit(EXIT_FAILURE);
    }

    HDF5ReadOnlyFile data(argv[optind]);
    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 (rotation && rotation.value() != 0.0)
        x.Goniometer(GoniometerAxis("omega", 0, rotation.value(), Coord(-1,0,0), std::nullopt));

    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
    x.FilePrefix(prefix);
    x.SetFileWriterFormat(format);
    x.OverwriteExistingFiles(true);

    if (images_per_file.has_value())
        x.ImagesPerFile(images_per_file.value());

    x.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++) {
        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");
}