"""
Визуализатор для Audio_INMP441: осциллограмма, спектр (БПФ) и спектрограмма
звука в реальном времени. Читает поток чисел из Serial Plotter-режима ESP32
(main.cpp, режим 'p' — см. README.md).

Запуск:
    pip install -r requirements.txt
    python3 visualizer.py [/dev/ttyUSB0]

Ось Y осциллограммы и подписи амплитуды — в LSB 16-битной шкалы (-32768..32767),
БЕЗ калибровки в дБ SPL: у INMP441 без заводского эталона нет способа перевести
отсчёт АЦП в паскали. Это осознанное ограничение, а не недоработка — см. THEORY.md,
раздел "Почему нет абсолютных дБ".
"""
import sys
import time
from collections import deque

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

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

# Должны совпадать с прошивкой: SAMPLE_RATE / kPlotDecimation в main.cpp.
MIC_SAMPLE_RATE_HZ = 16000
PLOT_DECIMATION = 8
EFFECTIVE_RATE_HZ = MIC_SAMPLE_RATE_HZ // PLOT_DECIMATION  # частота потока, реально идущего по Serial

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


def read_int_lines(ser, max_lines=4000):
    """Забирает всё, что накопилось на порту, и парсит числовые строки."""
    values = []
    if ser.in_waiting == 0:
        return values
    raw = ser.read(ser.in_waiting).decode("ascii", errors="ignore")
    for line in raw.strip().split("\n")[-max_lines:]:
        line = line.strip()
        try:
            values.append(int(line))
        except ValueError:
            continue  # строки метра громкости/диагностики — не число, пропускаем
    return values


def main():
    try:
        ser = serial.Serial(PORT, BAUD, timeout=0.05)
    except serial.SerialException as e:
        print(f"Не удалось открыть {PORT}: {e}")
        print("Проверьте порт (ls /dev/ttyUSB*) и что main.cpp переключён в режим 'p' (Serial Plotter).")
        sys.exit(1)

    print(f"Подключено к {PORT} @ {BAUD}. Эффективная частота потока: {EFFECTIVE_RATE_HZ} Гц.")
    print("Если график не двигается — отправьте 'p' в Serial Monitor, чтобы включить режим Plotter.")

    waveform = deque([0] * WINDOW_SAMPLES, maxlen=WINDOW_SAMPLES)
    spectrogram = np.zeros((WINDOW_SAMPLES // 2 + 1, SPECTROGRAM_COLUMNS))

    freqs = np.fft.rfftfreq(WINDOW_SAMPLES, d=1.0 / EFFECTIVE_RATE_HZ)
    time_axis = np.arange(-WINDOW_SAMPLES, 0) / EFFECTIVE_RATE_HZ

    fig, (ax_wave, ax_spec, ax_spectrogram) = plt.subplots(3, 1, figsize=(10, 8))
    fig.suptitle("INMP441 — звук в реальном времени (осциллограмма / спектр / спектрограмма)")

    (wave_line,) = ax_wave.plot(time_axis, list(waveform), lw=1)
    ax_wave.set_xlim(time_axis[0], time_axis[-1])
    ax_wave.set_ylim(-32768, 32767)
    ax_wave.set_xlabel("Время, с")
    ax_wave.set_ylabel("Амплитуда, LSB")
    ax_wave.set_title("Осциллограмма — давление воздуха на мембране во времени")
    ax_wave.grid(alpha=0.3)

    (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"Спектр (Найквист {EFFECTIVE_RATE_HZ / 2:.0f} Гц — выше не увидим, см. THEORY.md)")
    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.96])

    hann_window = np.hanning(WINDOW_SAMPLES)
    spec_smoothed_max = 1.0  # для плавной авто-нормировки шкалы спектра

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

        for v in values:
            waveform.append(v)

        data = np.array(waveform, dtype=np.float64)
        wave_line.set_ydata(data)

        spectrum = np.abs(np.fft.rfft(data * 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 wave_line, spec_line, spectrogram_img

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

    ser.close()


if __name__ == "__main__":
    main()