// Copyright (2019-2023) Paul Scherrer Institute

#include <iostream>

#include "PCIExpressDevice.h"
#include "MockAcquisitionDevice.h"
#include "JFJochReceiverTest.h"
#include "../tests/FPGAUnitTest.h"

void print_usage(Logger &logger) {
    logger.Info("Usage ./jfjoch_action_test {<options>} <path to repository>");
    logger.Info("Options:");
    logger.Info("   -C      conversion on CPU");
    logger.Info("   -M      use mock device");
    logger.Info("   -R      raw");
    logger.Info("   -v      verbose");
    logger.Info("   -H      mock aq. dev. with HBM (DL380 with Intel MAX only)");
    logger.Info("   -D      mock aq. dev. with DDR (2 NUMA node machines only)");
    logger.Info("   -s<num> number of data streams (acquisition devices)");
    logger.Info("   -m<num> number of modules");
    logger.Info("   -i<num> number of images");
    logger.Info("   -p<num> data processing period");
    logger.Info("   -N<num> number of image processing threads");
    logger.Info("   -P<txt> NUMA Policy (none|n2g2|n8g4|n8g4_hbm), none is default");
}

int main(int argc, char **argv) {
    Logger logger("ActionTest");
    logger.Verbose(true);

    constexpr uint64_t clock_MHz = 200;
    uint16_t nstreams = 1;
    uint16_t nmodules = 1;
    size_t nimages = 2;
    uint64_t processing_period = 20;
    uint16_t nthreads = 64;
    bool conversion_on_cpu = false;
    bool use_mock_device = false;
    bool verbose = false;
    std::string numa_policy_name;
    bool use_hbm_for_aq_dev = false;
    bool use_ddr_for_aq_dev = false;
    bool raw_data = false;

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

    int opt;
    while ((opt = getopt(argc, argv, "s:i:m:p:N:P:CMvHDR")) != -1) {
        switch (opt) {
            case 'C':
                conversion_on_cpu = true;
                break;
            case 'M':
                use_mock_device = true;
                break;
            case 'i':
                nimages = atol(optarg);
                break;
            case 'm':
                nmodules = atol(optarg);
                break;
            case 's':
                nstreams = atol(optarg);
                break;
            case 'p':
                processing_period = atol(optarg);
                break;
            case 'N':
                nthreads = atol(optarg);
                break;
            case 'v':
                verbose = true;
                break;
            case 'P':
                numa_policy_name = std::string(optarg);
                break;
            case 'H':
                use_hbm_for_aq_dev = true;
                break;
            case 'D':
                use_ddr_for_aq_dev = true;
                break;
            case 'R':
                raw_data = true;
                break;
            default: /* '?' */
                print_usage(logger);
                exit(EXIT_FAILURE);
        }
    }

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

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

    if (raw_data)
        x.Mode(DetectorMode::Raw);
    else
        x.Mode(DetectorMode::Conversion);

    x.ImagesPerTrigger(nimages).PedestalG0Frames(0).UseInternalPacketGenerator(true).PhotonEnergy_keV(12.4).NumTriggers(1);
    x.SpotFindingPeriod(std::chrono::milliseconds(processing_period)).MaskModuleEdges(false).MaskChipEdges(false).ConversionOnCPU(conversion_on_cpu);
    x.Compression(JFJochProtoBuf::BSHUF_LZ4).DataStreams(nstreams);

    logger.Info("Data streams {} Total modules {} Total images {} Threads {}", nstreams, nmodules, nimages, nthreads);

    std::vector<std::string> dev_name = {
            "/dev/jfjoch0",
            "/dev/jfjoch2",
            "/dev/jfjoch1",
            "/dev/jfjoch3"
    };

    logger.Verbose(verbose);

    std::vector<std::unique_ptr<PCIExpressDevice>> pcie_devices;
    std::vector<std::unique_ptr<MockAcquisitionDevice>> mock_devices;
    std::vector<AcquisitionDevice *> aq_devices;

    std::string image_path = std::string(argv[optind]) + "/tests/test_data/mod5_raw0.bin";
    std::vector<uint16_t> input(RAW_MODULE_SIZE, 0);
    LoadBinaryFile(image_path, input.data(), RAW_MODULE_SIZE);

    if (use_mock_device) {
        if (nmodules > 1) {
            logger.Warning("Conversion results might be wrong with more than 1 module per stream");
        }

        for (int i = 0; i < nstreams; i++) {
            int16_t numa_node = -1;

            if (use_hbm_for_aq_dev)
                numa_node = 2 + (i % 2);
            else if (use_ddr_for_aq_dev)
                numa_node = i % 2;

            if (numa_node != -1)
                logger.Info("Pinning stream {} to NUMA node {}", i, numa_node);

            mock_devices.push_back(std::make_unique<MockAcquisitionDevice>(i, 1024, numa_node));
            mock_devices[i]->SetCustomInternalGeneratorFrame(input);
            mock_devices[i]->EnableLogging(&logger);
            aq_devices.push_back(mock_devices[i].get());
        }

    } else {
        if (nstreams > dev_name.size()) {
            logger.Error("Only {} data streams allowed on this platform", dev_name.size());
            exit(EXIT_FAILURE);
        }

        for (int i = 0; i < nstreams; i++) {
            pcie_devices.push_back(std::make_unique<PCIExpressDevice>(dev_name[i], i));
            pcie_devices[i]->SetInternalGeneratorFrame(input);
            pcie_devices[i]->EnableLogging(&logger);
            pcie_devices[i]->SetDefaultMAC();
            pcie_devices[i]->SetIPv4Address((i << 24) + 0x010a0a0a);
            aq_devices.push_back(pcie_devices[i].get());
        }
    }

    volatile bool done = false;
    JFJochProtoBuf::ReceiverOutput output;
    bool ret;
    std::thread run_thread([&] {
        try {
            ret = JFJochReceiverTest(output, logger, aq_devices, x, nthreads, false, verbose, nullptr, numa_policy_name);
        } catch (std::exception &e) {
            logger.Error(e.what());
            ret = false;
        }
        done = true;
    });

    if (!use_mock_device) {
        while (!done) {
            for (int i = 0; i < nstreams; i++) {
                auto status = pcie_devices[i]->GetStatus();
                logger.Info("Device {}:  Head packet: {:8d}  Power: {:5.1f} W  FPGA Temp: {:d} degC  HBM Temp: {:d}/{:d} degC  Stalls: {:15d}",
                            i, status.slowest_head(), status.current_edge_12v_a() * status.voltage_edge_12v_v() +
                                                      status.current_edge_3p3v_a() * status.voltage_edge_3p3v_v(),
                            std::lround(status.fpga_temp_degc()), status.hbm_temp_0_degc(), status.hbm_temp_1_degc(),
                            status.stalls_hbm());
            }
            std::this_thread::sleep_for(std::chrono::seconds(1));
        }
    }

    run_thread.join();

    double receiving_time = static_cast<double>(output.end_time_ms() - output.start_time_ms())/1000.0;

    logger.Info("Efficiency: {:.2f}%", output.efficiency() * 100.f);
    logger.Info("Max delay: {}",output.max_receive_delay());
    logger.Info("Compression factor: {}x", output.compressed_ratio());
    logger.Info("Receiving time: {} s", receiving_time);
    logger.Info("Frame rate: {} Hz", static_cast<double>(nimages)/receiving_time);
    logger.Info("Total throughput: {:.2f} GB/s",
            static_cast<double>(nimages*x.GetModulesNum()*RAW_MODULE_SIZE*sizeof(uint16_t)) / (receiving_time * 1e9));

    if (!use_mock_device) {
        logger.Info("");
        for (int i = 0; i < nstreams; i++) {
            auto stalls_hbm = output.device_statistics(i).fpga_status().stalls_hbm();
            auto stalls_host = output.device_statistics(i).fpga_status().stalls_host();

            uint64_t throughput_MBs = nimages * x.GetModulesNum(i) * RAW_MODULE_SIZE * sizeof(uint16_t) * clock_MHz /
                                      (nimages * x.GetModulesNum(i) * 128 * 128 + stalls_hbm);
            double performance = static_cast<double>(throughput_MBs) / 1000;

            logger.Info("Device {}:  stalls HBM: {}  stalls host: {}   est. performance: {:.2f} GB/s", i, stalls_hbm,
                        stalls_host, performance);

            for (const auto &iter: output.device_statistics(i).fpga_status().host_writer_err())
                logger.Error("Device {}: FPGA host writer error {}", i, iter);

            if (output.device_statistics(i).fpga_status().mailbox_err_reg() != 0)
                logger.Error("Device {}: Mailbox error {:x}", i,
                             output.device_statistics(i).fpga_status().mailbox_err_reg());
        }
    }
    if (ret) {
        logger.Info("");
        logger.Info("Test properly executed! (check stall values manually)");
        exit(EXIT_SUCCESS);
    } else {
        logger.Info("Test finished with errors! (check stall values manually)");
        exit(EXIT_FAILURE);
    }
}