#
# Copyright 2018 Ettus Research, a National Instruments Company
#
# SPDX-License-Identifier: GPL-3.0-or-later
#
"""
Helper class to initialize a Rhodium daughterboard
"""

from __future__ import print_function

import time

from usrp_mpm.cores import ClockSynchronizer, nijesdcore
from usrp_mpm.cores.eyescan import EyeScanTool
from usrp_mpm.dboard_manager.gain_rh import GainTableRh
from usrp_mpm.dboard_manager.lmk_rh import LMK04828Rh
from usrp_mpm.dboard_manager.rh_periphs import DboardClockControl
from usrp_mpm.sys_utils.uio import open_uio


class RhodiumInitManager(object):
    """
    Helper class: Holds all the logic to initialize an N320/N321 (Rhodium)
    daughterboard.
    """

    # The Phase DAC is set at midscale, having its flatness validate +/- 1023 codes
    # from this initial value.
    INIT_PHASE_DAC_WORD = 32768
    PHASE_DAC_SPI_ADDR = 0x3
    # External PPS pipeline delay from the PPS captured at the FPGA to TDC input,
    # in reference clock ticks
    EXT_PPS_DELAY = 5
    # Variable PPS delay before the RP/SP pulsers begin. Fixed value for the N3xx devices.
    N3XX_INT_PPS_DELAY = 4
    # JESD core default configuration.
    JESD_DEFAULT_ARGS = {
        "lmfc_divider": 12,
        "rx_sysref_delay": 5,
        "tx_sysref_delay": 11,
        "tx_driver_swing": 0b1101,
        "tx_precursor": 0b00100,
        "tx_postcursor": 0b00100,
    }
    # After testing the roundtrip latency (i.e. FPGA -> TX -> RX -> FPGA),
    # it was found that a different value of RX SYSREF delay is required
    # for sampling_clock_rate = 400 MSPS to achieve latency consistency.
    RX_SYSREF_DLY_DIC = {400e6: 6, 491.52e6: 5, 500e6: 5}

    def __init__(self, rh_class, spi_ifaces):
        self.rh_class = rh_class
        self._spi_ifaces = spi_ifaces
        self.adc = rh_class.adc
        self.dac = rh_class.dac
        self.slot_idx = rh_class.slot_idx
        self.log = rh_class.log.getChild("init")

    def _init_lmk(
        self,
        lmk_spi,
        ref_clk_freq,
        sampling_clock_rate,
        pdac_spi,
        init_phase_dac_word,
        phase_dac_spi_addr,
    ):
        """
        Sets the phase DAC to initial value, and then brings up the LMK
        according to the selected ref clock frequency.
        Will throw if something fails.
        """
        self.log.trace("Initializing Phase DAC to d{}.".format(init_phase_dac_word))
        pdac_spi.poke16(phase_dac_spi_addr, init_phase_dac_word)
        return LMK04828Rh(self.slot_idx, lmk_spi, ref_clk_freq, sampling_clock_rate, self.log)

    def _sync_db_clock(self, dboard_ctrl_regs, ref_clk_freq, master_clock_rate, args):
        "Synchronizes the DB clock to the common reference"
        reg_offset = 0x200
        ext_pps_delay = self.EXT_PPS_DELAY
        if args.get("time_source", self.rh_class.default_time_source) == "sfp0":
            reg_offset = 0x400
            ref_clk_freq = 62.5e6
            ext_pps_delay = 1  # only 1 flop between the WR core output and the TDC input
        synchronizer = ClockSynchronizer(
            dboard_ctrl_regs,
            self.rh_class.lmk,
            self._spi_ifaces["phase_dac"],
            reg_offset,
            master_clock_rate,
            ref_clk_freq,
            1.116e-12,  # fine phase shift. TODO don't hardcode. This should live in the EEPROM
            self.INIT_PHASE_DAC_WORD,
            self.PHASE_DAC_SPI_ADDR,
            ext_pps_delay,
            self.N3XX_INT_PPS_DELAY,
            self.slot_idx,
        )
        # The radio clock traces on the motherboard are 69 ps longer for Daughterboard B
        # than Daughterboard A. We want both of these clocks to align at the converters
        # on each board, so adjust the target value for DB B. This is an N3xx series
        # peculiarity and will not apply to other motherboards.
        trace_delay_offset = {0: 0.0e-0, 1: 69.0e-12}[self.slot_idx]
        offset_error = abs(synchronizer.run(num_meas=[512, 128], target_offset=trace_delay_offset))
        if offset_error > 100e-12:
            self.log.error(
                "Residual clock synchronization offset error is invalid! "
                "Expected: <100 ps Actual: {:.1f} ps!".format(offset_error * 1e12)
            )
            raise RuntimeError("Clock synchronization offset error is greater than expected.")
        else:
            self.log.debug("Residual synchronization error: {:.1f} ps.".format(offset_error * 1e12))
        synchronizer = None
        self.log.debug("Sample Clock Synchronization Complete!")

    def set_jesd_rate(self, jesdcore, new_rate, current_jesd_rate, force=False):
        """
        Make the QPLL and GTX changes required to change the JESD204B core rate.
        """
        # The core is directly compiled for 500 MHz sample rate, which
        # corresponds to a lane rate of 5.0 Gbps. The same QPLL and GTX settings apply
        # for the 491.52 MHz sample rate.
        #
        # The lower non-LTE rate, 400 MHz, requires changes to the default configuration
        # of the MGT components. This function performs the required changes in the
        # following order (as recommended by UG476).
        #
        # 1) Modify any QPLL settings.
        # 2) Perform the QPLL reset routine by pulsing reset then waiting for lock.
        # 3) Modify any GTX settings.
        # 4) Perform the GTX reset routine by pulsing reset and waiting for reset done.

        assert new_rate in (4000e6, 4915.2e6, 5000e6)

        # On first run, we have no idea how the FPGA is configured... so let's force an
        # update to our rate.
        force = force or (current_jesd_rate is None)

        skip_drp = False
        if not force:
            #           Current     New       Skip?
            skip_drp = {
                4000.0e6: {4000.0e6: True, 4915.2e6: False, 5000.0e6: False},
                4915.2e6: {4000.0e6: False, 4915.2e6: True, 5000.0e6: True},
                5000.0e6: {4000.0e6: False, 4915.2e6: True, 5000.0e6: True},
            }[current_jesd_rate][new_rate]

        if skip_drp:
            self.log.trace(
                "Current lane rate is compatible with the new rate. Skipping " "reconfiguration."
            )

        # These are the only registers in the QPLL and GTX that change based on the
        # selected line rate. The MGT wizard IP was generated for each of the rates and
        # reference clock frequencies and then diffed to create this table.
        QPLL_CFG = {4000.0e6: 0x6801C1, 4915.2e6: 0x680181, 5000.0e6: 0x0680181}[new_rate]
        MGT_RX_CLK25_DIV = {4000.0e6: 8, 4915.2e6: 10, 5000.0e6: 10}[new_rate]
        MGT_TX_CLK25_DIV = {4000.0e6: 8, 4915.2e6: 10, 5000.0e6: 10}[new_rate]

        # 1-2) Do the QPLL first
        if not skip_drp:
            self.log.trace("Changing QPLL settings to support {} Gbps".format(new_rate / 1e9))
            jesdcore.set_drp_target("qpll", 0)
            # QPLL_CONFIG is spread across two regs: 0x32 (dedicated) and 0x33 (shared)
            reg_x32 = QPLL_CFG & 0xFFFF  # [16:0] -> [16:0]
            reg_x33 = jesdcore.drp_access(rd=True, addr=0x33)
            reg_x33 = (reg_x33 & 0xF800) | ((QPLL_CFG >> 16) & 0x7FF)  # [26:16] -> [11:0]
            jesdcore.drp_access(rd=False, addr=0x32, wr_data=reg_x32)
            jesdcore.drp_access(rd=False, addr=0x33, wr_data=reg_x33)

        # Run the QPLL reset sequence and prep the MGTs for modification.
        jesdcore.init()

        # 3-4) And the 4 MGTs second
        if not skip_drp:
            self.log.trace("Changing MGT settings to support {} Gbps".format(new_rate / 1e9))
            for lane in range(4):
                jesdcore.set_drp_target("mgt", lane)
                # MGT_RX_CLK25_DIV is embedded with others in 0x11. The
                # encoding for the DRP register value is one less than the
                # desired value.
                reg_x11 = jesdcore.drp_access(rd=True, addr=0x11)
                reg_x11 = (reg_x11 & 0xF83F) | ((MGT_RX_CLK25_DIV - 1 & 0x1F) << 6)  # [10:6]
                jesdcore.drp_access(rd=False, addr=0x11, wr_data=reg_x11)
                # MGT_TX_CLK25_DIV is embedded with others in 0x6A. The
                # encoding for the DRP register value is one less than the
                # desired value.
                reg_x6a = jesdcore.drp_access(rd=True, addr=0x6A)
                reg_x6a = (reg_x6a & 0xFFE0) | (MGT_TX_CLK25_DIV - 1 & 0x1F)  # [4:0]
                jesdcore.drp_access(rd=False, addr=0x6A, wr_data=reg_x6a)
            self.log.trace("GTX settings changed to support {} Gbps".format(new_rate / 1e9))
            jesdcore.disable_drp_target()

        self.log.trace("JESD204b Lane Rate set to {} Gbps!".format(new_rate / 1e9))
        return new_rate

    def init_jesd(self, jesdcore, sampling_clock_rate):
        """
        Bringup the JESD links between the ADC, DAC, and the FPGA.
        All clocks must be set up and stable before starting this routine.
        """
        jesdcore.check_core()

        # JESD Lane Rate only depends on the sampling_clock_rate selection, since all
        # other link parameters (LMFS,N) remain constant.
        L = 4
        M = 2
        F = 1
        S = 1
        N = 16
        new_rate = sampling_clock_rate * M * N * (10.0 / 8) / L / S
        self.log.trace("Calculated JESD204B lane rate is {} Gbps".format(new_rate / 1e9))
        self.rh_class.current_jesd_rate = self.set_jesd_rate(
            jesdcore, new_rate, self.rh_class.current_jesd_rate
        )

        self.log.trace("Setting up JESD204B TX blocks.")
        jesdcore.init_framer()  # Initialize FPGA's framer.
        self.adc.init_framer()  # Initialize ADC's framer.

        self.log.trace("Enabling SYSREF capture blocks.")
        self.dac.enable_sysref_capture(True)  # Enable DAC's SYSREF capture.
        self.adc.enable_sysref_capture(True)  # Enable ADC's SYSREF capture.
        jesdcore.enable_lmfc(True)  # Enable FPGA's SYSREF capture.

        self.log.trace("Setting up JESD204B DAC RX block.")
        self.dac.init_deframer()  # Initialize DAC's deframer.

        self.log.trace("Sending SYSREF to all devices.")
        jesdcore.send_sysref_pulse()  # Send SYSREF to all devices.

        self.log.trace("Setting up JESD204B FPGA RX block.")
        jesdcore.init_deframer()  # Initialize FPGA's deframer.

        self.log.trace("Disabling SYSREF capture blocks.")
        self.dac.enable_sysref_capture(False)  # Disable DAC's SYSREF capture.
        self.adc.enable_sysref_capture(False)  # Disable ADC's SYSREF capture.
        jesdcore.enable_lmfc(False)  # Disable FPGA's SYSREF capture.

        time.sleep(0.100)  # Allow time for CGS/ILA.

        self.log.trace("Verifying JESD204B link status.")
        error_flag = False
        if not jesdcore.get_framer_status():
            self.log.error("JESD204b FPGA Core Framer is not synced!")
            error_flag = True
        if not self.dac.check_deframer_status():
            self.log.error("DAC JESD204B Deframer is not synced!")
            error_flag = True
        if not self.adc.check_framer_status():
            self.log.error("ADC JESD204B Framer is not synced!")
            error_flag = True
        if not jesdcore.get_deframer_status():
            self.log.error("JESD204B FPGA Core Deframer is not synced!")
            error_flag = True
        if error_flag:
            raise RuntimeError("JESD204B Link Initialization Failed. See MPM logs for details.")
        self.log.info("JESD204B Link Initialization & Training Complete")

    def init(self, args):
        """
        Run the full initialization sequence. This will bring everything up
        from scratch: The LMK, JESD cores, the AD9695, the DAC37J82, and
        anything else that is clocking-related.
        Depending on the settings, this can take a fair amount of time.
        """
        # Input validation on RX margin tests (@ FPGA and DAC)
        # By accepting the rx_eyescan/tx_prbs argument being str or bool, one may
        # request an eyescan measurement to be performed from either the USRP's
        # shell (i.e. using --default-args) or from the host's MPM shell.
        perform_rx_eyescan = False
        if "rx_eyescan" in args:
            perform_rx_eyescan = (args["rx_eyescan"] == "True") or (args["rx_eyescan"] == True)
        if perform_rx_eyescan:
            self.log.trace("Adding RX eye scan PMA enable to JESD args.")
            self.JESD_DEFAULT_ARGS["enable_rx_eyescan"] = True
        perform_tx_prbs = False
        if "tx_prbs" in args:
            perform_tx_prbs = (args["tx_prbs"] == "True") or (args["tx_prbs"] == True)

        # Latency across the JESD204B TX/RX links should remain constant and
        # deterministic across the supported sampling_clock_rate values.
        # After testing the roundtrip latency (i.e. FPGA -> TX -> RX -> FPGA),
        # it was found that a different set of SYSREF delay values are required
        # for sampling_clock_rate = 400 MSPS to achieve latency consistency.
        self.JESD_DEFAULT_ARGS["rx_sysref_delay"] = self.RX_SYSREF_DLY_DIC[
            self.rh_class.sampling_clock_rate
        ]

        # Bringup Sequence.
        #   1. Prerequisites (include opening mmaps)
        #   2. Initialize LMK and bringup clocks.
        #   3. Synchronize DB Clocks.
        #   4. Initialize FPGA JESD IP.
        #   5. DAC Configuration.
        #   6. ADC Configuration.
        #   7. JESD204B Initialization.
        #   8. CPLD Gain Tables Initialization.

        # 1. Prerequisites
        # Open FPGA IP (Clock control and JESD core).
        self.log.trace("Creating gain table object...")
        self.gain_table_loader = GainTableRh(
            self._spi_ifaces["cpld"], self._spi_ifaces["cpld_gain_loader"], self.log
        )

        with open_uio(
            label="dboard-regs-{}".format(self.rh_class.slot_idx), read_only=False
        ) as radio_regs:
            self.log.trace("Creating dboard clock control object")
            db_clk_control = DboardClockControl(radio_regs, self.log)
            self.log.trace("Creating jesdcore object")
            jesdcore = nijesdcore.NIJESDCore(
                radio_regs, self.rh_class.slot_idx, **self.JESD_DEFAULT_ARGS
            )

            # 2. Initialize LMK and bringup clocks.
            # Disable FPGA MMCM's outputs, and assert its reset.
            db_clk_control.reset_mmcm()
            # Always place the JESD204b cores in reset before modifying the clocks,
            # otherwise high power or erroneous conditions could exist in the FPGA!
            jesdcore.reset()
            # Configure and bringup the LMK's clocks.
            self.log.trace("Initializing LMK...")
            self.rh_class.lmk = self._init_lmk(
                self._spi_ifaces["lmk"],
                self.rh_class.ref_clock_freq,
                self.rh_class.sampling_clock_rate,
                self._spi_ifaces["phase_dac"],
                self.INIT_PHASE_DAC_WORD,
                self.PHASE_DAC_SPI_ADDR,
            )
            self.log.trace("LMK Initialized!")
            # Deassert FPGA's MMCM reset, poll for lock, and enable outputs.
            db_clk_control.enable_mmcm()

            # 3. Synchronize DB Clocks.
            # The clock synchronzation driver receives the master_clock_rate, which for
            # Rhodium is half the sampling_clock_rate.
            self._sync_db_clock(
                radio_regs,
                self.rh_class.ref_clock_freq,
                self.rh_class.sampling_clock_rate / 2,
                args,
            )

            # 4. DAC Configuration.
            self.dac.config()

            # 5. ADC Configuration.
            self.adc.config()

            # 6-7. JESD204B Initialization.
            self.init_jesd(jesdcore, self.rh_class.sampling_clock_rate)
            # [Optional] Perform RX eyescan.
            if perform_rx_eyescan:
                self.log.info("Performing RX eye scan on ADC to FPGA link...")
                self._rx_eyescan(jesdcore, args)
            # [Optional] Perform TX PRBS test.
            if perform_tx_prbs:
                self.log.info("Performing TX PRBS-31 test on FPGA to DAC link...")
                self._tx_prbs_test(jesdcore, args)
            # Direct the garbage collector by removing our references
            jesdcore = None
            db_clk_control = None

        # 8. CPLD Gain Tables Initialization.
        self.gain_table_loader.init()

        return True

    ##########################################################################
    # JESD204B RX margin testing
    ##########################################################################

    def _rx_eyescan(self, jesdcore, args):
        """
        This function creates an eyescan object to perform this measurement with the
        given configuration and lanes.

        Parameters:
          prescale   -> Controls the prescaling of the sample count to keep both sample
                        count and error count in reasonable precision.
                        Valid values: from 0 to 31.
        """
        # The following constants must be defined according to GTs configuration
        # for each project. For further details, refer to the eyescan.py file.
        # For Rhodium, these parameters are based on the JESD core.
        rxout_div = 2
        rx_int_datawidth = 20
        eq_mode = "LPM"
        # The following variables define the GTs to be scanned and the range of the
        # measurement.
        prescale = 0
        scan_lanes = [0, 1, 2, 3]
        hor_range = {"start": -32, "stop": 32, "step": 2}
        ver_range = {"start": -127, "stop": 127, "step": 2}
        # Set default configuration values for Rhodium when the user is not intentionally
        # changing the constants/variables default values.
        for key in (
            "rxout_div",
            "rx_int_datawidth",
            "eq_mode",
            "prescale",
            "scan_lanes",
            "hor_range",
            "ver_range",
        ):
            if key not in args:
                self.log.trace(
                    "Setting Rh default value for {0}... val: {1}".format(key, locals()[key])
                )
                args[key] = locals()[key]
        #
        # Create an eyescan object.
        assert jesdcore is not None
        eyescan_tool = EyeScanTool(jesdcore, self.slot_idx, **args)
        # Put the ADC in pseudorandom test mode.
        adc_regs = self._spi_ifaces["adc"]
        # test_val = adc_regs.peek8(0x0550)
        # adc_regs.poke8(0x0550, 0x05)
        test_val = adc_regs.peek8(0x0573)
        adc_regs.poke8(0x0573, 0x13)
        # Perform eye scan on given lanes and range.
        file_name = eyescan_tool.eyescan_full_scan(
            args["scan_lanes"], args["hor_range"], args["ver_range"]
        )
        # Do some housekeeping...
        # adc_regs.poke8(0x0550, test_val) # Enable normal operation.
        adc_regs.poke8(0x0573, test_val)  # Enable normal operation.
        adc_regs.poke8(0x0000, 0x81)  # Reset.
        eyescan_tool = None
        return file_name

    def _tx_prbs_test(self, jesdcore, args):
        """
        This function allows to test the PRBS-31 pattern at the DAC.
        """

        def _test_lanes(**tx_settings):
            """
            This methods enables, monitors, and disables the PRBS-31 test.
            """
            results = []
            jesdcore.adjust_tx_phy(**tx_phy_settings)
            self.log.info(
                "Testing TX PHY settings: tx_driver_swing=0b{0:04b}"
                "  tx_precursor=0b{1:05b}"
                "  tx_postcursor=0b{2:05b}".format(
                    tx_phy_settings["tx_driver_swing"],
                    tx_phy_settings["tx_precursor"],
                    tx_phy_settings["tx_postcursor"],
                )
            )
            # Enable the GTs TX pattern generator in PRBS-31 mode.
            jesdcore.set_pattern_gen(mode="PRBS-31")
            # Monitor each receive lane at DAC.
            for lane_num in range(0, 4):
                self.dac.test_mode(mode="PRBS-31", lane=lane_num)  # Enable PRBS test mode.
                number_of_failures = 0
                for _ in range(0, POLLS_PER_GT):
                    time.sleep(WAIT_TIME_PER_POLL)
                    alarm_pin_dac = self.rh_class.cpld.get_dac_alarm()
                    if alarm_pin_dac:
                        number_of_failures += 1
                results.append(number_of_failures)
                if number_of_failures > 0:
                    self.log.error(
                        "PRBS-31 test for DAC lane {0} failed {1}/{2}!".format(
                            lane_num, number_of_failures, POLLS_PER_GT
                        )
                    )
                else:
                    self.log.info("PRBS-31 test for DAC lane {0} passed!".format(lane_num))
                self.dac.test_mode(mode="OFF", lane=lane_num)  # Disable PRBS test mode.
            # Disable TX pattern generator at FPGA
            jesdcore.set_pattern_gen(mode="OFF")
            return results

        #
        WAIT_TIME_PER_POLL = 0.001  # in seconds.
        POLLS_PER_GT = 100
        # Create the CSV file.
        f = open("tx_prbs_sweep.csv", "w")
        f.write("Swing,Precursor,Postcursor,Polls,Failures 0,Failures 1,Failures 2,Failures 3\n")
        # Default TX PHY settings.
        tx_phy_settings = {
            "tx_driver_swing": 0b1111,  # See UG476, TXDIFFCTRL
            "tx_precursor": 0b00000,  # See UG476, TXPRECURSOR
            "tx_postcursor": 0b00000,
        }  # See UG476, TXPOSTCURSOR
        # Define sweep ranges.
        DEFAULT_SWING_RANGE = {"start": 0b0000, "stop": 0b1111 + 0b1, "step": 1}
        DEFAULT_CURSOR_RANGE = {"start": 0b00000, "stop": 0b11111 + 0b1, "step": 2}
        swing_range = args.get("swing_range", DEFAULT_SWING_RANGE)
        precursor_range = args.get("precursor_range", DEFAULT_CURSOR_RANGE)
        postcursor_range = args.get("postcursor_range", DEFAULT_CURSOR_RANGE)
        # Test the TX margin across multiple PHY settings.
        for swing in range(swing_range["start"], swing_range["stop"], swing_range["step"]):
            tx_phy_settings["tx_driver_swing"] = swing
            for precursor in range(
                precursor_range["start"], precursor_range["stop"], precursor_range["step"]
            ):
                tx_phy_settings["tx_precursor"] = precursor
                for postcursor in range(
                    postcursor_range["start"], postcursor_range["stop"], postcursor_range["step"]
                ):
                    tx_phy_settings["tx_postcursor"] = postcursor
                    results = _test_lanes(**tx_phy_settings)
                    f.write(
                        "{},{},{},{},{},{},{},{}\n".format(
                            tx_phy_settings["tx_driver_swing"],
                            tx_phy_settings["tx_precursor"],
                            tx_phy_settings["tx_postcursor"],
                            POLLS_PER_GT,
                            results[0],
                            results[1],
                            results[2],
                            results[3],
                        )
                    )
        # Housekeeping...
        f.close()
