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

# --- НАСТРОЙКИ ---
PORT = 'COM3'
BAUD = 921600
FS = 1000       
SAMPLES = 512   
PACKET_SIZE = 12  # Изменено: 2(AA BB) + 2(gx) + 2(fgx) + 2(f1) + 2(f2) + 2(f3)

try:
    ser = serial.Serial(PORT, BAUD, timeout=0.01)
    ser.set_buffer_size(rx_size=12800, tx_size=12800)
    ser.reset_input_buffer()
    print(f"Подключено к {PORT}. Ожидание данных (12-байтные пакеты)...")
except Exception as e:
    print(f"Ошибка порта: {e}")
    exit()

# Буферы данных
raw_buffer = np.zeros(SAMPLES)
filt_buffer = np.zeros(SAMPLES)
peak_freqs = [0, 0, 0] # Для хранения частот из STM32
data_packet = bytearray()

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))

# 1. Временной график
line_raw, = ax1.plot(raw_buffer, color='red', lw=2.5, alpha=0.3, label='Грязный (Raw)')
line_filt, = ax1.plot(filt_buffer, color='blue', lw=1.2, label='Фильтр')  # Обновлено название
ax1.set_title("Осциллограф: Гироскоп X")
ax1.set_ylim(-2000, 2000)
ax1.legend(loc='upper right')
ax1.grid(True, alpha=0.3)

# Текст для отображения частот на графике
peak_text = ax1.text(0.02, 0.95, '', transform=ax1.transAxes, verticalalignment='top', 
                     bbox=dict(boxstyle='round', facecolor='white', alpha=0.7), fontsize=12, color='darkgreen')

# 2. Частотный график
freqs = np.fft.rfftfreq(SAMPLES, 1/FS)
line_spec_filt, = ax2.plot(freqs, np.zeros(len(freqs)), color='blue', lw=1.5, label='Спектр (После фильтров)', zorder=10)
fill_poly = None 

# Вертикальные линии для визуализации работы Notch-фильтров
peak_lines = [ax2.axvline(0, color='green', linestyle='--', alpha=0.0) for _ in range(3)]

ax2.set_title("Частотный анализ (FFT)")
ax2.set_xlabel("Частота (Гц)")
ax2.set_xlim(0, 500)
ax2.set_ylim(0, 30000)
ax2.xaxis.set_major_locator(MultipleLocator(50))
ax2.grid(True, alpha=0.3)

def update(frame):
    global raw_buffer, filt_buffer, data_packet, fill_poly, peak_freqs
    
    new_data = ser.read_all()
    if new_data:
        data_packet.extend(new_data)
    
    # Парсинг пакетов (длина 12)
    while len(data_packet) >= PACKET_SIZE:
        header_idx = data_packet.find(b'\xaa\xbb')
        if header_idx == -1:
            data_packet.clear()
            break
        
        if header_idx + PACKET_SIZE <= len(data_packet):
            # Читаем: 2(gx), 2(fgx), 2(f1), 2(f2), 2(f3)
            chunk = data_packet[header_idx+2 : header_idx+PACKET_SIZE]
            try:
                # <hhHHH -> int16, int16, uint16, uint16, uint16
                r_val, f_val, f1, f2, f3 = struct.unpack('<hhHHH', chunk)
                
                raw_buffer[:-1] = raw_buffer[1:]
                raw_buffer[-1] = r_val
                filt_buffer[:-1] = filt_buffer[1:]
                filt_buffer[-1] = f_val
                
                peak_freqs = [f1, f2, f3]
            except Exception as e:
                print(f"Unpack error: {e}")
            
            del data_packet[:header_idx+PACKET_SIZE]
        else:
            break

    # Обновление временных линий
    line_raw.set_ydata(raw_buffer)
    line_filt.set_ydata(filt_buffer)
    
    # Обновление текста с частотами
    peak_text.set_text(f"Notch 1: {peak_freqs[0]} Hz\nNotch 2: {peak_freqs[1]} Hz\nNotch 3: {peak_freqs[2]} Hz")
    
    # FFT
    sig_r = raw_buffer - np.mean(raw_buffer)
    sig_f = filt_buffer - np.mean(filt_buffer)
    fft_r = np.abs(np.fft.rfft(sig_r))
    fft_f = np.abs(np.fft.rfft(sig_f))
    
    line_spec_filt.set_ydata(fft_f)
    
    # Обновление вертикальных линий (показываем где фильтры)
    for i in range(3):
        if peak_freqs[i] > 0:
            peak_lines[i].set_xdata([peak_freqs[i], peak_freqs[i]])
            peak_lines[i].set_alpha(0.6)
        else:
            peak_lines[i].set_alpha(0.0)

    if fill_poly:
        fill_poly.remove()
    fill_poly = ax2.fill_between(freqs, 0, fft_r, color='red', alpha=0.15, zorder=1)

    # Автомасштаб спектра
    #m = np.max(fft_r) if len(fft_r) > 0 else 5000
    #curr_limit = ax2.get_ylim()[1]
    #if m > curr_limit:
    #    ax2.set_ylim(0, m * 1.2)
    #elif m < curr_limit * 0.2 and curr_limit > 5000:
    #    ax2.set_ylim(0, 5000)

    return line_raw, line_filt, line_spec_filt, peak_text

ani = FuncAnimation(fig, update, interval=30, cache_frame_data=False)
plt.tight_layout()
plt.show()
ser.close()