// How to build:
//  $ export UHD_INCLUDE_PATH=/path/to/include
//  $ export UHD_LIBRARY_PATH=/path/to/lib
//  $ g++ -std=c++17 twinrx_issue.cpp -o twinrx_issue -I$UHD_INCLUDE_PATH -L$UHD_LIBRARY_PATH -lboost_program_options -lpthread -luhd

#include <uhd/usrp/multi_usrp.hpp>
#include <uhd/utils/safe_main.hpp>
#include <uhd/utils/thread.hpp>
#include <boost/program_options.hpp>
#include <boost/thread.hpp>
#include <chrono>
#include <fstream>
#include <csignal>
#include <cerrno>
#include <poll.h>

namespace po = boost::program_options;

typedef std::vector<std::complex<float>> recv_buff_t;

// Global objects
static uhd::usrp::multi_usrp::sptr usrp;
static recv_buff_t buffer;
static recv_buff_t writing_buffer ;
static std::vector<double> rf_freqs;

static uhd::stream_cmd_t stream_start(uhd::stream_cmd_t::STREAM_MODE_START_CONTINUOUS);
static uhd::stream_cmd_t stream_stop(uhd::stream_cmd_t::STREAM_MODE_STOP_CONTINUOUS);

bool done ;
bool save_samples ;

void sig_int_handler(int) {
    done = true ;
}

static size_t recv_chan ;
static size_t scan_chan ;

void check_user_input() {
    struct pollfd pfd = { STDIN_FILENO, POLLIN, 0 } ;
    int ret = 0 ;
    while(!done) {
        // Poll for user hitting enter to stop gracefully
        ret = poll(&pfd, 1, 100) ;
        if( ret == 1 ) {
            std::cin.ignore() ;
            std::cout << "Saving Samples!" << std::endl ;
            save_samples = true ;
        } else if( ret == -1 ) {
            std::cout << "Error: " << strerror(errno) << std::endl ;
        }
    }
}

static void write_buffer_to_file(const std::string &path) {
    std::cout << str(boost::format("Writing file %s... ") % path) << std::flush ;
    std::ofstream ofile(path.c_str(), std::ios::binary) ;
    ofile.write((char*)writing_buffer.data(), sizeof(std::complex<float>) * writing_buffer.size()) ;
    std::cout << "done!" << std::endl ;
}

// This is a helper function for receiving samples from the USRP
void twinrx_recv()
{
    buffer = recv_buff_t(100e6/100) ;
    writing_buffer = recv_buff_t(buffer.size()) ;
    uhd::rx_metadata_t md;
    size_t file_number = 0 ;

    // Get an rx_streamer from the device
    uhd::stream_args_t stream_args("fc32", "sc16");
    stream_args.channels.push_back(recv_chan);
    auto rx_stream = usrp->get_rx_stream(stream_args);
    std::vector<std::thread> threads ;

    rx_stream->issue_stream_cmd(stream_start) ;

    // Repeatedly retrieve samples until the entire acquisition is received
    while (!done) {
        size_t num_recvd = rx_stream->recv(buffer.data(), buffer.size(), md, 1.0);

        if (md.error_code != uhd::rx_metadata_t::ERROR_CODE_NONE) {
            std::cout << md.strerror() << std::endl;
            continue ;
        }

        if( save_samples ) {
            std::string fname = str(boost::format("buffer-%03d.dat") % file_number)  ;
            // Copy the buffer to the writing_buffer
            writing_buffer = buffer ;
            // Non-blocking writing
            threads.push_back(std::thread(write_buffer_to_file,fname)) ;
            save_samples = false ;
            file_number++ ;
        }
    }

    std::cout << std::endl << "Stopping the stream..." ;
    // Stop the stream
    rx_stream->issue_stream_cmd(stream_stop) ;

    // Read all the samples left in the pipeline
    do {
        size_t num_recv = rx_stream->recv(buffer.data(), buffer.size(), md, 1.0) ;
    } while (md.error_code == uhd::rx_metadata_t::ERROR_CODE_NONE) ;
    // Rejoin all the file saving threads
    std::cout << "done!" << std::endl ;
    std::cout << "Cleaning up ... " ;
    for(auto &t : threads) {
        if( t.joinable() ) t.join() ;
    }
    std::cout << "done!" << std::endl ;
}

int UHD_SAFE_MAIN(int argc, char* argv[])
{
    // Program options
    std::string args, ant;
    double rate ;
    double start_freq, end_freq;
    done = false ;

    // Set up the program options
    po::options_description desc("Allowed options");

    // clang-format off
    desc.add_options()
            ("help", "Print this help message")
            ("args", po::value<std::string>(&args)->default_value(""), "UHD device args")
            ("chan", po::value<size_t>(&scan_chan)->default_value(1), "Scan Channel")
            ("ant", po::value<std::string>(&ant)->default_value("RX1"), "Antenna")
            ("start-freq", po::value<double>(&start_freq), "Start frequency (defaults to lowest valid frequency)")
            ("end-freq", po::value<double>(&end_freq), "End frequency (defaults to highest valid frequency)")
            ("rate", po::value<double>(&rate)->default_value(100e6), "Incoming sample rate")
            ;

    // clang-format on
    po::variables_map vm;
    po::store(po::parse_command_line(argc, argv, desc), vm);
    po::notify(vm);

    // Receive channel is opposite of scanning channel
    if( scan_chan == 0 ) recv_chan = 1 ;
    else                 recv_chan = 0 ;

    if (vm.count("help")) {
        std::cout << "TwinRX Noise Issue - " << desc << std::endl;
        return EXIT_SUCCESS;
    }

    // Create a USRP device
    std::cout << boost::format("Creating the USRP device with args: \"%s\"...\n")
                     % args;
    usrp = uhd::usrp::multi_usrp::make(args);

    // Make sure the USRP is an X3xx with a TwinRX
    uhd::dict<std::string, std::string> info = usrp->get_usrp_rx_info(scan_chan);
    if (info.get("mboard_id").find("X3") == std::string::npos) {
        throw uhd::runtime_error(
            "This example can only be used with an X-Series motherboard.");
    }

    usrp->set_rx_subdev_spec(uhd::usrp::subdev_spec_t("B:0 B:1")) ;
    usrp->set_rx_antenna(ant, scan_chan) ;
    std::cout << boost::format("Setting sample rate to: %d\n") % rate;
    usrp->set_rx_rate(rate);
    std::cout << boost::format("Actual sample rate:     %d\n") % usrp->get_rx_rate();

    std::cout << "info rx_id: " << info.get("rx_id") << std::endl ;
    if (info.get("rx_id").find("TwinRX") == std::string::npos) {
        throw uhd::runtime_error(
            "This example can only be used with a TwinRX daughterboard.");
    }

    // Validate frequency range
    uhd::freq_range_t rx_freq_range = usrp->get_rx_freq_range(scan_chan);
    if (!vm.count("start-freq")) {
        start_freq = rx_freq_range.start();
    }
    if (!vm.count("end-freq")) {
        end_freq = rx_freq_range.stop();
    }
    if (start_freq < rx_freq_range.start() or end_freq > rx_freq_range.stop()) {
        throw uhd::runtime_error(
            (boost::format("Start and stop frequencies must be between %d and %d MHz")
                % (rx_freq_range.start() / 1e6) % (rx_freq_range.stop() / 1e6))
                .str());
    }
    if (start_freq > end_freq) {
        throw uhd::runtime_error("Start frequency must be less than end frequency.");
    }
    if ((end_freq - start_freq) > 0 and (end_freq - start_freq) < rate) {
        throw uhd::runtime_error("The sample rate must be less than the range between "
                                 "the start and end frequencies.");
    }

    for (double rx_freq = start_freq ; rx_freq < end_freq ; rx_freq += rate) {
        rf_freqs.push_back(rx_freq);
    }
    std::cout << boost::format("Total Hops: %d\n") % rf_freqs.size();

    // Hop frequencies and acquire bursts of samples at each until done sweeping
    std::cout << "Scanning..." << std::endl ;
    auto start_time = std::chrono::system_clock::now() ;
    auto stop_time = std::chrono::system_clock::now() ;

    std::signal(SIGINT, &sig_int_handler) ;
    std::cout << "Press Ctrl+C to stop, ENTER to save samples ..." << std::endl ;

    auto user_input = std::thread(check_user_input) ;
    auto rx_stream = std::thread(twinrx_recv) ;

    while (!done) {
        auto locked = usrp->get_rx_sensor("lo_locked", scan_chan);

        start_time = std::chrono::system_clock::now() ;
        for (size_t i = 0; i < rf_freqs.size(); i++) {
            // Set the frequency
            usrp->set_rx_freq(rf_freqs[i], scan_chan) ;

            // Set maximum gain
            usrp->set_rx_gain(93, scan_chan) ;

            // Wait of the LO to be locked
            do {
                locked = usrp->get_rx_sensor("lo_locked", scan_chan) ;
            } while( !locked.to_bool() ) ;
        }
        stop_time = std::chrono::system_clock::now() ;
        auto us = std::chrono::duration_cast<std::chrono::microseconds>(stop_time - start_time).count() ;
        std::cout << boost::format("Total time: %7d microseconds  (%7.2f usec/hop)\r") % static_cast<float>(us) % (static_cast<float>(us) / rf_freqs.size()) ;
        std::cout << std::flush ;
    }

    user_input.join() ;
    rx_stream.join() ;

    std::cout << "Done!" << std::endl;

    usrp.reset();
    return EXIT_SUCCESS;
}