// Copyright (2019-2024) Paul Scherrer Institute

#include <iostream>

#include "NeuralNetResPredictor.h"
#include "../writer/HDF5Objects.h"

int main(int argc, char **argv) {
    if (argc != 2) {
        std::cerr << "Usage ./resonet_test <.pt file with traced model>" << std::endl;
        exit(EXIT_FAILURE);
    }

    RegisterHDF5Filter();

    DiffractionExperiment experiment(DetectorGeometry(8, 2, 8, 36));
    experiment.DetectorDistance_mm(75).PhotonEnergy_keV(12.4).BeamY_pxl(1136).BeamX_pxl(1090);

    NeuralNetResPredictor predictor(argv[1]);

    HDF5ReadOnlyFile data("../../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);

    std::cout << "Max pooling " << predictor.GetMaxPoolFactor(experiment) << std::endl;

    float x = 0;
    size_t niterations = 25;

    auto start_time = std::chrono::system_clock::now();
    for (int i = 0; i < niterations; i++)
        x += predictor.Inference(experiment, image_conv.data());
    auto end_time = std::chrono::system_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);
    int64_t frequency_Hz = (niterations * 1e6) / (double) (elapsed.count());

    std::cout << ((double) elapsed.count()) / (1e3 * niterations) << " ms" << std::endl;
    std::cout << frequency_Hz << " Hz" << std::endl;
    std::cout << "Resolution " << predictor.Inference(experiment, image_conv.data()) << std::endl;
}