"""
02_spectrum_fft.py — Спектр (БПФ) и спектрограмма в реальном времени.

Раскладывает звук на частотные составляющие преобразованием Фурье — см. THEORY.md,
раздел "Спектр и преобразование Фурье". Свист даёт один острый пик, голос — несколько
(форманты), хлопок — широкий "размазанный" спектр.

Запуск:
    python3 02_spectrum_fft.py [/dev/ttyUSB0]
"""
import sys

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

from serial_stream import open_port, read_int_lines, EFFECTIVE_RATE_HZ

PORT = sys.argv[1] if len(sys.argv) > 1 else "/dev/ttyUSB0"

WINDOW_SAMPLES = 2048          # окно анализа
SPECTROGRAM_COLUMNS = 150      # сколько последних БПФ-кадров хранит спектрограмма


def main():
    ser = open_port(PORT)
    print(f"Подключено к {PORT}. Найквист (верхняя граница спектра): {EFFECTIVE_RATE_HZ / 2:.0f} Гц.")
    print("Если график не двигается — отправьте 'p' в Serial Monitor (режим Plotter).")

    waveform = np.zeros(WINDOW_SAMPLES)
    spectrogram = np.zeros((WINDOW_SAMPLES // 2 + 1, SPECTROGRAM_COLUMNS))
    freqs = np.fft.rfftfreq(WINDOW_SAMPLES, d=1.0 / EFFECTIVE_RATE_HZ)
    hann_window = np.hanning(WINDOW_SAMPLES)

    fig, (ax_spec, ax_spectrogram) = plt.subplots(2, 1, figsize=(10, 7))
    fig.suptitle("ICS-43434 — спектр (БПФ) и спектрограмма")

    (spec_line,) = ax_spec.plot(freqs, np.zeros_like(freqs), lw=1, color="darkorange")
    ax_spec.set_xlim(0, EFFECTIVE_RATE_HZ / 2)
    ax_spec.set_ylim(0, 1)
    ax_spec.set_xlabel("Частота, Гц")
    ax_spec.set_ylabel("|БПФ|, отн. ед.")
    ax_spec.set_title(f"Спектр (окно {WINDOW_SAMPLES} сэмплов, Найквист {EFFECTIVE_RATE_HZ/2:.0f} Гц)")
    ax_spec.grid(alpha=0.3)

    spectrogram_img = ax_spectrogram.imshow(
        spectrogram, origin="lower", aspect="auto", cmap="magma",
        extent=[-SPECTROGRAM_COLUMNS, 0, 0, EFFECTIVE_RATE_HZ / 2],
        vmin=0, vmax=1,
    )
    ax_spectrogram.set_xlabel("Кадры БПФ назад")
    ax_spectrogram.set_ylabel("Частота, Гц")
    ax_spectrogram.set_title("Спектрограмма — как спектр меняется во времени")

    fig.tight_layout(rect=[0, 0, 1, 0.95])

    spec_smoothed_max = 1.0

    def update(_frame):
        nonlocal waveform, spec_smoothed_max
        values = read_int_lines(ser)
        if not values:
            return spec_line, spectrogram_img

        n_new = min(len(values), WINDOW_SAMPLES)
        waveform = np.roll(waveform, -n_new)
        waveform[-n_new:] = values[-n_new:]

        spectrum = np.abs(np.fft.rfft(waveform * hann_window))
        current_max = max(spectrum.max(), 1.0)
        spec_smoothed_max = 0.9 * spec_smoothed_max + 0.1 * current_max
        spec_norm = spectrum / spec_smoothed_max
        spec_line.set_ydata(spec_norm)
        ax_spec.set_ylim(0, max(1.0, spec_norm.max() * 1.1))

        spectrogram[:, :-1] = spectrogram[:, 1:]
        spectrogram[:, -1] = np.clip(spec_norm, 0, 1)
        spectrogram_img.set_data(spectrogram)

        return spec_line, spectrogram_img

    FuncAnimation(fig, update, interval=50, blit=False, cache_frame_data=False)
    plt.show()
    ser.close()


if __name__ == "__main__":
    main()