# Description: Reads FFT magnitudes from UART and plots the values in real time
#              using a maplotlib animation
#              The values received are assumed to be from a real FFT and only
#              frequencies from DC to Nyquist are included (the other half of
#              of the frequency domain is the complex conjugate, so no
#              no additional information is gained)
# Author: DAD Projects
# Date: 9/18/2024

import struct
import serial
import numpy as np
from matplotlib import pyplot as plt
import matplotlib.animation as animation
from threading import Thread

# Values must match the microcontroller
FFT_LENGTH = 2048
FFT_LENGTH_HALF = int(FFT_LENGTH/2)
SAMP_RATE = 20480

# Initialize the serial object
# Parameters must also match the microcontroller's UART parameters
ser = serial.Serial('COM6', 921600, timeout=1000,
                    parity=serial.PARITY_NONE, stopbits=serial.STOPBITS_ONE)

# Plotting up to Nyquist frequency
freqs = np.linspace(0, SAMP_RATE/2, FFT_LENGTH_HALF, endpoint=False)
# freqs = np.linspace(0, FFT_LENGTH_HALF-1, FFT_LENGTH_HALF)
fft_mags = np.zeros(FFT_LENGTH_HALF)
fig = plt.figure(figsize=(8, 4))
ax = plt.subplot(111)
line, = ax.plot(freqs, fft_mags)
ax.set_ylim(0, 1200)  # Revisit the limits
ax.set_xlim(np.min(freqs), np.max(freqs))


def read_uart():
    global fft_mags

    while True:
        # Read until packet termination xFFxFFxFFxFF
        uart_vals = ser.read_until(b'\xFF\xFF\xFF\xFF')

        if uart_vals.__len__() == FFT_LENGTH_HALF*4 + 4:
            fft_mags = struct.unpack(f"{FFT_LENGTH_HALF}f",
                                     np.flip(uart_vals[0:FFT_LENGTH_HALF*4]))


def animate(frame):
    global fft_mags
    line.set_ydata(fft_mags)
    print(f"Animating frame {frame}")
    return line,


# Refreshes the plot every 20ms
ani = animation.FuncAnimation(
    fig, animate, interval=200, blit=True, cache_frame_data=False)

if __name__ == "__main__":
    uart_thread = Thread(target=read_uart, daemon=True)
    uart_thread.start()
    plt.show()
    print("thread finished...exiting")