"""
Lab018. Формирование BPSK-радиокадра.

Программа:

1. Формирует внутренний пакет SDR Rover Link.
2. Добавляет радиопреамбулу.
3. Добавляет внешнее синхрослово.
4. Добавляет длину внутреннего пакета.
5. Преобразует радиокадр в биты.
6. Модулирует биты методом BPSK.
7. Увеличивает частоту дискретизации.
8. Применяет Root Raised Cosine-фильтр.
9. Формирует комплексные IQ-сэмплы.
10. Сохраняет IQ в форматах NPY и CF32.
11. Строит графики кадра, сигнала и спектра.

Передача через SDR в этой лабораторной не выполняется.
"""

from pathlib import Path
import struct

import matplotlib.pyplot as plt
import numpy as np

from protocol.packet import (
    MESSAGE_TYPE_TEXT,
    build_packet,
)


# ============================================================
# Настройки содержимого
# ============================================================

MESSAGE = "ПРИВЕТ SDR"

SEQUENCE_NUMBER = 18


# ============================================================
# Настройки радиокадра
# ============================================================

# Радиосинхрослово физического уровня.
#
# Не следует путать с SYNC_WORD = 0xAA55,
# который находится внутри packet.py.
RADIO_SYNC_WORD = 0xD391

# Преамбула 10101010...
PREAMBLE_BIT_COUNT = 64

# Пустой участок до и после радиокадра.
GUARD_SYMBOL_COUNT = 16


# ============================================================
# Настройки BPSK
# ============================================================

# BPSK передаёт один бит одним символом.
SYMBOL_RATE = 20_000

# На один BPSK-символ приходится 32 IQ-сэмпла.
SAMPLES_PER_SYMBOL = 32

SAMPLE_RATE = (
    SYMBOL_RATE
    * SAMPLES_PER_SYMBOL
)

# Коэффициент скругления RRC-фильтра.
RRC_ROLLOFF = 0.35

# Длина фильтра в символах.
#
# Используем чётное значение.
RRC_SPAN_SYMBOLS = 10

# Амплитудный запас.
#
# Не используем полную шкалу ±1,
# чтобы позднее не перегружать передатчик.
TX_AMPLITUDE = 0.70


# ============================================================
# Выходные файлы
# ============================================================

OUTPUT_DIRECTORY = Path(
    "data/processed/lab018"
)

OUTPUT_DIRECTORY.mkdir(
    parents=True,
    exist_ok=True,
)

NPY_PATH = (
    OUTPUT_DIRECTORY
    / "lab018_bpsk_tx_iq.npy"
)

CF32_PATH = (
    OUTPUT_DIRECTORY
    / "lab018_bpsk_tx_iq.cf32"
)

GRAPH_PATH = (
    OUTPUT_DIRECTORY
    / "lab018_bpsk_radio_frame.png"
)

INFO_PATH = (
    OUTPUT_DIRECTORY
    / "lab018_frame_info.txt"
)


# ============================================================
# Преобразование байтов в биты
# ============================================================

def bytes_to_bits(
    data: bytes,
) -> np.ndarray:
    """
    Преобразовать bytes в массив битов 0 и 1.

    Используется порядок от старшего бита к младшему.
    """

    if not isinstance(data, bytes):
        raise TypeError(
            "data должен иметь тип bytes"
        )

    byte_array = np.frombuffer(
        data,
        dtype=np.uint8,
    )

    return np.unpackbits(
        byte_array
    )


# ============================================================
# BPSK-модулятор
# ============================================================

def bpsk_modulate(
    bits: np.ndarray,
) -> np.ndarray:
    """
    Преобразовать биты в BPSK-символы.

        0 → -1
        1 → +1
    """

    bits = np.asarray(
        bits,
        dtype=np.uint8,
    )

    if bits.ndim != 1:
        raise ValueError(
            "bits должен быть одномерным массивом"
        )

    if not np.all(
        (bits == 0) | (bits == 1)
    ):
        raise ValueError(
            "bits должен содержать только 0 и 1"
        )

    symbols = (
        2.0 * bits.astype(np.float64)
        - 1.0
    )

    return symbols.astype(
        np.complex128
    )


# ============================================================
# Root Raised Cosine-фильтр
# ============================================================

def root_raised_cosine_taps(
    rolloff: float,
    samples_per_symbol: int,
    span_symbols: int,
) -> np.ndarray:
    """
    Рассчитать коэффициенты Root Raised Cosine-фильтра.

    Параметры
    ----------
    rolloff:
        Коэффициент скругления beta.

    samples_per_symbol:
        Количество сэмплов на символ.

    span_symbols:
        Полная длина фильтра в символах.

    Возвращает
    ----------
    Одномерный массив коэффициентов фильтра.
    """

    if not 0.0 < rolloff <= 1.0:
        raise ValueError(
            "rolloff должен находиться в диапазоне 0...1"
        )

    if samples_per_symbol <= 0:
        raise ValueError(
            "samples_per_symbol должен быть положительным"
        )

    if span_symbols <= 0:
        raise ValueError(
            "span_symbols должен быть положительным"
        )

    if span_symbols % 2 != 0:
        raise ValueError(
            "span_symbols должен быть чётным"
        )

    half_sample_count = (
        span_symbols
        * samples_per_symbol
        // 2
    )

    sample_indexes = np.arange(
        -half_sample_count,
        half_sample_count + 1,
        dtype=np.float64,
    )

    # Время нормировано к длительности одного символа.
    time_values = (
        sample_indexes
        / samples_per_symbol
    )

    taps = np.zeros_like(
        time_values
    )

    beta = rolloff

    for index, time_value in enumerate(
        time_values
    ):

        # Особая точка t = 0.
        if np.isclose(
            time_value,
            0.0,
        ):
            taps[index] = (
                1.0
                - beta
                + (
                    4.0
                    * beta
                    / np.pi
                )
            )

            continue

        # Особые точки t = ±1/(4 beta).
        if np.isclose(
            abs(time_value),
            1.0 / (4.0 * beta),
        ):
            taps[index] = (
                beta
                / np.sqrt(2.0)
                * (
                    (
                        1.0
                        + 2.0 / np.pi
                    )
                    * np.sin(
                        np.pi
                        / (4.0 * beta)
                    )
                    + (
                        1.0
                        - 2.0 / np.pi
                    )
                    * np.cos(
                        np.pi
                        / (4.0 * beta)
                    )
                )
            )

            continue

        numerator = (
            np.sin(
                np.pi
                * time_value
                * (1.0 - beta)
            )
            + (
                4.0
                * beta
                * time_value
                * np.cos(
                    np.pi
                    * time_value
                    * (1.0 + beta)
                )
            )
        )

        denominator = (
            np.pi
            * time_value
            * (
                1.0
                - (
                    4.0
                    * beta
                    * time_value
                ) ** 2
            )
        )

        taps[index] = (
            numerator
            / denominator
        )

    # Нормируем энергию фильтра.
    taps /= np.sqrt(
        np.sum(
            taps ** 2
        )
    )

    return taps


# ============================================================
# Увеличение частоты дискретизации
# ============================================================

def upsample_symbols(
    symbols: np.ndarray,
    samples_per_symbol: int,
) -> np.ndarray:
    """
    Вставить между символами нулевые сэмплы.

    Пример при SPS = 4:

        [1, -1]
            ↓
        [1, 0, 0, 0, -1, 0, 0, 0]
    """

    symbols = np.asarray(
        symbols,
        dtype=np.complex128,
    )

    upsampled = np.zeros(
        len(symbols)
        * samples_per_symbol,
        dtype=np.complex128,
    )

    upsampled[
        ::samples_per_symbol
    ] = symbols

    return upsampled


# ============================================================
# Спектр
# ============================================================

def calculate_normalized_spectrum(
    iq_samples: np.ndarray,
    sample_rate: float,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Рассчитать нормированный амплитудный спектр.
    """

    fft_size = 2 ** int(
        np.ceil(
            np.log2(
                max(
                    len(iq_samples),
                    1024,
                )
            )
        )
    )

    spectrum = np.fft.fftshift(
        np.fft.fft(
            iq_samples,
            n=fft_size,
        )
    )

    frequencies = np.fft.fftshift(
        np.fft.fftfreq(
            fft_size,
            d=1.0 / sample_rate,
        )
    )

    magnitude = np.abs(
        spectrum
    )

    maximum_magnitude = np.max(
        magnitude
    )

    magnitude_db = 20.0 * np.log10(
        magnitude / maximum_magnitude
        + 1e-12
    )

    return frequencies, magnitude_db


# ============================================================
# Формирование внутреннего пакета
# ============================================================

payload = MESSAGE.encode(
    "utf-8"
)

protocol_packet = build_packet(
    payload=payload,
    message_type=MESSAGE_TYPE_TEXT,
    sequence_number=SEQUENCE_NUMBER,
)


# ============================================================
# Формирование внешнего заголовка радиокадра
# ============================================================

if len(protocol_packet) > 65535:
    raise ValueError(
        "Внутренний пакет слишком велик "
        "для 16-битного поля длины"
    )

radio_header = struct.pack(
    ">HH",
    RADIO_SYNC_WORD,
    len(protocol_packet),
)

radio_header_bits = bytes_to_bits(
    radio_header
)

protocol_packet_bits = bytes_to_bits(
    protocol_packet
)

# Преамбула 101010...
preamble_bits = np.tile(
    np.array(
        [1, 0],
        dtype=np.uint8,
    ),
    PREAMBLE_BIT_COUNT // 2,
)

radio_frame_bits = np.concatenate(
    [
        preamble_bits,
        radio_header_bits,
        protocol_packet_bits,
    ]
)


# ============================================================
# BPSK-модуляция
# ============================================================

radio_frame_symbols = bpsk_modulate(
    radio_frame_bits
)

upsampled_symbols = upsample_symbols(
    radio_frame_symbols,
    SAMPLES_PER_SYMBOL,
)


# ============================================================
# Guard interval
# ============================================================

guard_samples = np.zeros(
    GUARD_SYMBOL_COUNT
    * SAMPLES_PER_SYMBOL,
    dtype=np.complex128,
)

guarded_upsampled_symbols = np.concatenate(
    [
        guard_samples,
        upsampled_symbols,
        guard_samples,
    ]
)


# ============================================================
# RRC-фильтрация
# ============================================================

rrc_taps = root_raised_cosine_taps(
    rolloff=RRC_ROLLOFF,
    samples_per_symbol=SAMPLES_PER_SYMBOL,
    span_symbols=RRC_SPAN_SYMBOLS,
)

filtered_iq = np.convolve(
    guarded_upsampled_symbols,
    rrc_taps,
    mode="full",
)


# ============================================================
# Нормирование амплитуды
# ============================================================

maximum_amplitude = np.max(
    np.abs(
        filtered_iq
    )
)

if maximum_amplitude == 0.0:
    raise RuntimeError(
        "Сформирован нулевой IQ-сигнал"
    )

normalized_iq = (
    filtered_iq
    / maximum_amplitude
    * TX_AMPLITUDE
)

tx_iq = normalized_iq.astype(
    np.complex64
)


# ============================================================
# Параметры сигнала
# ============================================================

frame_duration_seconds = (
    len(tx_iq)
    / SAMPLE_RATE
)

theoretical_baseband_bandwidth_hz = (
    SYMBOL_RATE
    * (1.0 + RRC_ROLLOFF)
)

average_power = float(
    np.mean(
        np.abs(tx_iq) ** 2
    )
)

peak_power = float(
    np.max(
        np.abs(tx_iq) ** 2
    )
)

papr_db = 10.0 * np.log10(
    peak_power / average_power
)


# ============================================================
# Сохранение IQ
# ============================================================

# Формат NPY сохраняет тип и форму массива.
np.save(
    NPY_PATH,
    tx_iq,
)

# complex64 записывается как чередующиеся:
#
# float32 I, float32 Q, float32 I, float32 Q...
tx_iq.tofile(
    CF32_PATH
)


# ============================================================
# Проверка сохранённых данных
# ============================================================

loaded_npy_iq = np.load(
    NPY_PATH
)

loaded_cf32_iq = np.fromfile(
    CF32_PATH,
    dtype=np.complex64,
)

assert np.array_equal(
    loaded_npy_iq,
    tx_iq,
)

assert np.array_equal(
    loaded_cf32_iq,
    tx_iq,
)


# ============================================================
# Спектр
# ============================================================

frequencies_hz, spectrum_db = (
    calculate_normalized_spectrum(
        iq_samples=tx_iq,
        sample_rate=SAMPLE_RATE,
    )
)


# ============================================================
# Графики
# ============================================================

figure, axes = plt.subplots(
    4,
    1,
    figsize=(13, 13),
)


# ------------------------------------------------------------
# 1. Структура битового радиокадра
# ------------------------------------------------------------

bit_indexes = np.arange(
    len(radio_frame_bits)
)

axes[0].step(
    bit_indexes,
    radio_frame_bits,
    where="post",
)

preamble_end = (
    PREAMBLE_BIT_COUNT
)

sync_end = (
    preamble_end + 16
)

length_end = (
    sync_end + 16
)

axes[0].axvline(
    preamble_end,
    linestyle="--",
)

axes[0].axvline(
    sync_end,
    linestyle="--",
)

axes[0].axvline(
    length_end,
    linestyle="--",
)

axes[0].text(
    PREAMBLE_BIT_COUNT / 2,
    1.08,
    "PREAMBLE",
    ha="center",
)

axes[0].text(
    preamble_end + 8,
    1.08,
    "SYNC",
    ha="center",
)

axes[0].text(
    sync_end + 8,
    1.08,
    "LENGTH",
    ha="center",
)

axes[0].text(
    length_end
    + len(protocol_packet_bits) / 2,
    1.08,
    "SDR ROVER PACKET",
    ha="center",
)

axes[0].set_ylim(
    -0.2,
    1.3,
)

axes[0].set_xlabel(
    "Номер бита"
)

axes[0].set_ylabel(
    "Бит"
)

axes[0].set_title(
    "Структура радиокадра"
)

axes[0].grid(
    True
)


# ------------------------------------------------------------
# 2. Первые BPSK-символы
# ------------------------------------------------------------

symbol_plot_count = min(
    100,
    len(radio_frame_symbols),
)

axes[1].step(
    np.arange(
        symbol_plot_count
    ),
    radio_frame_symbols[
        :symbol_plot_count
    ].real,
    where="post",
)

axes[1].set_xlabel(
    "Номер символа"
)

axes[1].set_ylabel(
    "Амплитуда BPSK"
)

axes[1].set_title(
    "Первые BPSK-символы"
)

axes[1].set_ylim(
    -1.3,
    1.3,
)

axes[1].grid(
    True
)


# ------------------------------------------------------------
# 3. Отфильтрованный IQ-сигнал
# ------------------------------------------------------------

filter_delay_samples = (
    len(rrc_taps) - 1
) // 2

frame_start_sample = (
    GUARD_SYMBOL_COUNT
    * SAMPLES_PER_SYMBOL
    + filter_delay_samples
)

waveform_plot_symbol_count = 20

waveform_plot_start = (
    frame_start_sample
)

waveform_plot_end = (
    waveform_plot_start
    + waveform_plot_symbol_count
    * SAMPLES_PER_SYMBOL
)

waveform_plot_samples = tx_iq[
    waveform_plot_start:
    waveform_plot_end
]

time_axis_ms = (
    np.arange(
        len(waveform_plot_samples)
    )
    / SAMPLE_RATE
    * 1000.0
)

axes[2].plot(
    time_axis_ms,
    waveform_plot_samples.real,
    label="I",
)

axes[2].plot(
    time_axis_ms,
    waveform_plot_samples.imag,
    label="Q",
)

axes[2].set_xlabel(
    "Время, мс"
)

axes[2].set_ylabel(
    "Амплитуда"
)

axes[2].set_title(
    "BPSK после Root Raised Cosine-фильтра"
)

axes[2].grid(
    True
)

axes[2].legend()


# ------------------------------------------------------------
# 4. Спектр
# ------------------------------------------------------------

axes[3].plot(
    frequencies_hz / 1000.0,
    spectrum_db,
)

axes[3].axvline(
    theoretical_baseband_bandwidth_hz
    / 2.0
    / 1000.0,
    linestyle="--",
)

axes[3].axvline(
    -theoretical_baseband_bandwidth_hz
    / 2.0
    / 1000.0,
    linestyle="--",
)

axes[3].set_xlim(
    -100,
    100,
)

axes[3].set_ylim(
    -100,
    5,
)

axes[3].set_xlabel(
    "Частота относительно несущей, кГц"
)

axes[3].set_ylabel(
    "Нормированный уровень, дБ"
)

axes[3].set_title(
    "Спектр сформированного IQ-сигнала"
)

axes[3].grid(
    True
)


figure.tight_layout()

figure.savefig(
    GRAPH_PATH,
    dpi=160,
)

plt.close(
    figure
)


# ============================================================
# Текстовый отчёт
# ============================================================

report_lines = [
    "Lab018. BPSK radio frame",
    "",
    f"Message: {MESSAGE}",
    f"Protocol packet bytes: {len(protocol_packet)}",
    f"Preamble bits: {len(preamble_bits)}",
    f"Radio header bits: {len(radio_header_bits)}",
    f"Protocol packet bits: {len(protocol_packet_bits)}",
    f"Total radio frame bits: {len(radio_frame_bits)}",
    f"Symbol rate: {SYMBOL_RATE} symbols/s",
    f"Samples per symbol: {SAMPLES_PER_SYMBOL}",
    f"Sample rate: {SAMPLE_RATE} samples/s",
    f"RRC rolloff: {RRC_ROLLOFF}",
    f"RRC taps: {len(rrc_taps)}",
    f"IQ sample count: {len(tx_iq)}",
    f"Frame duration: {frame_duration_seconds:.6f} s",
    (
        "Approximate occupied bandwidth: "
        f"{theoretical_baseband_bandwidth_hz:.1f} Hz"
    ),
    f"Average normalized power: {average_power:.6f}",
    f"Peak normalized power: {peak_power:.6f}",
    f"PAPR: {papr_db:.2f} dB",
    f"NPY file: {NPY_PATH}",
    f"CF32 file: {CF32_PATH}",
]

INFO_PATH.write_text(
    "\n".join(
        report_lines
    ),
    encoding="utf-8",
)


# ============================================================
# Вывод результатов
# ============================================================

print(
    "=== Lab018. Формирование BPSK-радиокадра ==="
)

print("\nИсходное сообщение:")

print(MESSAGE)

print("\nВнутренний пакет SDR Rover Link:")

print(
    len(protocol_packet),
    "байт",
)

print("\nРадиосинхрослово:")

print(
    f"0x{RADIO_SYNC_WORD:04X}"
)

print("\nСтруктура радиокадра:")

print(
    "Преамбула:",
    len(preamble_bits),
    "бит",
)

print(
    "Радиозаголовок:",
    len(radio_header_bits),
    "бит",
)

print(
    "Внутренний пакет:",
    len(protocol_packet_bits),
    "бит",
)

print(
    "Всего:",
    len(radio_frame_bits),
    "битов и BPSK-символов",
)

print("\nПараметры дискретизации:")

print(
    "Символьная скорость:",
    SYMBOL_RATE,
    "символов/с",
)

print(
    "Сэмплов на символ:",
    SAMPLES_PER_SYMBOL,
)

print(
    "Частота дискретизации:",
    SAMPLE_RATE,
    "сэмплов/с",
)

print("\nRRC-фильтр:")

print(
    "Rolloff:",
    RRC_ROLLOFF,
)

print(
    "Количество коэффициентов:",
    len(rrc_taps),
)

print("\nIQ-сигнал:")

print(
    "Количество сэмплов:",
    len(tx_iq),
)

print(
    "Длительность:",
    f"{frame_duration_seconds * 1000:.2f}",
    "мс",
)

print(
    "Максимальная амплитуда:",
    f"{np.max(np.abs(tx_iq)):.3f}",
)

print(
    "PAPR:",
    f"{papr_db:.2f}",
    "дБ",
)

print("\nОценочная ширина спектра:")

print(
    f"{theoretical_baseband_bandwidth_hz / 1000:.1f}",
    "кГц",
)

print("\nФайл NumPy:")

print(NPY_PATH)

print("\nФайл raw complex float32:")

print(CF32_PATH)

print("\nГрафик:")

print(GRAPH_PATH)

print("\nОтчёт:")

print(INFO_PATH)


# ============================================================
# Автоматические проверки
# ============================================================

assert len(preamble_bits) == PREAMBLE_BIT_COUNT

assert radio_header[:2] == struct.pack(
    ">H",
    RADIO_SYNC_WORD,
)

assert len(radio_header) == 4

assert len(radio_frame_symbols) == len(
    radio_frame_bits
)

assert np.iscomplexobj(
    tx_iq
)

assert tx_iq.dtype == np.complex64

assert np.max(
    np.abs(tx_iq)
) <= TX_AMPLITUDE + 1e-6

assert NPY_PATH.exists()
assert CF32_PATH.exists()
assert GRAPH_PATH.exists()
assert INFO_PATH.exists()

print(
    "\nПроверка пройдена: "
    "реальный BPSK-радиокадр и IQ-сэмплы сформированы."
)