// Copyright (2019-2022) Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-or-later

#include <iostream>
#include <iomanip>

#include "bitshuffle/bitshuffle_core.h"
#include "../writer/HDF5Objects.h"
#include "../common/RawToConvertedGeometry.h"
#include "../image_analysis/GPUImageAnalysis.h"

#include "../image_analysis/IndexerWrapper.h"

#include "../image_analysis/RadialIntegration.h"
#include "../common/Logger.h"
#include "../image_analysis/PredictSpotsOnDetector.h"

#define make_unit_cell(a1,a2,a3,a4,a5,a6) UnitCell{.a = a1, .b = a2, .c = a3, .alpha = a4, .beta = a5, .gamma = a6}

Logger logger{"DataAnalysisPerfTest"};

auto TestAll(const DiffractionExperiment &experiment, const JFJochProtoBuf::DataProcessingSettings &settings,
             GPUImageAnalysis &spot_finder, int16_t* image, size_t nimages) {
    IndexerWrapper indexer;
    indexer.Setup(experiment.GetUnitCell());

    spot_finder.SetInputBuffer(image);

    auto start_time = std::chrono::system_clock::now();
    for (int i = 0; i < nimages; i++) {
        std::vector<DiffractionSpot> spots;
        std::vector<float> result;
        spot_finder.LoadDataToGPU();
        spot_finder.RunSpotFinder(settings);
        spot_finder.GetSpotFinderResults(experiment, settings, spots);
        spot_finder.RunRadialIntegration();
        std::vector<Coord> recip;
        for (const auto& s: spots)
            recip.emplace_back(s.ReciprocalCoord(experiment));
        indexer.Run(recip);
        spot_finder.GetRadialIntegrationProfile(result);
    }

    auto end_time = std::chrono::system_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);

    std::ostringstream strstream;
    logger.Info("{:20s} {:8.1f} ms/image", "Full",
                elapsed.count() / (1000.0 * (double) nimages));

    return strstream.str();
}

auto TestAllWithROI(const DiffractionExperiment &experiment, const JFJochProtoBuf::DataProcessingSettings &settings,
             GPUImageAnalysis &spot_finder, int16_t* image, size_t nimages) {
    IndexerWrapper indexer;
    indexer.Setup(experiment.GetUnitCell());

    spot_finder.SetInputBuffer(image);

    std::vector<int16_t> roi_image(experiment.GetPixelsNum());

    auto start_time = std::chrono::system_clock::now();
    for (int i = 0; i < nimages; i++) {
        ROIFilter filter(experiment.GetXPixelsNum(), experiment.GetYPixelsNum());

        std::vector<DiffractionSpot> spots;
        std::vector<float> result;
        spot_finder.LoadDataToGPU();
        spot_finder.RunSpotFinder(settings);
        spot_finder.GetSpotFinderResults(experiment, settings, spots);
        spot_finder.RunRadialIntegration();
        std::vector<Coord> recip;
        for (const auto& s: spots)
            recip.emplace_back(s.ReciprocalCoord(experiment));
        auto indexer_ret = indexer.Run(recip);
        spot_finder.GetRadialIntegrationProfile(result);
        if (!indexer_ret.empty()) {
            PredictSpotsOnDetector(filter, experiment, indexer_ret[0].l);
            filter.Apply<int16_t>(roi_image, INT16_MIN);
        }
    }

    auto end_time = std::chrono::system_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);

    std::ostringstream strstream;
    logger.Info("{:20s} {:8.1f} ms/image", "Full+ROI",
                elapsed.count() / (1000.0 * (double) nimages));

    return strstream.str();
}

void TestIndexing() {
    constexpr const int nexec  = 5;

    std::vector<Coord> hkl;
    for (int i = 1; i < 7; i++)
        for (int j = 1; j<6; j++)
            for (int k = 1; k < 4; k++)
                hkl.emplace_back(i,j,k);

    std::vector<UnitCell> cells;
    cells.emplace_back(make_unit_cell(30,40,50,90,90,90));
    cells.emplace_back(make_unit_cell(80,80,90,90,90,120));
    cells.emplace_back(make_unit_cell(40,45,80,90,82.5,90));
    cells.emplace_back(make_unit_cell(80,80,150,90,90,90));
    cells.emplace_back(make_unit_cell(30,40,50,90,90,90));

    IndexerWrapper indexer;

    for (auto &c: cells) {
        CrystalLattice l(c);
        CrystalLattice recip_l = l.ReciprocalLattice();

        std::vector<Coord> recip;
        recip.reserve(hkl.size());
        for (const auto &i: hkl)
            recip.emplace_back(i.x * recip_l.Vec0() + i.y * recip_l.Vec1() + i.z * recip_l.Vec2());

        auto start_time = std::chrono::system_clock::now();
        indexer.Setup(c);
        for (int i = 0; i < nexec; i++) {
            if (indexer.Run(recip).empty())
                logger.Warning("WARNING! No solution found");
        }

        auto end_time = std::chrono::system_clock::now();
        auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);
        logger.Info("Fast feedback index. {:8.1f} ms/image", elapsed.count() / (1000.0 * nexec));
    }
}

auto TestSpotFinder(const DiffractionExperiment &experiment, const JFJochProtoBuf::DataProcessingSettings &settings,
                    GPUImageAnalysis &spot_finder, int16_t* image, size_t nimages) {

    std::vector<DiffractionSpot> spots;

    spot_finder.SetInputBuffer(image);

    auto start_time = std::chrono::system_clock::now();
    for (int i = 0; i < nimages; i++) {
        spot_finder.LoadDataToGPU();
        spot_finder.RunSpotFinder(settings);
        spot_finder.GetSpotFinderResults(experiment, settings, spots);
    }

    auto end_time = std::chrono::system_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);

    std::ostringstream strstream;
    logger.Info("{:20s} {:8.1f} ms/image   {:5d} spots", "Spot finding",
                elapsed.count() / (1000.0 * (double) nimages), spots.size());

    return strstream.str();
}

auto TestSpotFinderWithoutCopyToDevice(const DiffractionExperiment &experiment, const JFJochProtoBuf::DataProcessingSettings &settings,
                    GPUImageAnalysis &spot_finder, int16_t* image, size_t nimages) {

    std::vector<DiffractionSpot> spots;

    spot_finder.SetInputBuffer(image);
    spot_finder.LoadDataToGPU();

    auto start_time = std::chrono::system_clock::now();
    for (int i = 0; i < nimages; i++) {
        spot_finder.RunSpotFinder(settings);
        spot_finder.GetSpotFinderResults(experiment, settings, spots);
    }

    auto end_time = std::chrono::system_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);

    std::ostringstream strstream;
    logger.Info("{:20s} {:8.1f} ms/image   {:5d} spots", "Spot finding",
                elapsed.count() / (1000.0 * (double) nimages), spots.size());

    return strstream.str();
}

auto TestRadialIntegrationGPU(const DiffractionExperiment &x,
                              const JFJochProtoBuf::DataProcessingSettings &settings,
                              int16_t* image, size_t nimages) {

    uint32_t nredo = 20;
    RadialIntegrationMapping mapping(x);

    std::vector<uint8_t> one_byte_mask(x.GetPixelsNum(), 1);
    GPUImageAnalysis integration(x.GetXPixelsNum(), x.GetYPixelsNum(), one_byte_mask, mapping);

    integration.SetInputBuffer(image);

    std::vector<float> result;

    auto start_time = std::chrono::system_clock::now();
    for (int redo = 0; redo < nredo; redo++) {
        for (int i = 0; i < nimages; i++) {
            integration.LoadDataToGPU();
            integration.RunRadialIntegration();
            integration.GetRadialIntegrationProfile(result);
        }
    }
    auto end_time = std::chrono::system_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);

    return elapsed.count() / (1000.0 * (double) nimages * nredo);
}

auto TestRadialIntegrationGPUWithoutCopyToDevice(const DiffractionExperiment &x,
                              const JFJochProtoBuf::DataProcessingSettings &settings,
                              int16_t* image, size_t nimages) {

    uint32_t nredo = 20;
    RadialIntegrationMapping mapping(x);

    std::vector<uint8_t> one_byte_mask(x.GetPixelsNum(), 1);
    GPUImageAnalysis integration(x.GetXPixelsNum(), x.GetYPixelsNum(), one_byte_mask, mapping);

    integration.SetInputBuffer(image);

    std::vector<float> result;

    integration.LoadDataToGPU();

    auto start_time = std::chrono::system_clock::now();
    for (int redo = 0; redo < nredo; redo++) {
        for (int i = 0; i < nimages; i++) {
            integration.RunRadialIntegration();
            integration.GetRadialIntegrationProfile(result);
        }
    }
    auto end_time = std::chrono::system_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);

    return elapsed.count() / (1000.0 * (double) nimages * nredo);
}


auto TestRadialIntegration(const DiffractionExperiment &experiment, const JFJochProtoBuf::DataProcessingSettings &settings,
                           int16_t* image, size_t nimages) {

    uint32_t nredo = 20;
    RadialIntegrationMapping mapping(experiment);
    RadialIntegration integration(mapping);
    std::vector<float> result;

    auto start_time = std::chrono::system_clock::now();
    for (int redo = 0; redo < nredo; redo++) {
        for (int i = 0; i < nimages; i++) {
            integration.ProcessOneImage(image + i * experiment.GetPixelsNum(), experiment.GetPixelsNum());
            //integration.GetResult(result);
        }
    }
    auto end_time = std::chrono::system_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);

    return elapsed.count() / (1000.0 * (double) nimages * nredo);
}

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

    if (argc != 2) {
        logger.Info("Usage: ./DataAnalysisPerfTest <JF4M hdf5 file>");
        exit(EXIT_FAILURE);
    }

    HDF5ReadOnlyFile data(argv[1] );

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

    HDF5DataSpace file_space(dataset);

    if (file_space.GetNumOfDimensions() != 3) {
        logger.Error("/entry/data/data must be 3D");
        exit(EXIT_FAILURE);
    }

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

    if ((file_space.GetDimensions()[1] == 2164) && (file_space.GetDimensions()[2] == 2068)) {
        logger.Info("JF4M with gaps detected (2068 x 2164)");
    } else {
        logger.Error("Unknown geometry - exiting");
        exit(EXIT_FAILURE);
    }
    uint64_t nimages = file_space.GetDimensions()[0];
    uint64_t npixel = file_space.GetDimensions()[1] * file_space.GetDimensions()[2];

    logger.Info("Number of images in the dataset: {}", nimages);

    x.Mode(DetectorMode::Conversion);

    std::vector<int16_t> image_conv ( nimages * npixel);

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

    logger.Info("Images loaded");

    x.BeamX_pxl(1090).BeamY_pxl(1136).DetectorDistance_mm(75).PhotonEnergy_keV(WVL_1A_IN_KEV);
    x.SetUnitCell(UnitCell{.a = 78.90f, .b = 78.90f, .c = 36.94, .alpha = 90, .beta = 90, .gamma = 90});

    JFJochProtoBuf::DataProcessingSettings settings;
    settings.set_signal_to_noise_threshold(2.5);
    settings.set_photon_count_threshold(5);
    settings.set_min_pix_per_spot(3);
    settings.set_max_pix_per_spot(200);
    settings.set_low_resolution_limit(80.0);
    settings.set_high_resolution_limit(2.0);
    settings.set_local_bkg_size(5);

    std::vector<uint8_t> one_byte_mask(x.GetPixelsNum(), 1);
    RadialIntegrationMapping mapping(x);

    logger.Info("COLSPOT NBX=NBY=5 (GPU)");
    if (GPUImageAnalysis::GPUPresent()) {
        GPUImageAnalysis local_peakfinder_gpu(x.GetXPixelsNum(), x.GetYPixelsNum(), one_byte_mask);
        TestSpotFinder(x, settings, local_peakfinder_gpu,image_conv.data(), nimages);
    }
    logger.Info("COLSPOT NBX=NBY=5 (GPU/no copy to device)");
    if (GPUImageAnalysis::GPUPresent()) {
        GPUImageAnalysis local_peakfinder_gpu(x.GetXPixelsNum(), x.GetYPixelsNum(), one_byte_mask);
        TestSpotFinderWithoutCopyToDevice(x, settings, local_peakfinder_gpu,image_conv.data(), nimages);
    }

    settings.set_local_bkg_size(3);
    logger.Info("COLSPOT NBX=NBY=3 (GPU)");
    if (GPUImageAnalysis::GPUPresent()) {
        GPUImageAnalysis local_peakfinder_gpu(x.GetXPixelsNum(), x.GetYPixelsNum(), one_byte_mask);
        TestSpotFinder(x, settings, local_peakfinder_gpu,image_conv.data(), nimages);
    }

    logger.Info("{:20s} {:8.1f} ms/image", "Radial int. (GPU)", TestRadialIntegrationGPU(x, settings,
                                                                                                image_conv.data(), nimages));
    logger.Info("{:20s} {:8.1f} ms/image", "Radial int. (GPU/nocopy)", TestRadialIntegrationGPUWithoutCopyToDevice(x, settings,
                                                                                         image_conv.data(), nimages));

    logger.Info("{:20s} {:8.1f} ms/image", "Radial int. (CPU)", TestRadialIntegration(x, settings,
                                                                                        image_conv.data(), nimages));

    TestIndexing();

    logger.Info("Full package");
    if (GPUImageAnalysis::GPUPresent()) {
        GPUImageAnalysis local_peakfinder_gpu(x.GetXPixelsNum(), x.GetYPixelsNum(), one_byte_mask, mapping);
        TestAll(x, settings, local_peakfinder_gpu,image_conv.data(), nimages);
        TestAllWithROI(x, settings, local_peakfinder_gpu,image_conv.data(), nimages);
    }
}