import uhd
import argparse
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import cmath
import scipy
from scipy import signal
matplotlib.rcParams['figure.figsize'] = [15,5]

graph = uhd.rfnoc.RfnocGraph('addr = 192.168.10.2')

FFT_noc_block = graph.get_block('0/FFT#0')
## FFT BLOCK CONFIGURATION
# fft_length = 512
# fft_magnitude = uhd.libpyuhd.rfnoc.fft_magnitude.COMPLEX
# fft_direction = uhd.libpyuhd.rfnoc.fft_direction.FORWARD
# fft_shift = uhd.libpyuhd.rfnoc.fft_shift.NORMAL

# fft = uhd.rfnoc.FftBlockControl(FFT_noc_block)
# fft.set_length(fft_length)
# fft.set_magnitude(fft_magnitude)
# fft.set_direction(fft_direction)
# fft.set_shift_config(fft_shift)

# Stream Arguments
sa = uhd.usrp.StreamArgs('fc32','sc16')

radio_noc_block = graph.get_block("0/Radio#0")


graph.connect("0/Radio#0", 0, "0/DDC#0", 0, True)
rx_streamer = graph.create_rx_streamer(1, sa)

graph.connect("0/DDC#0", 0, '0/FFT#0', 0, True)
graph.connect('0/FFT#0', 0, rx_streamer, 0, True)
graph.commit()


fft_length = 512
fft_magnitude = uhd.libpyuhd.rfnoc.fft_magnitude.COMPLEX
fft_direction = uhd.libpyuhd.rfnoc.fft_direction.FORWARD
fft_shift = uhd.libpyuhd.rfnoc.fft_shift.NORMAL

fft = uhd.rfnoc.FftBlockControl(FFT_noc_block)
fft.set_length(fft_length)
fft.set_magnitude(fft_magnitude)
fft.set_direction(fft_direction)
fft.set_shift_config(fft_shift)

sa.args = "spp=" + str(fft_length)

radio_block = uhd.rfnoc.RadioControl(radio_noc_block)

radio_block.set_rx_frequency(853e6, 0)
radio_block.set_rx_gain(65, 0)
radio_block.set_rx_antenna("RX1", 0)
radio_block.set_rate(200e6)
radio_block.set_properties(f'spp={fft_length}', 0)

ddc_block = uhd.rfnoc.DdcBlockControl(graph.get_block("0/DDC#0"))
ddc_block.set_input_rate(200e6, 0)
ddc_block.set_output_rate(5e6, 0)




num_samples = fft_length  # Number of samples to receive at a time
radio_data = np.zeros((1, num_samples), dtype="complex64")  # Keep the buffer outside

# Initialize streaming command for continuous data reception
stream_cmd = uhd.types.StreamCMD(uhd.types.StreamMode.start_cont)
stream_cmd.stream_now = True
rx_streamer.issue_stream_cmd(stream_cmd)  # Start streaming

# Set up the plot
fig, ax = plt.subplots()
x_data = np.linspace(-0.5, 0.5, fft_length)
y_data = np.zeros(fft_length)
line, = ax.semilogy(x_data, np.abs(y_data))  # Initial empty plot

# Update function for the animation
def update_plot(frame):
    metadata = uhd.types.RXMetadata()

    # Receive new samples from the radio
    num_samples_received = rx_streamer.recv(radio_data, metadata, timeout=0.5)

    # Check for overflows or underflows
    if metadata.error_code == uhd.types.RXMetadataErrorCode.none:
        # Ensure that valid samples are received
        if num_samples_received > 0:
            # Update the y-data with the magnitude of the newly received samples
            line.set_ydata(np.abs(radio_data[0]))  # Update y-axis with new radio data
            
            # Adjust plot to fit new data if necessary
            ax.relim()
            ax.autoscale_view()
    else:
        print(f"Metadata error: {metadata.error_code}")

    return line,
while True: 
    # Set up the animation to update the plot continuously
    ani = animation.FuncAnimation(fig, update_plot, interval=100, blit=False)  # Update every 100 ms
    # Display the plot
    plt.show()


# Stop streaming when done (optional)
rx_streamer.issue_stream_cmd(uhd.types.StreamCMD(uhd.types.StreamMode.stop_cont))