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

#include <random>
#include <iostream>
#include <iomanip>
#include "../jungfrau/JFPedestalCalc.h"
#include "../common/Logger.h"
#include "../jungfrau/JFCalibration.h"
#include "../jungfrau/JFConversionFloatingPoint.h"
#include "../jungfrau/JFConversionFixedPoint.h"
#include "../tests/FPGAUnitTest.h"

void test_pedestal(Logger &logger) {
    size_t nframes = 5000;
    DiffractionExperiment x(1,{1});

    std::vector<uint16_t> data(nframes * RAW_MODULE_SIZE);

    x.Mode(DetectorMode::Conversion);
    double mean = 1000.0;
    double stddev = 50.0;

    // Predictable random number generator
    std::mt19937 g1(5423);
    std::normal_distribution<double> distribution(mean, stddev);

    for (size_t i = 0; i < nframes * RAW_MODULE_SIZE; i++) {
        double number = distribution(g1);
        if (number < 20) number = 20;
        if (number > 16300) number = 16300;
        data[i] = number;
    }

    JFPedestalCalc calc_cpu(x);
    auto start_time = std::chrono::system_clock::now();
    for (int i = 0; i < nframes; i++)
        calc_cpu.AnalyzeImage(data.data() + i * RAW_MODULE_SIZE);

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

    logger.Info("CPU pedestal performance: {:5d} us/module {:5.2f} GB/s", std::lround(elapsed.count() / ((double) nframes)),
                nframes * RAW_MODULE_SIZE * sizeof(uint16_t) * 1000 * 1000/ ((double) elapsed.count() * 1024 * 1024 * 1024));
}

template <class T> void test_conversion(Logger &logger) {
    size_t nframes = 128;
    int64_t nmodules = 8;
    int64_t ntries = 8;

    DiffractionExperiment x(1,{nmodules});

    std::vector<uint16_t> input(nframes * nmodules * RAW_MODULE_SIZE);
    std::vector<int16_t> output(nframes * nmodules * RAW_MODULE_SIZE);

    for (int i = 0; i < nmodules * nframes; i++) {
        std::string image_path = "../../tests/test_data/mod5_raw" + std::to_string(i % 20) + ".bin";
        LoadBinaryFile(image_path, input.data() + i * RAW_MODULE_SIZE, RAW_MODULE_SIZE);
    }

    std::vector<T> v(nmodules);

    JFModuleGainCalibration gain_calib = GainCalibrationFromTestFile();

    for (int m = 0; m < nmodules; m++) {
        JFModulePedestal pedestal_g0;
        JFModulePedestal pedestal_g1;
        JFModulePedestal pedestal_g2;

        for (int i = 0; i < RAW_MODULE_SIZE; i++) {
            pedestal_g0.GetPedestal()[i] = 3000 + i % 50 + m * 135;
            pedestal_g1.GetPedestal()[i] = 15000 + i % 50 - m * 135;
            pedestal_g2.GetPedestal()[i] = 14000 + i % 50 - m * 135;
        }
        v[m].Setup(gain_calib, pedestal_g0, pedestal_g1, pedestal_g2, 12.4);
    }

    x.Mode(DetectorMode::Conversion);

    logger.Info("JF FP conversion input prepared");
    auto start_time = std::chrono::system_clock::now();
    for (int z = 0; z < ntries; z++) {
        for (int i = 0; i < nframes; i++) {
            for (int m = 0; m < nmodules; m++) {
                v[m].Convert(output.data() + (i * nmodules + m) * RAW_MODULE_SIZE,
                             input.data() + (i * nmodules + m) * RAW_MODULE_SIZE);
            }
        }
    }
    auto end_time = std::chrono::system_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);

    logger.Info("Conversion performance: {:5d} us/module {:5.2f} GB/s", std::lround(elapsed.count() / ((double) (ntries * nframes * nmodules))),
                ntries * nframes * nmodules * RAW_MODULE_SIZE * sizeof(uint16_t) * 1000 * 1000/ ((double) elapsed.count() * 1024 * 1024 * 1024));
}


template <class T> void test_conversion_with_geom(Logger &logger) {
    size_t nframes = 128;
    int64_t nmodules = 8;
    int64_t ntries = 8;

    DiffractionExperiment x(1,{nmodules});

    std::vector<uint16_t> input(nframes * nmodules * RAW_MODULE_SIZE);
    std::vector<int16_t> output(nframes * nmodules * CONVERTED_MODULE_SIZE);

    for (int i = 0; i < nmodules * nframes; i++) {
        std::string image_path = "../../tests/test_data/mod5_raw" + std::to_string(i % 20) + ".bin";
        LoadBinaryFile(image_path, input.data() + i * RAW_MODULE_SIZE, RAW_MODULE_SIZE);
    }

    std::vector<T> v(nmodules);

    JFModuleGainCalibration gain_calib = GainCalibrationFromTestFile();

    for (int m = 0; m < nmodules; m++) {
        JFModulePedestal pedestal_g0;
        JFModulePedestal pedestal_g1;
        JFModulePedestal pedestal_g2;

        for (int i = 0; i < RAW_MODULE_SIZE; i++) {
            pedestal_g0.GetPedestal()[i] = 3000 + i % 50 + m * 135;
            pedestal_g1.GetPedestal()[i] = 15000 + i % 50 - m * 135;
            pedestal_g2.GetPedestal()[i] = 14000 + i % 50 - m * 135;
        }
        v[m].Setup(gain_calib, pedestal_g0, pedestal_g1, pedestal_g2, 12.4);
    }

    x.Mode(DetectorMode::Conversion);

    logger.Info("JF FP conversion input prepared");
    auto start_time = std::chrono::system_clock::now();
    for (int z = 0; z < ntries; z++) {
        for (int i = 0; i < nframes; i++) {
            for (int m = 0; m < nmodules; m++) {
                v[m].ConvertAdjustGeom(output.data() + (i * nmodules + m) * CONVERTED_MODULE_SIZE,
                             input.data() + (i * nmodules + m) * RAW_MODULE_SIZE, 1);
            }
        }
    }
    auto end_time = std::chrono::system_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);

    logger.Info("Conversion performance: {:5d} us/module {:5.2f} GB/s", std::lround(elapsed.count() / ((double) (ntries * nframes * nmodules))),
                ntries * nframes * nmodules * RAW_MODULE_SIZE * sizeof(uint16_t) * 1000 * 1000/ ((double) elapsed.count() * 1024 * 1024 * 1024));
}

int main () {
    Logger logger("JFCalibrationPerfTest");
    test_pedestal(logger);
    logger.Info("Floating point conversion");
    test_conversion<JFConversionFloatingPoint>(logger);
    logger.Info("Fixed point conversion");
    test_conversion<JFConversionFixedPoint>(logger);
    logger.Info("Fixed point conversion (with geom)");
    test_conversion_with_geom<JFConversionFixedPoint>(logger);
}