"""
Lab024A. Первый приём реальных IQ-сэмплов с Pluto+.

Схема:
    антенна 40–860 МГц -> RX1

Передатчики TX1 и TX2 не используются.
"""

from pathlib import Path
import json

import adi
import matplotlib.pyplot as plt
import numpy as np


# ---------------------------------------------------------------------
# Параметры приёмника
# ---------------------------------------------------------------------

PLUTO_URI = "ip:192.168.2.1"

CENTER_FREQUENCY_HZ = 100_000_000
SAMPLE_RATE_HZ = 2_400_000
RX_BANDWIDTH_HZ = 2_000_000

RX_BUFFER_SIZE = 262_144

OUTPUT_DIRECTORY = Path("data/raw")
IQ_FILE_PATH = OUTPUT_DIRECTORY / "lab024a_pluto_rx_100mhz.npy"
METADATA_FILE_PATH = OUTPUT_DIRECTORY / "lab024a_pluto_rx_100mhz.json"
SPECTRUM_FILE_PATH = OUTPUT_DIRECTORY / "lab024a_pluto_rx_100mhz_spectrum.png"


def calculate_spectrum(
    samples: np.ndarray,
    sample_rate_hz: float,
    center_frequency_hz: float,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Рассчитывает спектр принятого комплексного IQ-сигнала.

    Возвращает:
        frequencies_hz — абсолютные радиочастоты;
        power_db       — относительная мощность спектра в дБ.
    """

    sample_count = len(samples)

    window = np.hanning(sample_count)
    windowed_samples = samples * window

    spectrum = np.fft.fftshift(np.fft.fft(windowed_samples))
    power = np.abs(spectrum) ** 2

    power_db = 10.0 * np.log10(power + 1e-12)
    power_db -= np.max(power_db)

    baseband_frequencies_hz = np.fft.fftshift(
        np.fft.fftfreq(sample_count, d=1.0 / sample_rate_hz)
    )

    absolute_frequencies_hz = (
        center_frequency_hz + baseband_frequencies_hz
    )

    return absolute_frequencies_hz, power_db


def main() -> None:
    OUTPUT_DIRECTORY.mkdir(parents=True, exist_ok=True)

    print("Подключение к Pluto+...")
    sdr = adi.Pluto(uri=PLUTO_URI)

    # Используем только первый приёмный канал RX1.
    sdr.rx_enabled_channels = [0]

    sdr.sample_rate = SAMPLE_RATE_HZ
    sdr.rx_lo = CENTER_FREQUENCY_HZ
    sdr.rx_rf_bandwidth = RX_BANDWIDTH_HZ

    # Автоматическая регулировка усиления.
    sdr.gain_control_mode_chan0 = "slow_attack"

    sdr.rx_buffer_size = RX_BUFFER_SIZE

    print()
    print("Параметры приёмника:")
    print(f"  URI:                    {PLUTO_URI}")
    print(f"  Центральная частота:    {sdr.rx_lo / 1e6:.3f} МГц")
    print(f"  Частота дискретизации:  {sdr.sample_rate / 1e6:.3f} Мвыб/с")
    print(f"  Полоса RX:              {sdr.rx_rf_bandwidth / 1e6:.3f} МГц")
    print(f"  Режим усиления:         {sdr.gain_control_mode_chan0}")
    print(f"  Размер буфера:          {sdr.rx_buffer_size} отсчётов")

    print()
    print("Получение IQ-сэмплов...")

    # Первый буфер после перенастройки иногда содержит переходный процесс.
    _ = sdr.rx()

    # Второй буфер сохраняем и анализируем.
    samples = np.asarray(sdr.rx(), dtype=np.complex64)

    print("IQ-сэмплы получены.")

    np.save(IQ_FILE_PATH, samples)

    mean_value = np.mean(samples)
    rms_value = np.sqrt(np.mean(np.abs(samples) ** 2))
    peak_value = np.max(np.abs(samples))

    metadata = {
        "pluto_uri": PLUTO_URI,
        "center_frequency_hz": int(sdr.rx_lo),
        "sample_rate_hz": int(sdr.sample_rate),
        "rx_bandwidth_hz": int(sdr.rx_rf_bandwidth),
        "gain_control_mode": sdr.gain_control_mode_chan0,
        "sample_count": int(len(samples)),
        "sample_dtype": str(samples.dtype),
        "mean_i": float(np.real(mean_value)),
        "mean_q": float(np.imag(mean_value)),
        "rms": float(rms_value),
        "peak": float(peak_value),
    }

    with METADATA_FILE_PATH.open("w", encoding="utf-8") as metadata_file:
        json.dump(metadata, metadata_file, ensure_ascii=False, indent=4)

    frequencies_hz, power_db = calculate_spectrum(
        samples=samples,
        sample_rate_hz=float(sdr.sample_rate),
        center_frequency_hz=float(sdr.rx_lo),
    )

    plt.figure(figsize=(12, 6))
    plt.plot(frequencies_hz / 1e6, power_db)

    plt.title("Lab024A. Спектр сигнала, принятого Pluto+")
    plt.xlabel("Частота, МГц")
    plt.ylabel("Относительная мощность, дБ")
    plt.grid(True)
    plt.ylim(-100, 5)
    plt.tight_layout()

    plt.savefig(SPECTRUM_FILE_PATH, dpi=150)
    plt.show()

    print()
    print("Статистика:")
    print(f"  Количество сэмплов: {len(samples)}")
    print(f"  Тип данных:         {samples.dtype}")
    print(f"  Среднее I:          {np.real(mean_value):.3f}")
    print(f"  Среднее Q:          {np.imag(mean_value):.3f}")
    print(f"  RMS:                {rms_value:.3f}")
    print(f"  Пиковая амплитуда:  {peak_value:.3f}")

    print()
    print("Созданы файлы:")
    print(f"  IQ:        {IQ_FILE_PATH}")
    print(f"  Метаданные:{METADATA_FILE_PATH}")
    print(f"  Спектр:    {SPECTRUM_FILE_PATH}")

    print()
    print("Lab024A выполнена успешно.")


if __name__ == "__main__":
    main()