// Copyright (2019-2024) Paul Scherrer Institute

#include <iostream>
#include <cmath>
#include <vector>
#include <future>

#include "../writer/HDF5Objects.h"
#include "../image_analysis/AzimuthalIntegration.h"
#include "../compression/JFJochDecompress.h"

std::mutex m;

void CalcRadialIntegration(const DiffractionExperiment &x, const AzimuthalIntegrationMapping &mapping,
                           std::vector<double> *result, char *filename, int run, int stride) {
    std::unique_lock<std::mutex> ul(m);
    HDF5ReadOnlyFile x_data_file(filename);
    HDF5DataSet x_dataset(x_data_file, "/entry/data/data");
    HDF5DataSpace x_file_space(x_dataset);

    std::cout << "# " << run << " " << filename << " images " << x_file_space.GetDimensions()[0] << " stride " << stride
              << std::endl;

    AzimuthalIntegration rad_int(x, mapping);

    std::vector<uint8_t> buffer;
    std::vector<int16_t> uncompressed(x.GetPixelsNum());

    for (hsize_t i = 0; i < x_file_space.GetDimensions()[0]; i += stride) {
        std::vector<hsize_t> start = {i, 0, 0};
        bool image_loaded = true;
        try {
            x_dataset.ReadDirectChunk(buffer, start);
        } catch (const std::exception &e) {
            std::cerr << "Error reading image " << i << " " << e.what() << std::endl;
            image_loaded = false;
        }
        if (image_loaded) {
            ul.unlock();
            JFJochDecompress(uncompressed, CompressionAlgorithm::BSHUF_LZ4, buffer, x.GetPixelsNum());
            rad_int.Process(uncompressed.data(), x.GetPixelsNum());
            ul.lock();
        }

    }

    rad_int.GetResult(*result);
}

void GetGeometry(DiffractionExperiment &x, HDF5Object &master_file) {
    x.BeamX_pxl(
            HDF5DataSet(master_file, "/entry/instrument/detector/beam_center_x").ReadScalar<double>());
    x.BeamY_pxl(
            HDF5DataSet(master_file, "/entry/instrument/detector/beam_center_y").ReadScalar<double>());
    x.DetectorDistance_mm(
            HDF5DataSet(master_file, "/entry/instrument/detector/detector_distance").ReadScalar<double>() *1e3);
    x.PhotonEnergy_keV(WVL_1A_IN_KEV /
                       HDF5DataSet(master_file, "/entry/instrument/beam/incident_wavelength").ReadScalar<double>());
}

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

    if (argc < 4) {
        std::cout << "Usage ./ImageAverage <name of master dataset> <stride> {<name of data dataset>+}" << std::endl;
        exit(EXIT_FAILURE);
    }

    int stride = atoi(argv[2]);
    size_t nfiles = argc - 3;

    if (stride <= 0) {
        std::cout << "Stride cannot be less than 1" << std::endl;
        exit(EXIT_FAILURE);
    }

    HDF5ReadOnlyFile master_file(argv[1]);
    HDF5ReadOnlyFile data_file(argv[3]);
    HDF5DataSet dataset(data_file, "/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;

    if ((file_space.GetDimensions()[1] == 2164) && (file_space.GetDimensions()[2] == 2068)) {
        std::cout << "# JF4M with gaps detected (2068 x 2164)" << std::endl;
        x = DiffractionExperiment(DetectorGeometry(8, 2, 8, 36));
    } else if ((file_space.GetDimensions()[1] == 3264) && (file_space.GetDimensions()[2] == 3106)) {
        std::cout << "# JF9M with gaps detected (3106x3264)" << std::endl;
        x = DiffractionExperiment(DetectorGeometry(18, 3, 8, 36));
    } else {
        std::cout << "# Unknown geometry - exiting" << std::endl;
        exit(EXIT_FAILURE);
    }

    GetGeometry(x, master_file);

    std::cout << "# beam_center " << x.GetBeamX_pxl() << " " << x.GetBeamY_pxl() << std::endl;
    std::cout << "# detector_distance " << x.GetDetectorDistance_mm() << std::endl;
    std::cout << "# energy " << x.GetPhotonEnergyForConversion_keV() << std::endl;

    AzimuthalIntegrationMapping rad_int_map(x);

    auto result_x = rad_int_map.GetBinToQ();

    std::vector<std::vector<double> > result(nfiles);

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

    for (int i = 0; i < nfiles; i++)
        futures.emplace_back(std::async(std::launch::async, &CalcRadialIntegration,
                                        std::ref(x), std::ref(rad_int_map), &result[i], argv[i+3], i, stride));

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

    for (int i = 0; i < result_x.size(); i++) {
        std::cout << result_x[i] << " ";
        for (int f = 0; f < nfiles; f++)
            std::cout << result[f][i] << " ";
        std::cout << std::endl;
    }
}