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

#include "../common/print_license.h"
#include "../common/Logger.h"
#include "../common/DiffractionExperiment.h"
#include "../writer/HDF5Objects.h"
#include "../common/PixelMask.h"
#include "../image_pusher/HDF5FilePusher.h"
#include "../image_puller/TestImagePuller.h"
#include "../acquisition_device/AcquisitionDeviceGroup.h"
#include "../receiver/JFJochReceiverService.h"
#include "JFJochCompressor.h"


void print_usage(Logger &logger) {
    logger.Info("Usage ./jfjoch_fpga_test {<options>} <path to repository>");
    logger.Info("Options:");
    logger.Info("   -i<num>      Number of images");
    logger.Info("   -N<num>      Number of image processing threads (default: 8)");
    logger.Info("   -P<txt>      NUMA policy: none|n2g2|n8g4|n8g4_hbm (default: none)");
    logger.Info("   -F{<txt>}    Write file, optional parameter is name (default: lyso_lite_perf_test)");
    logger.Info("   -X<txt>      Indexing (none|fft|ffbidx), ffbidx is default");
    logger.Info("   -t<num>      Indexing thread pool size (default: 4)");
    logger.Info("   -f<num>      FFT indexing search vectors");
    logger.Info("   -Q           Quick integration");
}

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

    Logger logger("jfjoch_fpga_test");
    logger.Verbose(false);

    bool use_ml = false;
    uint16_t nthreads = 8;
    size_t nimages = 1000;
    std::string numa_policy_name;
    std::string filename = "";
    std::string ml_model = "";
    IndexingAlgorithmEnum indexing = IndexingAlgorithmEnum::FFBIDX;
    std::optional<int64_t> fft_num_vectors;
    uint16_t indexing_threads = 4;
    bool quick_integrate = false;

    RegisterHDF5Filter();

    if (argc == 1) {
        print_usage(logger);
        exit(EXIT_FAILURE);
    }

    int opt;
    while ((opt = getopt(argc, argv, "N:P:i:F::QvX:t:f:")) != -1) {
        switch (opt) {
            case 'N':
                nthreads = atol(optarg);
                break;
            case 'P':
                numa_policy_name = std::string(optarg);
                break;
            case 'i':
                nimages = atol(optarg);
                break;
            case 'X':
                if (std::string(optarg) == "none")
                    indexing = IndexingAlgorithmEnum::None;
                else if (std::string(optarg) == "fft" || std::string(optarg) == "FFT")
                    indexing = IndexingAlgorithmEnum::FFT;
                else if (std::string(optarg) == "ffbidx" || std::string(optarg) == "FFBIDX")
                    indexing = IndexingAlgorithmEnum::FFBIDX;
                break;
            case 't':
                indexing_threads = atol(optarg);
                break;
            case 'f':
                fft_num_vectors = atol(optarg);
                break;
            case 'F':
                if (optarg)
                    filename = std::string(optarg);
                else
                    filename = "lyso_lite_perf_test";
                break;
            case 'Q':
                quick_integrate = true;
                break;
            case 'v':
                logger.Verbose(true);
                break;
            default: /* '?' */
                print_usage(logger);
                exit(EXIT_FAILURE);
        }
    }

    if (optind != argc - 1) {
        print_usage(logger);
        exit(EXIT_FAILURE);
    }

    std::string jfjoch_path = std::string(argv[optind]) + "/";


    DiffractionExperiment experiment(DetDECTRIS(2068, 2164, "Test", {}));

    experiment.ImagesPerTrigger(nimages).NumTriggers(1).UseInternalPacketGenerator(true).ImagesPerFile(1000)
            .FilePrefix(filename).JungfrauConvPhotonCnt(false).SetFileWriterFormat(FileWriterFormat::NXmxVDS).OverwriteExistingFiles(true)
            .DetectorDistance_mm(75).BeamY_pxl(1136).BeamX_pxl(1090).IncidentEnergy_keV(12.4)
            .SetUnitCell(UnitCell{.a = 36.9, .b = 78.95, .c = 78.95, .alpha =90, .beta = 90, .gamma = 90}).PixelSigned(true);

    PixelMask pixel_mask(experiment);

    experiment.ROI().SetROI(ROIDefinition{.boxes = {ROIBox("ROI1", 123, 180, 500,800) }});

    // Load example image
    HDF5ReadOnlyFile data(jfjoch_path + "tests/test_data/compression_benchmark.h5");
    HDF5DataSet dataset(data, "/entry/data/data");
    HDF5DataSpace file_space(dataset);

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

    std::vector<hsize_t> start = {4,0,0};
    std::vector<hsize_t> file_size = {1, file_space.GetDimensions()[1], file_space.GetDimensions()[2]};
    dataset.ReadVector(image_conv, start, file_size);

    JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4);
    auto image_compressed = compressor.Compress(image_conv);

    HDF5FilePusher pusher;

    auto puller = std::make_shared<TestImagePuller>(nimages + 5);

    StartMessage start_msg;
    experiment.FillMessage(start_msg);

    puller->Put(ImagePullerOutput{
            .cbor = std::make_shared<CBORStream2DeserializerOutput>(start_msg)
    });

    DataMessage data_msg;
    data_msg.image = CompressedImage(image_compressed,
                                     file_space.GetDimensions()[2],
                                     file_space.GetDimensions()[1],
                                     CompressedImageMode::Int16,
                                     CompressionAlgorithm::BSHUF_LZ4);

    for (int i = 0; i < experiment.GetImageNum(); i++) {
        data_msg.number = i;
        puller->Put(ImagePullerOutput{
                .cbor = std::make_shared<CBORStream2DeserializerOutput>(data_msg)
        });
    }

    EndMessage end_msg{};

    puller->Put(ImagePullerOutput{
            .cbor = std::make_shared<CBORStream2DeserializerOutput>(end_msg)
    });

    AcquisitionDeviceGroup group;
    JFJochReceiverService service(group, logger, pusher);

    service.NUMAPolicy(numa_policy_name);
    service.NumThreads(nthreads);

    IndexingSettings i_settings;
    i_settings.Algorithm(indexing);
    if (fft_num_vectors)
        i_settings.FFT_NumVectors(fft_num_vectors.value());
    i_settings.IndexingThreads(indexing_threads);
    experiment.ImportIndexingSettings(i_settings);
    service.Indexing(i_settings);

    SpotFindingSettings settings = DiffractionExperiment::DefaultDataProcessingSettings();
    settings.signal_to_noise_threshold = 2.5;
    settings.photon_count_threshold = 5;
    settings.min_pix_per_spot = 1;
    settings.max_pix_per_spot = 200;
    settings.high_resolution_limit = 2.0;
    settings.low_resolution_limit = 50.0;
    settings.quick_integration = quick_integrate;
    service.SetSpotFindingSettings(settings);

    auto start_time = std::chrono::system_clock::now();
    service.Start(experiment, pixel_mask, nullptr, puller);
    auto output = service.Stop();
    auto end_time = std::chrono::system_clock::now();

    double receiving_time_s = static_cast<double>(output.end_time_ms - output.start_time_ms) / 1000.0;
    logger.Info("Throughput {:.1f} Hz", experiment.GetImageNum() / receiving_time_s);
    if (output.status.indexing_rate)
        logger.Info("Indexing rate {:.0f}%", output.status.indexing_rate.value() * 100.0);
    if (output.status.max_receive_delay)
        logger.Info("Max delay {}", output.status.max_receive_delay.value());
}