#!/usr/bin/env python3
#
# Copyright 2020 Ettus Research, a National Instruments Brand
#
# SPDX-License-Identifier: GPL-3.0-or-later
#
"""Utility to run power calibrations with USRPs.
"""

import argparse
import math
import pickle
import sys
import time

import uhd


def parse_args():
    """Parse args and return args object."""
    parser = argparse.ArgumentParser(
        description="Run power level calibration for supported USRPs",
    )
    parser.add_argument(
        "--args",
        default="",
        help="USRP Device Args",
    )
    parser.add_argument(
        "-d",
        "--dir",
        default="tx",
        help="Direction: Must be either rx or tx. From the perspective of the "
        "device: rx == receive power levels, tx == transmit power levels.",
        choices=["tx", "rx"],
    )
    parser.add_argument(
        "--start",
        help="Start Frequency. Defaults to lowest available frequency on device. "
        "Note that this is only a hint for the device object, which can choose "
        "to override this value.",
        type=float,
    )
    parser.add_argument(
        "--stop",
        help="Stop Frequency. Defaults to highest available frequency on device. "
        "Note that this is only a hint for the device object, which can choose "
        "to override this value.",
        type=float,
    )
    parser.add_argument(
        "--step",
        help="Frequency Step. Defaults are device-specific. "
        "Note that this is only a hint for the device object, which can choose "
        "to override this value. Devices can also measure at non-regular "
        "frequencies, e.g., to more accurately cover differences between bands.",
        type=float,
    )
    parser.add_argument(
        "--gain-step",
        type=float,
        default=1,
        help="Gain Step (dB). Defaults to 1 dB. "
        "Note that this is only a hint for the device object, which can choose "
        "to override this value. Devices can also measure at non-regular "
        "gain intervals.",
    )
    parser.add_argument(
        "--lo-offset",
        type=float,
        help="LO Offset. This gets applied to every tune request. Note that for "
        "TX measurements, there is also an offset applied by --tone-freq. "
        "The default value is device-dependent.",
    )
    parser.add_argument(
        "--amplitude",
        type=float,
        default=1.0 / math.sqrt(2),
        help="Amplitude of the tone that is generated for tx measurements. "
        "Default is 1/sqrt(2), or -3 dBFS.",
    )
    parser.add_argument(
        "--attenuation",
        type=float,
        default=0.0,
        help="Amount of attenuation between measurement device and DUT. This will "
        "be accounted for by simple addition, it is treated like a measurement error. "
        "The argument is generally positive, e.g. 30 means 30 dB of attenuation.",
    )
    parser.add_argument(
        "--tone-freq",
        type=float,
        default=1e6,
        help="Frequency of the tone that is generated for Tx measurements. This "
        "has the same effect as setting an LO offset, except in software.",
    )
    parser.add_argument(
        "--antenna",
        default="*",
        help="Select antenna port. A value of '*' means that the calibration "
        "will be repeated on all appropriate antenna ports.",
    )
    parser.add_argument(
        "--channels",
        default="*",
        help="Select channel. A value of '*' means that the calibration "
        "will be repeated on all appropriate channels.",
    )
    parser.add_argument(
        "--meas-dev", default="manual", help="Type of measurement device that is used"
    )
    parser.add_argument(
        "-o",
        "--meas-option",
        default=[],
        action="append",
        help="Options that are passed to the measurement device",
    )
    parser.add_argument(
        "--switch", default="manual", help="Type of switch to be used to connect antennas"
    )
    parser.add_argument(
        "--switch-option", default=[], action="append", help="Options that are passed to the switch"
    )
    parser.add_argument(
        "-r", "--rate", type=float, help="Sampling rate at which the calibration is performed"
    )
    parser.add_argument(
        "--store",
        metavar="filename.pickle",
        help="If provided, will store intermediate cal data. This can be analyzed "
        "separately, or loaded into the tool with --load.",
    )
    parser.add_argument(
        "--load",
        metavar="filename.pickle",
        help="If provided, will load intermediate cal data instead of running a " "measurement.",
    )
    return parser.parse_args()


def sanitize_args(usrp, args, default_rate):
    """Check the args against the USRP object."""
    assert (
        usrp.get_num_mboards() == 1
    ), "Power calibration tools are designed for a single motherboard!"
    available_chans = getattr(usrp, "get_{}_num_channels".format(args.dir))()
    if args.channels == "*":
        # * means all channels
        channels = list(range(available_chans))
    else:
        try:
            channels = [int(c) for c in args.channels.split(",")]
        except ValueError:
            raise ValueError("Invalid channel list: {}".format(args.channels))
        for chan in channels:
            assert chan in range(
                available_chans
            ), "ERROR: Invalid channel: {}. Should be in 0..{}.".format(chan, available_chans)
    print("=== Calibrating for channels:", ", ".join([str(x) for x in channels]))
    available_ants = getattr(usrp, "get_{}_antennas".format(args.dir))()
    if args.antenna == "*":
        invalid_antennas = ("CAL", "LOCAL", "CAL_LOOPBACK", "TERMINATION")
        antennas = [x for x in available_ants if x not in invalid_antennas]
    else:
        try:
            antennas = args.antenna.split(",")
        except ValueError:
            raise ValueError("Invalid antenna list: {}".format(args.antenna))
        for ant in antennas:
            assert ant in available_ants, "Invalid antenna: {}. Should be in {}.".format(
                ant, available_ants
            )
    print("=== Calibrating for antennas:", ", ".join([str(x) for x in antennas]))
    rate = args.rate or default_rate
    getattr(usrp, "set_{}_rate".format(args.dir))(rate)
    actual_rate = getattr(usrp, "get_{}_rate".format(args.dir))()
    print(
        "=== Requested sampling rate: {} Msps, actual rate: {} Msps".format(
            rate / 1e6, actual_rate / 1e6
        )
    )
    if args.dir == "tx" and abs(args.tone_freq) > actual_rate:
        raise ValueError(
            "The TX tone frequency offset of {} kHz is greater than the sampling rate.".format(
                args.tone_freq / 1e3
            )
        )
    return channels, antennas, actual_rate


def init_results(pickle_file):
    """Initialize results from pickle file, or empty dict."""
    if pickle_file is None:
        return {}
    with open(pickle_file, "rb") as results_file:
        return pickle.load(results_file)


class CalRunner:
    """Executor of the calibration routines."""

    def __init__(self, usrp, usrp_cal, meas_dev, args):
        """Initialize the cal runner."""
        self.usrp = usrp
        self.usrp_cal = usrp_cal
        self.meas_dev = meas_dev
        self.dir = args.dir
        self.tone_offset = args.tone_freq if args.dir == "tx" else 0.0
        self.lo_offset = args.lo_offset if args.lo_offset else usrp_cal.lo_offset
        if self.lo_offset:
            print("=== Using USRP LO offset: {:.2f} MHz".format(self.lo_offset / 1e6))

    def run(self, chan, freq):
        """Run all cal steps for a single frequency."""
        print("=== Running calibration at frequency {:.3f} MHz...".format(freq / 1e6))
        tune_req = uhd.types.TuneRequest(freq, self.lo_offset)
        getattr(self.usrp, "set_{}_freq".format(self.dir))(tune_req, chan)
        time.sleep(self.usrp_cal.tune_settling_time)
        actual_freq = getattr(self.usrp, "get_{}_freq".format(self.dir))(chan)
        if abs(actual_freq - freq) > 1.0:
            print(
                "WARNING: Frequency was coerced from {:.2f} MHz to {:.2f} MHz!".format(
                    freq / 1e6, actual_freq / 1e6
                )
            )
        self.meas_dev.set_frequency(actual_freq + self.tone_offset)
        getattr(self.usrp_cal, "run_{}_cal".format(self.dir))(freq)


def main():
    """Go, go, go!"""
    args = parse_args()
    print("=== Detecting USRP...")
    usrp = uhd.usrp.MultiUSRP(args.args)
    print("=== Measurement direction:", args.dir)
    print("=== Initializing measurement device...")
    meas_dev = uhd.usrp.cal.get_meas_device(args.dir, args.meas_dev, args.meas_option)
    meas_dev.power_offset = args.attenuation
    # If we're transmitting, then we need to factor in the "attenuation" from us
    # not transmitting at full scale
    if args.dir == "tx":
        meas_dev.power_offset -= 20 * math.log10(args.amplitude)
    print("=== Initializing port connector...")
    switch = uhd.usrp.cal.get_switch(args.dir, args.switch, args.switch_option)
    print("=== Initializing USRP calibration object...")
    usrp_cal = uhd.usrp.cal.get_usrp_calibrator(
        usrp,
        meas_dev,
        args.dir,
        gain_step=args.gain_step,
    )
    channels, antennas, rate = sanitize_args(usrp, args, usrp_cal.default_rate)
    results = init_results(args.load)
    usrp_cal.init(
        rate=rate,
        tone_freq=args.tone_freq,
        amplitude=args.amplitude,
    )
    print("=== Launching calibration...")
    cal_runner = CalRunner(usrp, usrp_cal, meas_dev, args)
    for chan in channels:
        if chan not in results:
            results[chan] = {}
        for ant in antennas:
            if ant in results[chan]:
                print("=== Using pickled data for channel {}, antenna {}.".format(chan, ant))
                continue
            print("=== Running calibration for channel {}, antenna {}.".format(chan, ant))
            # Set up all the objects
            getattr(usrp, "set_{}_antenna".format(args.dir))(ant, chan)
            switch.connect(chan, ant)
            usrp_cal.update_port(chan, ant)
            freqs = usrp_cal.init_frequencies(args.start, args.stop, args.step)
            usrp_cal.start()  # This will activate siggen
            # Now calibrate
            for freq in freqs:
                try:
                    cal_runner.run(chan, freq)
                except RuntimeError as ex:
                    print("ERROR: Stopping calibration due to exception: {}".format(str(ex)))
                    usrp_cal.stop()
                    return 1
            # Store results for pickling and shut down for next antenna port
            results[chan][ant] = usrp_cal.results
            usrp_cal.stop()  # This will deactivate siggen and store the data
    if args.store:
        print("=== Storing pickled calibration data to {}...".format(args.store))
        with open(args.store, "wb") as results_file:
            pickle.dump(results, results_file)
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except (RuntimeError, ValueError) as ex:
        print("ERROR:", str(ex))
        sys.exit(1)
