//
// Copyright 2011 Ettus Research LLC
// Copyright 2018 Ettus Research, a National Instruments Company
//
// SPDX-License-Identifier: GPL-3.0-or-later
//

#include <uhd/usrp/multi_usrp.hpp>
#include <uhd/utils/safe_main.hpp>
#include <uhd/utils/thread.hpp>
#include <boost/format.hpp>
#include <boost/program_options.hpp>
#include <complex>
#include <iostream>

namespace po = boost::program_options;

int UHD_SAFE_MAIN(int argc, char* argv[])
{
    const std::string program_doc =
        "This example program takes a round-trip time (RTT) value as input and\n"
        "probes whether the actual system round-trip time from USRP to host and\n"
        "back to USRP is smaller than this value. For precise timing, the USRP\n"
        "device internal clock is used. The reference point for all timing\n"
        "measurements is the timestamp generated by the Radio (RFNoC) block in\n"
        "the USRP FPGA, which tags each received sample burst with the precise\n"
        "hardware time when the first sample enters the digital receive path.\n"
        "\n"
        "Procedure:\n"
        "  - The test requests reception of a packet. When it is received at the\n"
        "    host, the USRP-provided receive timestamp t is read.\n"
        "  - Next, a transmission is scheduled for USRP time t + rtt.\n"
        "  - If the transmission fails (arrives too late), it indicates that the\n"
        "    actual round-trip latency is greater than the RTT being tested.\n"
        "\n"
        "Results:\n"
        "  - The test is repeated multiple times, and statistics are collected on\n"
        "    ACKed transmissions (successful sends), late packets (time errors),\n"
        "    and underruns.\n"
        "  - If most transmissions are ACKed and few or no late packets are\n"
        "    reported, the specified RTT is sufficient and the system round-trip\n"
        "    time is less than the tested value.\n"
        "  - If many late packets occur, it indicates that the actual system RTT\n"
        "    is greater than the specified value and the system cannot reliably\n"
        "    meet that timing constraint.\n"
        "\n"
        "Usage example:\n"
        "  Test with RTT value 1 ms and enable verbose output:\n"
        "    latency_test --args=\"addr=192.168.10.2\" --rtt=0.001 --verbose\n";
    // variables to be set by po
    std::string args;
    size_t nsamps;
    double rate;
    double rtt;
    size_t nruns;

    // setup the program options
    po::options_description desc("Allowed options");
    // clang-format off
    desc.add_options()
        ("help", "Show this help message and exit.")
        ("args",   po::value<std::string>(&args)->default_value(""), "Single USRP device selection and "
            "configuration arguments."
            "\nSpecify key-value pairs (e.g., addr, serial, type, master_clock_rate) separated by commas."
            "\nSee the UHD manual for model-specific options."
            "\nExamples:"
            "\n  --args \"addr=192.168.10.2\""
            "\n  --args \"addr=192.168.10.2,master_clock_rate=200e6\""
            "\nIf not specified, UHD connects to the first available device.")
        ("nsamps", po::value<size_t>(&nsamps)->default_value(100),   "Number of samples per packet for TX/RX "
            "per run.")
        ("nruns",  po::value<size_t>(&nruns)->default_value(1000),   "Number of test runs to perform.")
        ("rtt",    po::value<double>(&rtt)->default_value(0.001),    "The user-specified round-trip time (in "
            "seconds) for which the benchmark will test whether the system can reliably complete a USRP-to-host-to-USRP "
            "cycle.")
        ("rate",   po::value<double>(&rate)->default_value(100e6/4), "Sample rate (in samples/second) used "
            "for both TX and RX. Determines the speed of packet transfers and is used to calculate RX transfer time and "
            "schedule the TX burst."
            "\nWhen --from-eob is set, the RX transfer time (nsamps/rate) is added to the TX scheduling, making the test "
            "independent of RX transfer duration. If an unsupported rate is requested, the closest available rate is "
            "used.")
        ("from-eob", "Exclude the time required to transfer RX samples from the USRP to the host from the RTT "
            "measurement, so the test only measures the processing and scheduling delay (removes dependence on nsamps and "
            "rate).")
        ("verbose", "Enable inner-loop verbose reporting.")
    ;
    // clang-format on
    po::variables_map vm;
    po::store(po::parse_command_line(argc, argv, desc), vm);
    po::notify(vm);

    // print the help message
    if (vm.count("help")) {
        std::cout << program_doc << std::endl;
        std::cout << desc << std::endl;
        return ~0;
        return EXIT_SUCCESS;
    }

    bool verbose  = vm.count("verbose") != 0;
    bool from_eob = vm.count("from-eob") != 0;

    // create a usrp device
    std::cout << std::endl;
    // std::cout << boost::format("Creating the usrp device with: %s...") % args <<
    // std::endl;
    uhd::usrp::multi_usrp::sptr usrp = uhd::usrp::multi_usrp::make(args);
    // std::cout << boost::format("Using Device: %s") % usrp->get_pp_string() <<
    // std::endl;

    usrp->set_time_now(uhd::time_spec_t(0.0));

    // set the tx sample rate
    usrp->set_tx_rate(rate);
    std::cout << boost::format("Actual TX Rate: %f Msps...") % (usrp->get_tx_rate() / 1e6)
              << std::endl;

    // set the rx sample rate
    usrp->set_rx_rate(rate);
    double actual_rx_rate = usrp->get_rx_rate();
    std::cout << boost::format("Actual RX Rate: %f Msps...") % (actual_rx_rate / 1e6)
              << std::endl;

    double rx_time = nsamps / actual_rx_rate;
    if (from_eob) {
        std::cout << boost::format("Will add %f seconds to timespec for RX samples...")
                         % (rx_time)
                  << std::endl;
    }

    // allocate a buffer to use
    std::vector<std::complex<float>> buffer(nsamps);

    // create RX and TX streamers
    uhd::stream_args_t stream_args("fc32"); // complex floats
    uhd::rx_streamer::sptr rx_stream = usrp->get_rx_stream(stream_args);
    uhd::tx_streamer::sptr tx_stream = usrp->get_tx_stream(stream_args);

    // initialize result counts
    int time_error = 0;
    int ack        = 0;
    int underflow  = 0;
    int other      = 0;

    for (size_t nrun = 0; nrun < nruns; nrun++) {
        /***************************************************************
         * Issue a stream command some time in the near future
         **************************************************************/
        uhd::stream_cmd_t stream_cmd(uhd::stream_cmd_t::STREAM_MODE_NUM_SAMPS_AND_DONE);
        stream_cmd.num_samps  = buffer.size();
        stream_cmd.stream_now = false;
        stream_cmd.time_spec  = usrp->get_time_now() + uhd::time_spec_t(0.01);
        rx_stream->issue_stream_cmd(stream_cmd);

        /***************************************************************
         * Receive the requested packet
         **************************************************************/
        uhd::rx_metadata_t rx_md;
        size_t num_rx_samps = rx_stream->recv(&buffer.front(), buffer.size(), rx_md);

        if (verbose) {
            std::cout << boost::format(
                             "Run %d: Got packet: %u samples, %u full secs, %f frac secs")
                             % nrun % num_rx_samps % rx_md.time_spec.get_full_secs()
                             % rx_md.time_spec.get_frac_secs()
                      << std::endl;
        } else {
            std::cout << "." << std::flush;
        }

        /***************************************************************
         * Transmit a packet with delta time after received packet
         **************************************************************/
        uhd::tx_metadata_t tx_md;
        tx_md.start_of_burst = true;
        tx_md.end_of_burst   = true;
        tx_md.has_time_spec  = true;
        tx_md.time_spec      = rx_md.time_spec + uhd::time_spec_t(rtt);
        if (from_eob) {
            tx_md.time_spec += uhd::time_spec_t(rx_time);
        }
        size_t num_tx_samps = tx_stream->send(&buffer.front(), buffer.size(), tx_md);
        if (verbose) {
            std::cout << boost::format("Sent %d samples") % num_tx_samps << std::endl;
        }

        /***************************************************************
         * Check the async messages for result
         **************************************************************/
        uhd::async_metadata_t async_md;
        if (not tx_stream->recv_async_msg(async_md)) {
            std::cout << boost::format("failed:\n    Async message recv timed out.\n")
                      << std::endl;
            continue;
        }
        switch (async_md.event_code) {
            case uhd::async_metadata_t::EVENT_CODE_TIME_ERROR:
                time_error++;
                break;

            case uhd::async_metadata_t::EVENT_CODE_BURST_ACK:
                ack++;
                break;

            case uhd::async_metadata_t::EVENT_CODE_UNDERFLOW:
                underflow++;
                break;

            default:
                std::cerr << boost::format(
                                 "failed:\n    Got unexpected event code 0x%x.\n")
                                 % async_md.event_code
                          << std::endl;
                other++;
                break;
        }
    }

    while (true) {
        uhd::async_metadata_t async_md;
        if (not tx_stream->recv_async_msg(async_md)) {
            break;
        }
        switch (async_md.event_code) {
            case uhd::async_metadata_t::EVENT_CODE_TIME_ERROR:
                time_error++;
                break;

            case uhd::async_metadata_t::EVENT_CODE_BURST_ACK:
                ack++;
                break;

            case uhd::async_metadata_t::EVENT_CODE_UNDERFLOW:
                underflow++;
                break;

            default:
                std::cerr << boost::format(
                                 "failed:\n    Got unexpected event code 0x%x.\n")
                                 % async_md.event_code
                          << std::endl;
                other++;
                break;
        }
    }
    if (!verbose) {
        std::cout << std::endl;
    }

    /***************************************************************
     * Print the summary
     **************************************************************/
    std::cout << "Summary\n"
              << "================\n"
              << "Number of runs:   " << nruns << std::endl
              << "RTT value tested: " << (rtt * 1e3) << " ms" << std::endl
              << "ACKs received:    " << ack << "/" << nruns << std::endl
              << "Underruns:        " << underflow << std::endl
              << "Late packets:     " << time_error << std::endl
              << "Other errors:     " << other << std::endl
              << std::endl;
    return EXIT_SUCCESS;
}
