"""Проверенные примитивы BPSK-радиотракта, вынесенные из лабораторных.

Модуль не проводит экспериментов и ничего не пишет на диск: его можно
безопасно импортировать. Это и было причиной его появления — Lab018,
Lab019 и Lab023 при импорте создают каталоги и сохраняют файлы, поэтому
переиспользовать их функции напрямую невозможно.

Тела функций перенесены дословно из соответствующих лабораторных, чтобы
результаты остались сопоставимыми. Источник указан над каждой функцией.
Совпадение с оригиналами проверяется в tests/test_bpsk_radio.py.

Универсальные функции формирования и разбора радиокадра добавлены здесь
заново: в лабораторных они были жёстко привязаны к демонстрационному
текстовому сообщению и не принимали произвольные данные.
"""

from __future__ import annotations

import struct

import numpy as np

# Параметры радиокадра. Значения совпадают с Lab018, Lab019 и Lab023.
RADIO_SYNC_WORD = 0xD391
PREAMBLE_BIT_COUNT = 64
RADIO_HEADER_BIT_COUNT = 32
MARKER_BIT_COUNT = PREAMBLE_BIT_COUNT + 16

# Параметры оценки частотного рассогласования. Значения из Lab023.
SYMBOL_RATE = 20_000
CFO_REFINEMENT_HALF_WIDTH_HZ = 200.0
CFO_REFINEMENT_STEP_HZ = 1.0

MAXIMUM_PROTOCOL_PACKET_BYTES = 4096

# Перенесено дословно из Lab023 (tests/lab023_guarded_cfo_correction.py).
def bytes_to_bits(
    data: bytes,
) -> np.ndarray:
    """
    Преобразовать bytes в одномерный массив битов.
    """

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

    return np.unpackbits(
        np.frombuffer(
            data,
            dtype=np.uint8,
        )
    )


# Перенесено дословно из Lab023 (tests/lab023_guarded_cfo_correction.py).
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()


# Перенесено дословно из Lab023 (tests/lab023_guarded_cfo_correction.py).
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
    )


# Перенесено дословно из Lab023 (tests/lab023_guarded_cfo_correction.py).
def bpsk_demodulate(
    symbols: np.ndarray,
) -> np.ndarray:
    """
    Демодулировать BPSK по знаку компоненты I.
    """

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


# Перенесено дословно из Lab023 (tests/lab023_guarded_cfo_correction.py).
def estimate_carrier_parameters(
    received_symbols: np.ndarray,
    marker_symbols: np.ndarray,
) -> dict:
    """
    Оценить фазу и частотное рассогласование по маркеру.

    Алгоритм:

    1. Удалить известные BPSK-знаки маркера.
    2. Получить грубую оценку CFO по соседним символам.
    3. Выполнить уточняющий частотный поиск.
    4. Оценить начальную фазу после компенсации CFO.
    5. Рассчитать достоверность оценки.
    """

    marker_length = len(
        marker_symbols
    )

    received_marker = received_symbols[
        :marker_length
    ]

    if len(received_marker) != marker_length:
        raise ValueError(
            "Недостаточно символов маркера"
        )

    # Известные BPSK-знаки равны -1 или +1.
    # Умножение удаляет переданную манипуляцию,
    # оставляя фазу канала и шум.
    despread_marker = (
        received_marker
        * marker_symbols
    )

    marker_indexes = np.arange(
        marker_length,
        dtype=np.float64,
    )

    marker_magnitude_sum = float(
        np.sum(
            np.abs(
                despread_marker
            )
        )
    )

    # ========================================================
    # Постоянная фаза без CFO-компенсации
    # ========================================================

    constant_coherent_sum = np.sum(
        despread_marker
    )

    constant_phase = float(
        np.angle(
            constant_coherent_sum
        )
    )

    constant_coherence = float(
        np.abs(
            constant_coherent_sum
        )
        / (
            marker_magnitude_sum
            + 1e-12
        )
    )

    # ========================================================
    # Грубая оценка CFO
    # ========================================================

    adjacent_products = (
        despread_marker[1:]
        * np.conj(
            despread_marker[:-1]
        )
    )

    adjacent_sum = np.sum(
        adjacent_products
    )

    coarse_phase_increment = float(
        np.angle(
            adjacent_sum
        )
    )

    coarse_frequency_hz = (
        coarse_phase_increment
        * SYMBOL_RATE
        / (2.0 * np.pi)
    )

    phase_consistency = float(
        np.abs(
            adjacent_sum
        )
        / (
            np.sum(
                np.abs(
                    adjacent_products
                )
            )
            + 1e-12
        )
    )

    # ========================================================
    # Уточняющий поиск CFO
    # ========================================================

    frequency_candidates_hz = np.arange(
        (
            coarse_frequency_hz
            - CFO_REFINEMENT_HALF_WIDTH_HZ
        ),
        (
            coarse_frequency_hz
            + CFO_REFINEMENT_HALF_WIDTH_HZ
            + CFO_REFINEMENT_STEP_HZ / 2.0
        ),
        CFO_REFINEMENT_STEP_HZ,
        dtype=np.float64,
    )

    phase_increment_candidates = (
        2.0
        * np.pi
        * frequency_candidates_hz
        / SYMBOL_RATE
    )

    candidate_compensation = np.exp(
        -1j
        * phase_increment_candidates[
            :, np.newaxis
        ]
        * marker_indexes[
            np.newaxis, :
        ]
    )

    coherent_sums = np.sum(
        despread_marker[
            np.newaxis, :
        ]
        * candidate_compensation,
        axis=1,
    )

    best_candidate_index = int(
        np.argmax(
            np.abs(
                coherent_sums
            )
        )
    )

    estimated_frequency_hz = float(
        frequency_candidates_hz[
            best_candidate_index
        ]
    )

    phase_increment = float(
        phase_increment_candidates[
            best_candidate_index
        ]
    )

    best_coherent_sum = (
        coherent_sums[
            best_candidate_index
        ]
    )

    initial_phase_after_cfo = float(
        np.angle(
            best_coherent_sum
        )
    )

    cfo_coherence = float(
        np.abs(
            best_coherent_sum
        )
        / (
            marker_magnitude_sum
            + 1e-12
        )
    )

    coherence_gain = (
        cfo_coherence
        - constant_coherence
    )

    return {
        "constant_phase": constant_phase,
        "constant_coherence": constant_coherence,
        "phase_increment": phase_increment,
        "estimated_frequency_hz": (
            estimated_frequency_hz
        ),
        "initial_phase_after_cfo": (
            initial_phase_after_cfo
        ),
        "phase_consistency": phase_consistency,
        "cfo_coherence": cfo_coherence,
        "coherence_gain": coherence_gain,
    }


# Перенесено дословно из Lab023 (tests/lab023_guarded_cfo_correction.py).
def correct_phase_and_frequency(
    received_symbols: np.ndarray,
    initial_phase_radians: float,
    phase_increment: float,
) -> np.ndarray:
    """
    Компенсировать постоянную фазу и CFO.
    """

    symbol_indexes = np.arange(
        len(received_symbols),
        dtype=np.float64,
    )

    phase_model = (
        initial_phase_radians
        + phase_increment
        * symbol_indexes
    )

    return (
        received_symbols
        * np.exp(
            -1j * phase_model
        )
    )


# Перенесено дословно из Lab019 (tests/lab019_bpsk_receiver.py).
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


# Перенесено дословно из Lab019 (tests/lab019_bpsk_receiver.py).
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


# Перенесено дословно из Lab018 (tests/lab018_bpsk_radio_frame.py).
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


# Ниже — функции, добавленные при выделении модуля. В Lab018, Lab019 и
# Lab023 формирование кадра не принимало аргументов и собирало жёстко
# заданное текстовое сообщение, а разбор возвращал строку состояния.
# Структура кадра сохранена без изменений:
#
#     [преамбула 64 бита] [синхрослово 16 | длина 16] [пакет протокола]


def build_radio_frame(
    protocol_packet: bytes,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Сформировать радиокадр вокруг готового пакета протокола.

    Возвращает биты кадра, символы кадра и известные символы маркера,
    состоящего из преамбулы и синхрослова.
    """

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

    protocol_packet = bytes(protocol_packet)

    if not 1 <= len(protocol_packet) <= MAXIMUM_PROTOCOL_PACKET_BYTES:
        raise ValueError(
            "Длина пакета протокола должна быть от 1 до "
            f"{MAXIMUM_PROTOCOL_PACKET_BYTES} байт"
        )

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

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

    frame_bits = np.concatenate(
        [
            preamble_bits,
            bytes_to_bits(radio_header),
            bytes_to_bits(protocol_packet),
        ]
    )

    frame_symbols = bpsk_modulate(frame_bits)
    marker_symbols = frame_symbols[:MARKER_BIT_COUNT]

    return frame_bits, frame_symbols, marker_symbols


def parse_radio_frame(received_bits: np.ndarray) -> bytes | None:
    """Извлечь пакет протокола из битов принятого кадра.

    Возвращает пакет либо None, если заголовок кадра не разобран.
    Проверка контрольной суммы самого пакета в задачу не входит и
    выполняется отдельно через protocol.packet.parse_packet.
    """

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

    header_start = PREAMBLE_BIT_COUNT
    header_end = header_start + RADIO_HEADER_BIT_COUNT

    if len(received_bits) < header_end:
        return None

    try:
        radio_header = bits_to_bytes(received_bits[header_start:header_end])
        received_sync, packet_length = struct.unpack(">HH", radio_header)
    except (ValueError, struct.error):
        return None

    if received_sync != RADIO_SYNC_WORD:
        return None

    if not 1 <= packet_length <= MAXIMUM_PROTOCOL_PACKET_BYTES:
        return None

    packet_start = header_end
    packet_end = packet_start + packet_length * 8

    if len(received_bits) < packet_end:
        return None

    try:
        return bits_to_bytes(received_bits[packet_start:packet_end])
    except ValueError:
        return None


def radio_frame_bit_count(protocol_packet_bytes: int) -> int:
    """Вернуть длину радиокадра в битах для пакета заданного размера."""

    return (
        PREAMBLE_BIT_COUNT
        + RADIO_HEADER_BIT_COUNT
        + protocol_packet_bytes * 8
    )