"""
03_tuning_fork.py — Камертон под микроскопом: точная частота и биения.

Два независимых режима одного скрипта (переключаются флагом --mode):

  pitch  (по умолчанию) — непрерывно определяет основную частоту звучащего
         камертона алгоритмом pYIN (librosa) и показывает ближайшую музыкальную
         ноту. Метод устойчивее к шуму, чем поиск пика в БПФ, потому что ищет
         именно периодичность сигнала, а не самый громкий частотный бин.

  beats  — запись двух РЕАЛЬНЫХ источников по очереди (например, два камертона
         разного строя, или один камертон и ноту с телефона) и построение
         биений — как реально складываются их волны, а не симуляция синусоид.
         Физика — THEORY.md, раздел "Биения: что происходит, когда звучат два
         источника почти в унисон".

Запуск:
    python3 03_tuning_fork.py --mode pitch [/dev/ttyUSB0]
    python3 03_tuning_fork.py --mode beats [/dev/ttyUSB0]
"""
import argparse
import time

import numpy as np
import matplotlib.pyplot as plt

from serial_stream import open_port, read_block_blocking, EFFECTIVE_RATE_HZ

try:
    import librosa
    LIBROSA_OK = True
except ImportError:
    LIBROSA_OK = False

NOTE_NAMES_RU = ["До", "До#", "Ре", "Ре#", "Ми", "Фа", "Фа#", "Соль", "Соль#", "Ля", "Ля#", "Си"]


def freq_to_note_ru(f_hz: float) -> str:
    """Ближайшая нота при A4=440 Гц: n = 12*log2(f/440), номер MIDI = 69+n."""
    if f_hz <= 0:
        return "—"
    midi = 69 + 12 * np.log2(f_hz / 440.0)
    midi_round = int(round(midi))
    cents_off = (midi - midi_round) * 100  # отклонение от идеальной ноты, в центах
    name = NOTE_NAMES_RU[midi_round % 12]
    octave = midi_round // 12 - 1
    return f"{name}{octave} ({cents_off:+.0f} центов)"


def run_pitch_mode(port: str):
    if not LIBROSA_OK:
        print("Нужна librosa: pip install librosa")
        return

    ser = open_port(port)
    print(f"Режим PITCH. Частота потока: {EFFECTIVE_RATE_HZ} Гц. Звучите камертоном рядом с микрофоном.")
    print("Ctrl+C для выхода.\n")

    fig, ax = plt.subplots(figsize=(7, 4))
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.axis("off")
    big_text = ax.text(0.5, 0.6, "—", ha="center", va="center", fontsize=42)
    sub_text = ax.text(0.5, 0.3, "", ha="center", va="center", fontsize=16, color="gray")
    fig.suptitle("Камертон: частота в реальном времени (pYIN, librosa)")
    plt.ion()
    fig.show()

    try:
        while plt.fignum_exists(fig.number):
            block = read_block_blocking(ser, n_samples=EFFECTIVE_RATE_HZ // 2, timeout_s=2.0)
            if block.size < EFFECTIVE_RATE_HZ // 4:
                plt.pause(0.05)
                continue

            block = block / 32768.0  # librosa ожидает float в диапазоне [-1, 1]
            f0, voiced_flag, _ = librosa.pyin(
                block, sr=EFFECTIVE_RATE_HZ,
                fmin=librosa.note_to_hz("C2"), fmax=librosa.note_to_hz("C7"),
            )
            voiced = f0[voiced_flag]
            if voiced.size == 0:
                big_text.set_text("—")
                sub_text.set_text("нет устойчивого тона")
            else:
                f_hz = float(np.median(voiced))
                big_text.set_text(f"{f_hz:.1f} Гц")
                sub_text.set_text(freq_to_note_ru(f_hz))

            fig.canvas.draw_idle()
            plt.pause(0.01)
    except KeyboardInterrupt:
        pass
    ser.close()


def run_beats_mode(port: str):
    ser = open_port(port)
    print(f"Режим BEATS. Частота потока: {EFFECTIVE_RATE_HZ} Гц.")
    print("По команде запишет 3 секунды звука — это время дайте прозвучать ОБОИМ источникам вместе.")
    input("Нажмите Enter, чтобы начать запись...")

    duration_s = 3.0
    n_samples = int(EFFECTIVE_RATE_HZ * duration_s)
    print("Запись...")
    block = read_block_blocking(ser, n_samples=n_samples, timeout_s=duration_s + 3.0)
    ser.close()

    if block.size < n_samples // 2:
        print("Слишком мало данных — проверьте, что стенд в режиме Plotter ('p' в Serial Monitor).")
        return

    t = np.arange(block.size) / EFFECTIVE_RATE_HZ

    # Огибающая через преобразование Гильберта — быстрое затухание/нарастание амплитуды,
    # то самое "у-у-у-у", которое слышно на биениях (см. THEORY.md).
    from scipy.signal import hilbert
    envelope = np.abs(hilbert(block))

    fig, (ax_wave, ax_spec) = plt.subplots(2, 1, figsize=(10, 6))
    fig.suptitle("Биения двух источников — записанный сигнал")

    ax_wave.plot(t, block, lw=0.5, alpha=0.6, label="сигнал")
    ax_wave.plot(t, envelope, lw=2, color="crimson", label="огибающая |Гильберт|")
    ax_wave.plot(t, -envelope, lw=2, color="crimson")
    ax_wave.set_xlabel("Время, с")
    ax_wave.set_ylabel("Амплитуда, LSB")
    ax_wave.legend(loc="upper right")
    ax_wave.grid(alpha=0.3)

    spectrum = np.abs(np.fft.rfft(block * np.hanning(block.size)))
    freqs = np.fft.rfftfreq(block.size, d=1.0 / EFFECTIVE_RATE_HZ)
    mask = freqs < 2000  # камертоны обычно ниже 2 кГц
    ax_spec.plot(freqs[mask], spectrum[mask])
    ax_spec.set_xlabel("Частота, Гц")
    ax_spec.set_ylabel("|БПФ|")
    ax_spec.set_title("Спектр записи — два пика f₁ и f₂, разница между ними и есть частота биений")
    ax_spec.grid(alpha=0.3)

    # Найти два самых высоких пика спектра — приближённая оценка f1, f2.
    peak_idx = np.argsort(spectrum[mask])[-2:]
    peak_freqs = sorted(freqs[mask][peak_idx])
    if len(peak_freqs) == 2 and peak_freqs[0] > 20:
        beat_hz = abs(peak_freqs[1] - peak_freqs[0])
        print(f"\nНайдены пики: f1={peak_freqs[0]:.1f} Гц, f2={peak_freqs[1]:.1f} Гц")
        print(f"Ожидаемая частота биений |f1-f2| = {beat_hz:.2f} Гц")
        print("Сравните с числом провалов огибающей в секунду на верхнем графике.")

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


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=["pitch", "beats"], default="pitch")
    parser.add_argument("port", nargs="?", default="/dev/ttyUSB0")
    args = parser.parse_args()

    if args.mode == "pitch":
        run_pitch_mode(args.port)
    else:
        run_beats_mode(args.port)


if __name__ == "__main__":
    main()