# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.19.5
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# %%
"""
Import python libraries, and input pl_quiesce.py.  This is a python script Claude wrote to make sure that the PS and PL are in sync when redownloading.  

When PYNQ downloads a bitstream, any AXI master still actively transacting on an HP port belongs to the *previous* design.  Reconfiguration erases the logic but not the outstanding transaction, which can leave the new design's first memory-side access wedged.  Nothing in the new design's RTL can prevent this -- the offending logic no longer exists by the time the new design does -- so the fix has to happen on the PS side, before download.

What it does is:

1. Ask the PL server for ip_dict — the metadata for the design loaded right now, i.e. the previous one, since you haven't downloaded yet. If unavailable, print and return.
2. Filter it to IPs that are named in the parameter DMA_NAMES and have axi_dma in their type field.
3. Print matched names. If nothing matched, return — that's the normal path on the first download after boot.

Then for each matched DMA, map its registers from phys_addr, and run these steps on MM2S, then again on S2MM:

4. Read DMACR and check bit 0, Run/Stop. Record it as was_running.
5. If it's set, clear it. Per PG021 the engine finishes its current transaction and then asserts Halted — this is the graceful stop, and it's what lets an in-flight memory transaction complete rather than being abandoned mid-burst.
6. Poll DMASR bit 0 for up to 500 ms waiting for Halted. Record whether it got there.
7. Set DMACR bit 2, the soft reset, whether or not step 6 succeeded. It's the harder stop if the graceful one timed out, and free if it didn't.
8. Poll for that bit to self-clear, up to 500 ms. Record it.
9. Print the per-channel line and append the result dict.

Every step that can fail is caught and recorded rather than raised, so a partial teardown still lets the download proceed — with a visible INCOMPLETE in the log.

This was developed for the Loopback projects, which uses names "AXI_DMA_SG" and "AXI_DMA_NOSG" for the 2 DMA engines.  if you want to add to the list edit DMA_NAMES below
-----

Here is the mapping from board label to tile/block in the radio function:

Board label Tile Driver path ADC_A 226 adc_tiles[2].blocks[1] ADC_B 226 adc_tiles[2].blocks[0] ADC_C 224 adc_tiles[0].blocks[1] ADC_D 224 adc_tiles[0].blocks[0] DAC_A 230 dac_tiles[2].blocks[0] DAC_B 228 dac_tiles[0].blocks[0]

For this project we are using ADC on tile 226, adc 1, so that's ADC_A, and DAC on tile 228, dac 0, so that's DAC_B
"""


from pynq import Overlay, MMIO, allocate, Clocks
from pynq.lib import AxiGPIO
from pynq.overlays.base import BaseOverlay
from rfsoc4x2 import oled
from xrfdc import RFdc
import time
import xrfclk
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from pl_quiesce import quiesce_pl
from hwh_check import check_hwh
#
# parameters for the FPGA project Loopback, version 0xa5000001
#
# AXI GPIO:  (do gen.ip_dict to verify)
version_name = "VERSION"     # for the firmware version
rfdc_name = "RFDC"           # the RFDC data converter
rxfifo_reset_name = "RX_FIFO_RESET"   # to reset the RX_FIFO to keep it empty until we want to start reading
tlast_gen_name = "RX_FIFO_TLAST"      # to set the number of transfers to generate TLAST
#
# adc tiles are number 0-3 for 224-227 and block is 0/1 for ADC 0/ADC 1.  we are using 226, adc 1
#
adc_tile = 226
adc_tile_n = adc_tile - 224
adc_block_n = 1
#
# dac tiles are number 0-3 for 228-231 and same for block.  we are using 228, dac 0
#
dac_tile = 228
dac_tile_n = dac_tile - 228
dac_block_n = 0
#
# frequencies in MHz, decimation and interpolation values
#
rfdc_dac_sampling = 6881.28
rfdc_adc_sampling = 4915.2
rfdc_dac_interpolation = 40
rfdc_adc_decimation = 40
#
# clock frequency for the PS in MHz
#
ps_clock = 0
#
# firmware next:
#
firmware = "Loopback_v1.bit"
#
# DMA instance names to quiesce before download (see gen.ip_dict to verify)
#
dma_names = ["AXI_DMA_SG", "AXI_DMA_NOSG",   # Loopback
             "axi_dma_1", "axi_dma_2"]       # cyclic_clock

class Gen(BaseOverlay):
    def __init__(self, quiesce=True):
        # 
        # Run pl_quiesce BEFORE Overlay.__init__() reconfigures the PL: the master we
        # need to stop belongs to the design loaded right now, not the new one.
        #
        if quiesce:
            quiesce_pl(dma_names)
        #
        # load the new firmware
        #
        Overlay.__init__(self,firmware)
        #
        # Check to see if we forgot to connect anything obvious in the project!
        #
        check_hwh(self.bitfile_name)
        #
        # Dig out the firmware version 
        #
        version_ip = self.ip_dict[version_name]
        version_ptr = AxiGPIO(version_ip).channel1
        self.version = version_ptr.read()
        #
        # Do some housekeeping
        #
        if self.is_loaded():
            #
            # self.radio is used to point to the radio class
            #
            self.radio = getattr(self, rfdc_name)   # i.e., self.RFDC
            #
            # initialize the LMK and LMX chips
            #
            self.init_rf_clks()
            #
            # set default carrier frequencies (not really needed...)
            #
            self.freq_adc = 700.0      # MHz
            self.freq_dac = 700.0      # MHz
                
    def get_version(self):
        return self.version
    
    def set_external_rf_clock(self):
            print("Switching to external 10MHz clock...")
            for lmk in xrfclk.lmk_devices:
                with open(lmk['spi_device'], 'rb+', buffering=0) as f:
                    data = b'\x01\x47\x0A' # Only works for LMK0482x series
                    f.write(data)

    def set_internal_rf_clock(self):
        print("Switching to internal 10MHz clock...")
        for lmk in xrfclk.lmk_devices:
            with open(lmk['spi_device'], 'rb+', buffering=0) as f:
                data = b'\x01\x47\x1A' # Only works for LMK0482x series
                f.write(data)                
            
    def dac_status(self, settle=0.02, verbose=False):
        #
        # so we can tell the AI if something goes wrong what happened!
        #
        rm   = self.AXI_DMA_SG.register_map
        tile = self.RFDC.dac_tiles[dac_tile_n]
        blk  = tile.blocks[dac_block_n]

        cr, sr = rm.MM2S_DMACR, rm.MM2S_DMASR

        # poll until the descriptor pointer moves, or until 'settle' elapses.
        # (only 2 descriptors in the ring, so a fixed 2-sample gap aliases badly)
        d0 = int(rm.MM2S_CURDESC.Current_Descriptor_Pointer)
        d1 = d0
        advancing = False
        t0 = time.time()
        while time.time() - t0 < settle:
            d1 = int(rm.MM2S_CURDESC.Current_Descriptor_Pointer)
            if d1 != d0:
                advancing = True
                break
            time.sleep(0.0005)

        ts = self.RFDC.IPStatus['DACTileStatus'][dac_tile_n]
        bs = blk.BlockStatus

        r = {
            'rs':        int(cr.RS),
            'cyclic':    int(cr.Cyclic_BD_Enable),
            'halted':    int(sr.Halted),
            'idle':      int(sr.Idle),
            'errors':    int(sr.DMAIntErr or sr.DMASlvErr or sr.DMADecErr or
                             sr.SGIntErr  or sr.SGSlvErr  or sr.SGDecErr),
            'curdesc0':  hex(d0),
            'curdesc1':  hex(d1),
            'advancing': advancing,
            'tile_state':   ts['TileState'],
            'tile_power':   ts['PowerUpState'],
            'tile_pll':     ts['PLLState'],
            'dp_clocks':    bs['DataPathClocksStatus'],
            'fifo_flags':   bs['IsFIFOFlagsAsserted'],   # latches; informational only
            'fifo_enabled': tile.FIFOStatus,             # informational until baselined
            'nco_mhz':      blk.MixerSettings['Freq'],
        }
        r['ok'] = (r['rs'] and r['cyclic'] and not r['halted'] and not r['idle']
                   and not r['errors'] and r['advancing']
                   and r['tile_state'] == 15 and r['tile_pll'] == 1
                   and r['dp_clocks'] == 1)
        if verbose:
            for k, v in r.items():
                print(f"   {k:12s} {v}")
        return r

    def dac_hard_stop(self, timeout=0.05, reset_timeout=2.0):
        #
        # Stop MM2S with a timeout. pynq's stop() spins forever on 'Halted',
        # which never arrives while the stream side is backpressured.
        #
        rm = self.AXI_DMA_SG.register_map
        ch = self.AXI_DMA_SG.sendchannel
        rm.MM2S_DMACR.RS = 0
        t0 = time.time()
        while not rm.MM2S_DMASR.Halted:
            if time.time() - t0 > timeout:
                print("MM2S would not halt; forcing soft reset. DMASR:", rm.MM2S_DMASR)
                rm.MM2S_DMACR.Reset = 1
                t1 = time.time()
                while rm.MM2S_DMACR.Reset:
                    if time.time() - t1 > reset_timeout:
                        raise RuntimeError("MM2S soft reset did not complete")
                    time.sleep(0.001)
                print(f"soft reset completed in {time.time() - t1:.3f} s")
                break
            time.sleep(0.001)
        # keep pynq's channel bookkeeping consistent, so transfer() works next
        ch._cyclic = False
        ch._transfer_started = False
        return bool(rm.MM2S_DMASR.Halted)
    

    @property 
    def freq_dac(self): # MHz
        return self.radio.dac_tiles[dac_tile_n].blocks[dac_block_n].MixerSettings['Freq']
    
    @freq_dac.setter
    def freq_dac(self,freq): # MHz
        blk = self.radio.dac_tiles[dac_tile_n].blocks[dac_block_n]
        # get the setting and create the dictionary, then change the values and write it back, 1 write
        s = blk.MixerSettings      # single read
        s['EventSource'] = 0       # same dict object → goes out with the Freq write
        s['Freq'] = freq
        blk.UpdateEvent(1)


    @property 
    def freq_adc(self): # MHz
        return self.radio.adc_tiles[adc_tile_n].blocks[adc_block_n].MixerSettings['Freq']
    
    
    @freq_adc.setter
    def freq_adc(self, freq):  # MHz
        blk = self.radio.adc_tiles[adc_tile_n].blocks[adc_block_n]
        # get the setting and create the dictionary, then change the values and write it back, 1 write
        s = blk.MixerSettings      # single read
        s['EventSource'] = 2       # ADC needs 2, not 0!  same dict object → goes out with the Freq write
        s['Freq'] = freq
        blk.UpdateEvent(1)
            
gen = Gen()
ps_clock = Clocks.fclk0_mhz
version = gen.version
oled = oled.oled_display()
oled.write("Version \n"+hex(version))
print("Firmware version: ",hex(version))
print("PS Clock: ", ps_clock, " MHz")
print("DAC: sampling at "+str(rfdc_dac_sampling)+" MHz and interoplation "+str(rfdc_dac_interpolation))
print("ADC: sampling at "+str(rfdc_adc_sampling)+" MHz and decimation "+str(rfdc_adc_decimation))
print("all done")
gen.ip_dict

# %%
"""
get relevant parameters from the actual values that the hardware rf converter drivers have.  
these are based on the actual values of the reference clocks.  
compare to the parameters at the beginning of the first cell above
"""
rfdc = gen.RFDC
st = rfdc.IPStatus

def report(tiles, tstat, label, base, factor_attr, verbose=True):
    """Returns {tile_idx: {'phys':..., 'fs_msps':..., 'blocks':
                           {blk_idx: {'factor':..., 'fs_msps':...,
                                      'fabric_msps':..., 'nco_mhz':...}}}}"""
    out = {}
    for t, tile in enumerate(tiles):
        s = tstat[t]
        if not s['IsEnabled']:
            continue
        try:
            fs = tile.PLLConfig['SampleRate'] * 1000.0     # GHz -> MSps
        except RuntimeError:
            fs = None
        rec = {'phys': base + t, 'fs_msps': fs, 'blocks': {}}
        for b in range(4):
            if not (s['BlockStatusMask'] >> b) & 1:
                continue
            blk = tile.blocks[b]
            factor = getattr(blk, factor_attr)
            bfs = blk.BlockStatus['SamplingFreq'] * 1000.0
            rec['blocks'][b] = {
                'factor':      factor,
                'fs_msps':     bfs,
                'fabric_msps': bfs / factor if factor else None,
                'nco_mhz':     blk.MixerSettings['Freq'],
            }
        out[t] = rec
        if verbose:
            print(f"{label} tile {t} (phys {rec['phys']}): Fs = {fs:.4f} MSps, "
                  f"blocks enabled = {list(rec['blocks'])}")
            for b, d in rec['blocks'].items():
                print(f"    block {b}: {factor_attr} = {d['factor']}, "
                      f"fabric = {d['fabric_msps']:.4f} MSps, "
                      f"NCO = {d['nco_mhz']} MHz")
    return out

def only_block(rec, label):
    blocks = rec['blocks']
    if not blocks:
        raise RuntimeError(f"{label}: no enabled blocks")
    if len(blocks) != 1:
        print(f"note: {label} has {len(blocks)} enabled blocks {list(blocks)}; taking lowest")
    return blocks[min(blocks)]

dac = report(rfdc.dac_tiles, st['DACTileStatus'], "DAC", dac_tile, "InterpolationFactor")
adc = report(rfdc.adc_tiles, st['ADCTileStatus'], "ADC", adc_tile, "DecimationFactor")

d = only_block(dac[0], "DAC tile 0")
a = only_block(adc[2], "ADC tile 2")

dac_fs, dac_interp = dac[0]['fs_msps'], d['factor']
adc_fs, adc_decim  = adc[2]['fs_msps'], a['factor']

adc_fabric = adc_fs / adc_decim
dac_fabric = dac_fs / dac_interp

print(f"\nDAC: {dac_fs:.4f} MSps / {dac_interp}x  -> {dac_fabric:.4f} MSps fabric")
print(f"ADC: {adc_fs:.4f} MSps / {adc_decim}x  -> {adc_fabric:.4f} MSps fabric")

#
# compare to parameters in the first cell
#
def check(name, live, expected, rtol=1e-6):
    if expected is None:
        print(f"{name}: no reference value")
    elif abs(live - expected) <= rtol * abs(expected):
        print(f"{name} is consistent ({live:.6g})")
    else:
        print(f"================> {name} is INCONSISTENT!!! "
              f"live={live:.6g}  expected={expected:.6g}")

check("DAC sampling", dac_fs,     rfdc_dac_sampling)
check("DAC interpolation", dac_interp, rfdc_dac_interpolation)
check("ADC sampling", adc_fs,     rfdc_adc_sampling)
check("ADC decimation", adc_decim, rfdc_adc_decimation)


# %%
"""
set up the GPIOs for the RX_FIFO_RESET, RX_FIFO_TLAST over DMA, and vanilla DMA
"""
#
# setup for the RX_FIFO_RESET
#
rx_fifo_reset_ip = gen.ip_dict["RX_FIFO_RESET"]   # this setups up the pointer to the address in the dictionary
rx_fifo_reset = AxiGPIO(rx_fifo_reset_ip).channel1 # we are using channel 1 of the 2 channels
def RX_FIFO_RESET():
    rx_fifo_reset.write(0,0x1)
def RX_FIFO_RELEASE():
    rx_fifo_reset.write(1,0x1)
#
# now for TLAST
#
tlast_ip = gen.ip_dict["RX_FIFO_TLAST"]
tlast = AxiGPIO(tlast_ip).channel1
def SET_TLAST(n):
    tlast.write(n,0xFFFFFF)
#
# setup for DMA 
#
dma_tx = gen.AXI_DMA_SG
dma_rx = gen.AXI_DMA_NOSG

def DMA_SS2M_SOFT_RESET():
    dma_rx.register_map.S2MM_DMACR.Reset = 1
    t0 = time.time()
    print("DMA_SS2M_SOFT_RESET: S2MM status register: "+hex(dma_rx.register_map.S2MM_DMASR))
    while dma_rx.register_map.S2MM_DMACR.Reset:
        if time.time() - t0 > 0.1:
            raise RuntimeError("DMA soft reset did not complete!")
        time.sleep(0.001)
    print("S2MM control register: "+hex(dma_rx.register_map.S2MM_DMACR))
    print("S2MM status register: "+hex(dma_rx.register_map.S2MM_DMASR))

#
# set the RX_FIFO to be in the reset state
#
RX_FIFO_RESET()
#RX_FIFO_RELEASE()
#
# reset the S2MM DMA just to be sure we are in a known state
#
DMA_SS2M_SOFT_RESET()

# %%
"""
set the mixer frequencies for the ADC and DAC.  these should be the same for a loopback.  
100 MHz is a safe number, far enough away from the baluns limit of round 20 MHz
"""

gen.freq_adc = 100
gen.freq_dac = 100

# %%
"""
check for free space. summer 2026 changed one of the 4x2 to 500G...
"""
# !cat /proc/meminfo | grep -i cma

# %%
"""
Allocate buffers for the DAC.

data_size are actually number of 16 bit words, for I and for Q. this number has to be smaller than the
number of free bytes from the command above.

The maximum DMA is 67 MB (2^26) but for DMA scatter-gather it can do a series of transfers.

Make the send and receive buffers. We make want some number of samples, with I and Q each as 16 bits. 
If we define 16 bit buffers, and put the I on the even addresses and Q on the odd addresses, 
then when the buffer is sent over DMA, since we've configured that for 32 bit "beats", 
it will put them together correctly.

set the number of data points we want to send to be e an even multiple of 2. 
then data_size will be twice that since we are interleaving I and Q
"""
#npoints = 2**23
npoints = 2**21
data_size = 2*npoints 
print(f"Array will hold {npoints} = {npoints/1e6:.2f}M data points")
send_DAC_buffer = allocate(shape=(data_size,), dtype=np.int16)
DAC_channel = gen.AXI_DMA_SG.sendchannel


# %%
"""
now make an array that has the waveform, and take care that it fits nicely into "npoints" 
so there's no phase discontinuity
"""

f_DAC = 1.0E6*rfdc_dac_sampling
print(f"DAC sampling at {f_DAC/1E9:.6f} GSps")
fs = f_DAC/rfdc_dac_interpolation
print(f"DAC 'fabric' rate at {fs/1e6:.3f} MHz")
f_mod = 5.0e6   # in MHz
print(f"modulation tone at {f_mod:.8f} Hz = {f_mod/1e6:.1f} MHz")
#
# modify the modulation tone so there are an integer number of periods in the array to avoid phase discontinuties when we cycle it into the DAC
#
m = round(f_mod*npoints/fs)
print("Number of cycles is set to "+str(m))
f = m*fs/npoints
print(f"Modified modulation tone at {f:.8f} Hz = {f/1e6:.1f} MHz,  differs by {(f_mod-f):.3f} Hz")
#
# now make an i and q
#
duration = npoints/fs
print(f"Array duration is {duration:.7f} sec")
ampl = 0.9
n     = np.arange(npoints, dtype=np.float64)
start = time.perf_counter()
phase = (2*np.pi*m/npoints)*n              # exact, no float drift from t
i_a   = ampl*np.sin(phase)
q_a   = ampl*np.cos(phase)
#i_a = [ampl] * npoints
#q_a = [ampl] * npoints
stop = time.perf_counter()
print(f"Took {stop-start:.3f} seconds to fill both arrays")

# %%
"""
pack send_DAC_buffer[::2] to copy bytes to even addresses, and 1::2 to copy to odd, since this is how the DCU expects I and Q data.

and send the buffer into memory
"""

send_DAC_buffer[::2]  = np.rint(i_a*(2**15-1)).astype(np.int16)
send_DAC_buffer[1::2] = np.rint(q_a*(2**15-1)).astype(np.int16)

# %%
"""
start the DAC
"""

tile = gen.RFDC.dac_tiles[dac_tile_n]
gen.dac_hard_stop()               # bounded stop, never hangs
#
# reset the internal FIFOs in the RF Converter DAC side, so it's in a known state
#
tile.SetupFIFO(False)             # safe ONLY because MM2S is halted
tile.SetupFIFO(True)              # bring the tile FIFO up clean
send_DAC_buffer.flush()
DAC_channel.transfer(send_DAC_buffer, cyclic=True)

status = gen.dac_status()
if status['ok']:
    print("DAC seems to be running...")
    oled.write("DAC DMA cyclic on")
else:
    print("Problems with the DAC status.  Here is a printout:")
    print(status)
    oled.write("DAC PROBLEM!!!")

# %%
"""
Get the ADC data.  This cell loops a few seconds to make sure the transfer completed, which means
that the TLAST fix worked!
"""

#
# Grabs 2**16 I/Q pairs instead of what the DAC uses (e.g. 2**21 or so)
# At that size the whole spectrum is only 65,536 bins, so it plots directly 
# with no decimation and no envelope trickery.
#
# RBW = 122.88 MHz / 65536 = 1.875 kHz, capture duration = 533 us.
#
N_CAP = 2**16          # I/Q pairs to capture
SAMPLE_SUBSET = 200    # samples drawn in the time-domain panel
N_PEAKS = 6            # rows in the peak table

# ---------------------------------------------------------------- capture --
# release any buffer from a previous run so repeated executions don't eat CMA
# use try/except because this buffer is defined below, so the first time you execute
# it, it will bomb unless you use try/except
#
try:
    recv_ADC_buffer.freebuffer()
except NameError:
    pass
#
# reset RX_FIFO so the FIFO is empty, and set up for the TLAST sovAXI_DMA_NOSG 
# knows when the right number of beats have happened and the DMA ends
#
RX_FIFO_RESET()
SET_TLAST(N_CAP)
print("TLAST target set to " + hex(tlast.read()))
#
# allocate the receive buffer and setup the channel for DMA, start the DMA so that
# the state machines are ready, and send the address using .transfer.  
#
# Note that at this point, no DMA will happen because the FIFO is reset, which means it's empty,
# and this will keep the AXIS stream bus in a wait mode until data can state flowing
#
recv_ADC_buffer = allocate(shape=(2 * N_CAP,), dtype=np.int16)
ADC_channel = gen.AXI_DMA_NOSG.recvchannel
ADC_channel.start()
ADC_channel.transfer(recv_ADC_buffer)
print("S2MM DMASR armed: " + hex(dma_rx.register_map.S2MM_DMASR))
#
# release the FIFO reset, and data will start flowing
#
RX_FIFO_RELEASE()
blk = rfdc.adc_tiles[2].blocks[1]
print("BlockStatus: ",blk.BlockStatus)   # DataPathClocksStatus, FIFOFlagsAsserted
#
# wait for the DMA to finish.  instead of a wait, which is blocking, we set the maximum
# time we will need.   we know the clock frequency (ps_clock) and we know the number of
# words so this is easy to calculate:
#
total_time = max(1.0,N_CAP*1e-6/ps_clock)
print("Will wait ",total_time," seconds for the data")
#
to = time.time()
ok = False
while True:
    if dma_rx.register_map.S2MM_DMASR.Idle:
        ok = True
        break
    if time.time() - to > total_time:
        print("=== S2MM TIMEOUT ===")
        print("DMASR:", dma_rx.register_map.S2MM_DMASR)
        print("BlockSttus after timeout:", blk.BlockStatus)
        break
    time.sleep(0.01)
#
# Reset the FIFO so we are in a known state
#
RX_FIFO_RESET()
#
# this next command has to do with the fact that writing INTO DDR from the FPGA
# does not go through the ARM chip, but directly to DDR through the HP ports in
# the PS.  but the ARM chip might still hold state lines for some of those memory
# locations you allocated (and would hold zeros).  calling invalidate() for the
# receive buffer marks them dead so the next numpy read refetches and doesn't use
# ARM cache.   
#
recv_ADC_buffer.invalidate()
#
# check if we got anything.  if so, do the analysis
nonzero =  np.count_nonzero(recv_ADC_buffer)
print("nonzero words:", nonzero, "of", len(recv_ADC_buffer))
if nonzero == 0:
    print("Something happened, no data arrived.  Bail!")
else:
    print("ADC data arrived ok")

# %%
"""
    Analyze!

    this next command has to do with the fact that writing INTO DDR from the FPGA 
    does not go through the ARM chip, but directly to DDR through the HP ports in 
    the PS.  but the ARM chip might still hold state lines for some of those memory 
    locations you allocated (and would hold zeros).  calling invalidate() for the
    receive buffer marks them dead so the next numpy read refetches and doesn't use
    ARM cache.   

"""
recv_ADC_buffer.invalidate()
#
# check if we got anything.  if so, do the analysis
nonzero =  np.count_nonzero(recv_ADC_buffer)
print("nonzero words:", nonzero, "of", len(recv_ADC_buffer))
if nonzero == 0:
    print("Something happened, no data arrived.  Bail!")
else:
    #
    # de-interleave the received I and Q from recv_ADC_buffer and form the real data
    #
    i_rx = recv_ADC_buffer[0::2].astype(np.float64)
    q_rx = recv_ADC_buffer[1::2].astype(np.float64)
    iq_data = i_rx + 1j * q_rx
    #
    # type out some useful information
    #
    N = len(iq_data)
    f_fabric = rfdc_adc_sampling / rfdc_adc_decimation           # MHz
    rbw_hz = f_fabric * 1e6 / N
    print(f"\nCaptured {N} pairs at {f_fabric} MHz -> {N/f_fabric:.1f} us, "f"RBW = {rbw_hz:.1f} Hz")
    print(f"I range [{i_rx.min():.0f}, {i_rx.max():.0f}]   "
      f"Q range [{q_rx.min():.0f}, {q_rx.max():.0f}]   "
      f"(full scale +/-32767)")
    #
    # FFT with a hanning window since we are reading a chunck of continuous
    # data, and without a window we'd see the spectral leakate from the endpoints.
    # Also, shift the FFT spectrum to go from -Nyquist to +Nyquist
    #
    window = np.hanning(N)
    fft_shifted = np.fft.fftshift(np.fft.fft(iq_data * window))
    fft_db = 20 * np.log10(np.abs(fft_shifted) / np.abs(fft_shifted).max() + 1e-15)
    freqs = np.fft.fftshift(np.fft.fftfreq(N, d=1.0 / f_fabric))   # MHz
    #
    # make a table of the prominent peaks.  we want to find the peaks without accidentally
    # counting theshoulders or skirts of a single peak as a separate distinct peak
    #
    # rbw_hz is the frequency resolution, guard is the equivalent number 
    guard = max(1, int(0.1e6 / rbw_hz))
    print("guard = ",guard)
    work = fft_db.copy()
    print(f"\n  {'freq (MHz)':>12}   {'dBc':>7}")
    for _ in range(N_PEAKS):
        k = int(np.argmax(work))
        if not np.isfinite(work[k]):
            break
        print(f"  {freqs[k]:12.4f}   {fft_db[k]:7.2f}")
        work[max(0, k - guard):k + guard] = -np.inf

    # ------------------------------------------------------------------ plot ---
    fig = make_subplots(
        rows=2, cols=1,
        subplot_titles=(
            f"ADC Time-Domain Signal (First {SAMPLE_SUBSET} Samples)",
            f"Baseband Spectrum -- all {N:,} bins "
            f"(f_Fabric = {f_fabric:.2f} MHz, RBW = {rbw_hz:.0f} Hz)",
        ),
        vertical_spacing=0.12,
    )

    fig.add_trace(go.Scatter(y=i_rx[:SAMPLE_SUBSET], mode="lines",
                         name="In-Phase (I)",
                         line=dict(color="rgb(31, 119, 180)")), row=1, col=1)
    fig.add_trace(go.Scatter(y=q_rx[:SAMPLE_SUBSET], mode="lines",
                         name="Quadrature (Q)",
                         line=dict(color="rgb(255, 127, 14)")), row=1, col=1)

    fig.add_trace(go.Scattergl(x=freqs, y=fft_db, mode="lines",
                           name="Spectrum (dBc)",
                           line=dict(color="rgb(214, 39, 40)", width=1)),
              row=2, col=1)

    fig.update_xaxes(title_text="Sample Index", row=1, col=1)
    fig.update_yaxes(title_text="Raw ADC (int16)", row=1, col=1)
    fig.update_xaxes(title_text="Offset Frequency (MHz)", row=2, col=1)

    floor = np.floor((fft_db.min() - 5) / 10) * 10        # from the data, not -100
    fig.update_yaxes(title_text="Normalized Amplitude (dBc)",
                 range=[floor, 5], row=2, col=1)

    fig.update_layout(height=800, width=950,
                  title_text="RFSoC ADC Capture Analysis",
                  template="plotly_white", showlegend=True)
    fig.show()

# %%
