"""
Lab019. Программный приёмник BPSK-радиокадра.

Программа:

1. Загружает IQ-сэмплы, созданные в Lab018.
2. Намеренно добавляет неизвестную задержку.
3. Намеренно поворачивает фазу сигнала.
4. Выполняет согласованную RRC-фильтрацию.
5. Перебирает все возможные фазы дискретизации.
6. Ищет PREAMBLE + RADIO SYNC корреляционным методом.
7. Оценивает фазовый поворот BPSK.
8. Восстанавливает биты радиокадра.
9. Читает длину внутреннего пакета.
10. Разбирает пакет SDR Rover Link.
11. Проверяет CRC-32.
12. Восстанавливает исходное сообщение.

Упрощения:

- частота передатчика и приёмника совпадает;
- частота дискретизации совпадает;
- доплеровский сдвиг отсутствует;
- многолучёвость отсутствует;
- шум пока не добавляется.
"""

from pathlib import Path
import struct

import matplotlib.pyplot as plt
import numpy as np

from protocol.packet import (
    MESSAGE_TYPE_TEXT,
    parse_packet,
)


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

RADIO_SYNC_WORD = 0xD391

PREAMBLE_BIT_COUNT = 64

SAMPLES_PER_SYMBOL = 32

RRC_ROLLOFF = 0.35
RRC_SPAN_SYMBOLS = 10


# ============================================================
# Контрольные данные
# ============================================================

EXPECTED_MESSAGE = "ПРИВЕТ SDR"
EXPECTED_SEQUENCE_NUMBER = 18


# ============================================================
# Искусственные искажения для проверки приёмника
# ============================================================

# Приёмнику заранее неизвестно,
# с какого конкретно сэмпла начинается сигнал.
TEST_SAMPLE_DELAY = 11

# Приёмнику также заранее неизвестна
# фаза несущей BPSK.
TEST_PHASE_OFFSET_DEGREES = 37.0


# ============================================================
# Пути
# ============================================================

INPUT_IQ_PATH = Path(
    "data/processed/lab018/"
    "lab018_bpsk_tx_iq.npy"
)

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

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

GRAPH_PATH = (
    OUTPUT_DIRECTORY
    / "lab019_bpsk_receiver.png"
)

REPORT_PATH = (
    OUTPUT_DIRECTORY
    / "lab019_receiver_report.txt"
)

RECOVERED_PACKET_PATH = (
    OUTPUT_DIRECTORY
    / "lab019_recovered_packet.bin"
)


# ============================================================
# Преобразование bytes → bits
# ============================================================

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
    )


# ============================================================
# Преобразование bits → bytes
# ============================================================

def bits_to_bytes(
    bits: np.ndarray,
) -> bytes:
    """
    Упаковать отдельные биты обратно в bytes.
    """

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

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

    if len(bits) % 8 != 0:
        raise ValueError(
            "Количество битов должно быть кратно восьми"
        )

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

    return np.packbits(
        bits
    ).tobytes()


# ============================================================
# BPSK-модулятор для опорной последовательности
# ============================================================

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

        0 → -1
        1 → +1
    """

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

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

    return symbols.astype(
        np.complex128
    )


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

def bpsk_demodulate(
    symbols: np.ndarray,
) -> np.ndarray:
    """
    Принять решение по знаку компоненты I.

        I < 0  → 0
        I >= 0 → 1
    """

    return (
        symbols.real >= 0.0
    ).astype(np.uint8)


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

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

    Формула и параметры совпадают с Lab018.
    """

    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
    ):

        if np.isclose(
            time_value,
            0.0,
        ):
            taps[index] = (
                1.0
                - beta
                + 4.0 * beta / np.pi
            )

            continue

        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 build_frame_marker(
) -> tuple[np.ndarray, np.ndarray]:
    """
    Сформировать:

        PREAMBLE + RADIO SYNC

    Возвращает биты и BPSK-символы маркера.
    """

    preamble_bits = np.tile(
        np.array(
            [1, 0],
            dtype=np.uint8,
        ),
        PREAMBLE_BIT_COUNT // 2,
    )

    sync_bytes = struct.pack(
        ">H",
        RADIO_SYNC_WORD,
    )

    sync_bits = bytes_to_bits(
        sync_bytes
    )

    marker_bits = np.concatenate(
        [
            preamble_bits,
            sync_bits,
        ]
    )

    marker_symbols = bpsk_modulate(
        marker_bits
    )

    return marker_bits, marker_symbols


# ============================================================
# Корреляционный поиск радиокадра
# ============================================================

def find_radio_frame(
    matched_iq: np.ndarray,
    marker_symbols: np.ndarray,
    samples_per_symbol: int,
) -> dict:
    """
    Найти фазу дискретизации и начало радиокадра.

    Для каждой возможной фазы:

        0, 1, 2, ... SPS - 1

    берём по одному сэмплу на символ и вычисляем
    нормированную корреляцию с известным маркером.

    Использование комплексной корреляции позволяет
    одновременно оценить постоянный фазовый поворот.
    """

    marker_energy = float(
        np.sum(
            np.abs(marker_symbols) ** 2
        )
    )

    best_result = None

    for sample_phase in range(
        samples_per_symbol
    ):

        symbol_samples = matched_iq[
            sample_phase::samples_per_symbol
        ]

        if len(symbol_samples) < len(
            marker_symbols
        ):
            continue

        correlation = np.correlate(
            symbol_samples,
            marker_symbols,
            mode="valid",
        )

        window_energy = np.convolve(
            np.abs(symbol_samples) ** 2,
            np.ones(
                len(marker_symbols)
            ),
            mode="valid",
        )

        denominator = (
            np.sqrt(
                window_energy
                * marker_energy
            )
            + 1e-12
        )

        normalized_correlation = (
            np.abs(correlation)
            / denominator
        )

        start_symbol_index = int(
            np.argmax(
                normalized_correlation
            )
        )

        correlation_score = float(
            normalized_correlation[
                start_symbol_index
            ]
        )

        complex_correlation = correlation[
            start_symbol_index
        ]

        if (
            best_result is None
            or correlation_score
            > best_result["score"]
        ):
            best_result = {
                "score": correlation_score,
                "sample_phase": sample_phase,
                "start_symbol_index": (
                    start_symbol_index
                ),
                "symbol_samples": (
                    symbol_samples
                ),
                "correlation": correlation,
                "normalized_correlation": (
                    normalized_correlation
                ),
                "complex_correlation": (
                    complex_correlation
                ),
            }

    if best_result is None:
        raise RuntimeError(
            "Не удалось выполнить поиск радиокадра"
        )

    return best_result


# ============================================================
# Загрузка IQ из Lab018
# ============================================================

if not INPUT_IQ_PATH.exists():
    raise FileNotFoundError(
        f"Не найден IQ-файл: {INPUT_IQ_PATH}. "
        "Сначала необходимо выполнить Lab018."
    )

transmitted_iq = np.load(
    INPUT_IQ_PATH
)

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

if not np.iscomplexobj(
    transmitted_iq
):
    raise ValueError(
        "Файл не содержит комплексные IQ-сэмплы"
    )


# ============================================================
# Имитация неизвестных параметров канала
# ============================================================

phase_offset_radians = np.deg2rad(
    TEST_PHASE_OFFSET_DEGREES
)

phase_rotated_iq = (
    transmitted_iq
    * np.exp(
        1j * phase_offset_radians
    )
)

delayed_iq = np.concatenate(
    [
        np.zeros(
            TEST_SAMPLE_DELAY,
            dtype=np.complex64,
        ),
        phase_rotated_iq.astype(
            np.complex64
        ),
    ]
)

received_iq = delayed_iq


# ============================================================
# Согласованный RRC-фильтр
# ============================================================

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

matched_iq = np.convolve(
    received_iq,
    rrc_taps,
    mode="full",
)


# ============================================================
# Поиск PREAMBLE + RADIO SYNC
# ============================================================

marker_bits, marker_symbols = (
    build_frame_marker()
)

search_result = find_radio_frame(
    matched_iq=matched_iq,
    marker_symbols=marker_symbols,
    samples_per_symbol=SAMPLES_PER_SYMBOL,
)

correlation_score = search_result[
    "score"
]

selected_sample_phase = search_result[
    "sample_phase"
]

frame_start_symbol = search_result[
    "start_symbol_index"
]

symbol_samples_before_correction = (
    search_result["symbol_samples"]
)

complex_correlation = search_result[
    "complex_correlation"
]


# ============================================================
# Оценка и компенсация фазы
# ============================================================

estimated_phase_radians = np.angle(
    complex_correlation
)

estimated_phase_degrees = np.rad2deg(
    estimated_phase_radians
)

phase_correction = np.exp(
    -1j * estimated_phase_radians
)

corrected_symbol_samples = (
    symbol_samples_before_correction
    * phase_correction
)


# ============================================================
# Демодуляция
# ============================================================

all_received_bits = bpsk_demodulate(
    corrected_symbol_samples
)

available_frame_bits = all_received_bits[
    frame_start_symbol:
]


# ============================================================
# Проверка преамбулы
# ============================================================

received_marker_bits = available_frame_bits[
    :len(marker_bits)
]

marker_bit_errors = int(
    np.count_nonzero(
        received_marker_bits
        != marker_bits
    )
)


# ============================================================
# Чтение радиозаголовка
# ============================================================

RADIO_HEADER_BIT_COUNT = 32

radio_header_start = (
    PREAMBLE_BIT_COUNT
)

radio_header_end = (
    radio_header_start
    + RADIO_HEADER_BIT_COUNT
)

if len(available_frame_bits) < radio_header_end:
    raise RuntimeError(
        "Недостаточно битов для чтения радиозаголовка"
    )

received_radio_header_bits = (
    available_frame_bits[
        radio_header_start:
        radio_header_end
    ]
)

received_radio_header = bits_to_bytes(
    received_radio_header_bits
)

(
    received_radio_sync,
    protocol_packet_length,
) = struct.unpack(
    ">HH",
    received_radio_header,
)


# ============================================================
# Извлечение внутреннего пакета
# ============================================================

if received_radio_sync != RADIO_SYNC_WORD:
    raise RuntimeError(
        "Получено неверное радиосинхрослово: "
        f"0x{received_radio_sync:04X}"
    )

protocol_packet_bit_count = (
    protocol_packet_length * 8
)

protocol_packet_start = (
    radio_header_end
)

protocol_packet_end = (
    protocol_packet_start
    + protocol_packet_bit_count
)

if len(available_frame_bits) < protocol_packet_end:
    raise RuntimeError(
        "Принятый радиокадр короче значения LENGTH"
    )

received_protocol_packet_bits = (
    available_frame_bits[
        protocol_packet_start:
        protocol_packet_end
    ]
)

received_protocol_packet = bits_to_bytes(
    received_protocol_packet_bits
)

RECOVERED_PACKET_PATH.write_bytes(
    received_protocol_packet
)


# ============================================================
# Разбор внутреннего пакета
# ============================================================

parsed_packet = parse_packet(
    received_protocol_packet
)

restored_message = (
    parsed_packet.payload.decode(
        "utf-8"
    )
)


# ============================================================
# Вывод
# ============================================================

print(
    "=== Lab019. Программный BPSK-приёмник ==="
)

print("\nВходной IQ-файл:")

print(INPUT_IQ_PATH)

print("\nКоличество входных IQ-сэмплов:")

print(len(transmitted_iq))

print("\nИскусственно внесено:")

print(
    "Задержка:",
    TEST_SAMPLE_DELAY,
    "сэмплов",
)

print(
    "Поворот фазы:",
    TEST_PHASE_OFFSET_DEGREES,
    "градусов",
)


print("\nРезультат синхронизации:")

print(
    "Найденная фаза дискретизации:",
    selected_sample_phase,
    "из",
    SAMPLES_PER_SYMBOL,
)

print(
    "Начало кадра в потоке символов:",
    frame_start_symbol,
)

print(
    "Корреляционная оценка:",
    f"{correlation_score:.6f}",
)

print(
    "Оценённый фазовый поворот:",
    f"{estimated_phase_degrees:.2f}",
    "градуса",
)

print(
    "Ошибки в PREAMBLE + SYNC:",
    marker_bit_errors,
)


print("\nРадиозаголовок:")

print(
    "RADIO SYNC:",
    f"0x{received_radio_sync:04X}",
)

print(
    "LENGTH:",
    protocol_packet_length,
    "байт",
)


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

print(
    "Версия:",
    parsed_packet.version,
)

print(
    "Тип сообщения:",
    parsed_packet.message_type,
)

print(
    "Порядковый номер:",
    parsed_packet.sequence_number,
)

print(
    "Размер PAYLOAD:",
    len(parsed_packet.payload),
    "байт",
)


print("\nВосстановленное сообщение:")

print(restored_message)


print("\nВосстановленный бинарный пакет:")

print(RECOVERED_PACKET_PATH)


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

figure, axes = plt.subplots(
    2,
    2,
    figsize=(13, 10),
)


# ------------------------------------------------------------
# 1. Корреляция
# ------------------------------------------------------------

normalized_correlation = search_result[
    "normalized_correlation"
]

axes[0, 0].plot(
    normalized_correlation
)

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

axes[0, 0].set_xlabel(
    "Предполагаемое начало маркера, символ"
)

axes[0, 0].set_ylabel(
    "Нормированная корреляция"
)

axes[0, 0].set_title(
    "Поиск PREAMBLE + RADIO SYNC"
)

axes[0, 0].grid(
    True
)


# ------------------------------------------------------------
# 2. Созвездие до коррекции
# ------------------------------------------------------------

marker_start = frame_start_symbol

marker_end = (
    marker_start
    + len(marker_symbols)
)

marker_samples_before = (
    symbol_samples_before_correction[
        marker_start:
        marker_end
    ]
)

axes[0, 1].scatter(
    marker_samples_before.real,
    marker_samples_before.imag,
    s=20,
)

axes[0, 1].axhline(
    0.0,
    linewidth=1,
)

axes[0, 1].axvline(
    0.0,
    linewidth=1,
)

axes[0, 1].set_xlabel(
    "I"
)

axes[0, 1].set_ylabel(
    "Q"
)

axes[0, 1].set_title(
    "BPSK до компенсации фазы"
)

axes[0, 1].axis(
    "equal"
)

axes[0, 1].grid(
    True
)


# ------------------------------------------------------------
# 3. Созвездие после коррекции
# ------------------------------------------------------------

marker_samples_after = (
    corrected_symbol_samples[
        marker_start:
        marker_end
    ]
)

axes[1, 0].scatter(
    marker_samples_after.real,
    marker_samples_after.imag,
    s=20,
)

axes[1, 0].axhline(
    0.0,
    linewidth=1,
)

axes[1, 0].axvline(
    0.0,
    linewidth=1,
)

axes[1, 0].set_xlabel(
    "I"
)

axes[1, 0].set_ylabel(
    "Q"
)

axes[1, 0].set_title(
    "BPSK после компенсации фазы"
)

axes[1, 0].axis(
    "equal"
)

axes[1, 0].grid(
    True
)


# ------------------------------------------------------------
# 4. Восстановленные биты кадра
# ------------------------------------------------------------

received_frame_bits = available_frame_bits[
    :protocol_packet_end
]

axes[1, 1].step(
    np.arange(
        len(received_frame_bits)
    ),
    received_frame_bits,
    where="post",
)

axes[1, 1].axvline(
    PREAMBLE_BIT_COUNT,
    linestyle="--",
)

axes[1, 1].axvline(
    radio_header_end,
    linestyle="--",
)

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

axes[1, 1].set_ylabel(
    "Бит"
)

axes[1, 1].set_title(
    "Восстановленный радиокадр"
)

axes[1, 1].set_ylim(
    -0.2,
    1.2,
)

axes[1, 1].grid(
    True
)


figure.tight_layout()

figure.savefig(
    GRAPH_PATH,
    dpi=160,
)

plt.close(
    figure
)


# ============================================================
# Отчёт
# ============================================================

report_lines = [
    "Lab019. BPSK receiver",
    "",
    f"Input IQ: {INPUT_IQ_PATH}",
    f"Artificial sample delay: {TEST_SAMPLE_DELAY}",
    (
        "Artificial phase offset: "
        f"{TEST_PHASE_OFFSET_DEGREES:.2f} deg"
    ),
    (
        "Detected sample phase: "
        f"{selected_sample_phase}"
    ),
    (
        "Detected frame start symbol: "
        f"{frame_start_symbol}"
    ),
    (
        "Correlation score: "
        f"{correlation_score:.6f}"
    ),
    (
        "Estimated phase: "
        f"{estimated_phase_degrees:.2f} deg"
    ),
    (
        "Marker bit errors: "
        f"{marker_bit_errors}"
    ),
    (
        "Radio sync: "
        f"0x{received_radio_sync:04X}"
    ),
    (
        "Protocol packet length: "
        f"{protocol_packet_length} bytes"
    ),
    (
        "Sequence number: "
        f"{parsed_packet.sequence_number}"
    ),
    f"Restored message: {restored_message}",
]

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


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

assert received_radio_sync == RADIO_SYNC_WORD

assert protocol_packet_length == len(
    received_protocol_packet
)

assert marker_bit_errors == 0

assert parsed_packet.message_type == (
    MESSAGE_TYPE_TEXT
)

assert parsed_packet.sequence_number == (
    EXPECTED_SEQUENCE_NUMBER
)

assert restored_message == EXPECTED_MESSAGE

assert abs(
    estimated_phase_degrees
    - TEST_PHASE_OFFSET_DEGREES
) < 1.0

assert correlation_score > 0.95

assert GRAPH_PATH.exists()
assert REPORT_PATH.exists()
assert RECOVERED_PACKET_PATH.exists()


print("\nГрафик приёмника:")

print(GRAPH_PATH)

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

print(REPORT_PATH)

print(
    "\nПроверка пройдена: "
    "радиокадр синхронизирован, "
    "BPSK демодулирован, "
    "внутренний пакет восстановлен и прошёл CRC."
)