import serial
import struct
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import sys

# --- KONFIGURACIJA---
SERIAL_PORT = '/dev/ttyACM0'   # Preveri svoj port
BAUD_RATE = 115200     
FFT_SIZE = 1024
BIN_COUNT = FFT_SIZE // 2  # 512 Bins
DISPLAY_BINS = 100  # Only display first 100 bins
SWEEP_SIZE = 2
BYTES_PER_FRAME = BIN_COUNT * 4
WATERFALL_HISTORY = 100  # Number of frames to display in waterfall

try:
    ser = serial.Serial(SERIAL_PORT, BAUD_RATE, timeout=2)
    ser.reset_input_buffer()
    print("Povezava z Arduino Due vzpostavljena")
except Exception as e:
    print(f"Error: {e}")
    sys.exit()

# Setup Plot with two subplots
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 10))

x_freq = np.zeros(DISPLAY_BINS)
y_data = np.zeros(DISPLAY_BINS)

# FFT plot (top)
bar_plot, = ax1.plot(x_freq, y_data, color="#ff0000", lw=1)
ax1.set_xlabel('Razdalja [m]')
ax1.set_ylabel('Amplituda')

# Waterfall plot (bottom)
waterfall_data = np.zeros((WATERFALL_HISTORY, DISPLAY_BINS))
waterfall_img = ax2.imshow(waterfall_data, aspect='auto', cmap='viridis', 
                           interpolation='nearest', origin='upper')
ax2.set_xlabel('Razdalja [m]')
ax2.set_ylabel('Čas (frame)')
ax2.set_title('Waterfall prikaz')
cbar = plt.colorbar(waterfall_img, ax=ax2, label='Amplituda')

def update(frame):
    global waterfall_data
    
    if ser.in_waiting >= BYTES_PER_FRAME:
        raw_data = ser.read(BYTES_PER_FRAME)
        
        try:
            fm_sweep = [2]

            x_dist = np.linspace(0, 1024*3e8/100e6/2, BIN_COUNT)
            x_dist_display = x_dist[:DISPLAY_BINS]  # Only first 100 bins
            ax1.set_title(f'FFT prikaz radarskega signala FM={2}')
            bar_plot.set_xdata(x_dist_display)
            ax1.set_xlim(x_dist_display[2], x_dist_display[-1]) # Skip DC

            fft_data = struct.unpack(f'{BIN_COUNT}f', raw_data[:])
            fft_data_display = fft_data[:DISPLAY_BINS]  # Only use first 100 bins

            # Update FFT plot
            bar_plot.set_ydata(fft_data_display)

            # Auto-scale FFT plot
            peak = np.max(fft_data_display[2:]) 
            if peak > ax1.get_ylim()[1]:
                ax1.set_ylim(0, peak * 1.2)
            elif peak < ax1.get_ylim()[1] * 0.5 and peak > 10:
                ax1.set_ylim(0, peak * 1.5)
            
            # Update waterfall (scroll up and add new data at bottom)
            waterfall_data = np.roll(waterfall_data, -1, axis=0)
            waterfall_data[-1, :] = fft_data[:DISPLAY_BINS]
            
            waterfall_img.set_data(waterfall_data)
            waterfall_img.set_clim(vmin=np.min(waterfall_data), vmax=np.max(waterfall_data))
            
            # Update waterfall x-axis to match distance
            waterfall_img.set_extent([x_dist_display[0], x_dist_display[-1], WATERFALL_HISTORY, 0])
                
        except Exception as e:
            print(f"Frame Error: {e}")
            ser.reset_input_buffer()

    return bar_plot, waterfall_img

ani = FuncAnimation(fig, update, interval=0, blit=True)
plt.tight_layout()
plt.show()