"""
Lab013. Первая BPSK-передача цифрового пакета в Python.

Программа:
1. Формирует пакет SDR Rover Link с CRC-32.
2. Преобразует байты пакета в отдельные биты.
3. Преобразует биты в BPSK-символы.
4. Представляет символы как комплексные IQ-сэмплы.
5. Добавляет комплексный белый гауссов шум.
6. Демодулирует BPSK.
7. Восстанавливает байты пакета.
8. Проверяет пакет и CRC.
9. Сравнивает результат при разных уровнях SNR.
"""

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

from protocol.packet import (
    CRCError,
    MESSAGE_TYPE_TEXT,
    PacketError,
    build_packet,
    parse_packet,
)


# ============================================================
# Настройки эксперимента
# ============================================================

MESSAGE = "ПРИВЕТ SDR"

SEQUENCE_NUMBER = 13

# Проверим несколько уровней отношения сигнал/шум.
SNR_VALUES_DB = [
    12.0,
    6.0,
    2.0,
    0.0,
]

# Фиксированное значение обеспечивает повторяемость.
RANDOM_SEED = 2026

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

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


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

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

    bits = np.unpackbits(
        byte_array
    )

    return bits


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

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(
            "Массив должен содержать только 0 и 1"
        )

    packed_bytes = np.packbits(
        bits
    )

    return packed_bytes.tobytes()


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

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

    Используем отображение:

        0 → -1 + 0j
        1 → +1 + 0j
    """

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

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

    iq_symbols = real_symbols.astype(
        np.complex128
    )

    return iq_symbols


# ============================================================
# Модель радиошума
# ============================================================

def add_awgn(
    iq_samples: np.ndarray,
    snr_db: float,
    random_generator: np.random.Generator,
) -> np.ndarray:
    """
    Добавить комплексный белый гауссов шум AWGN.

    AWGN:
        Additive White Gaussian Noise —
        аддитивный белый гауссов шум.

    snr_db:
        Отношение средней мощности сигнала
        к средней мощности шума в децибелах.
    """

    signal_power = np.mean(
        np.abs(iq_samples) ** 2
    )

    snr_linear = 10.0 ** (
        snr_db / 10.0
    )

    noise_power = (
        signal_power / snr_linear
    )

    # Комплексный шум имеет две составляющие:
    # действительную I и мнимую Q.
    #
    # Поэтому мощность делится между ними пополам.
    noise_sigma = np.sqrt(
        noise_power / 2.0
    )

    noise = noise_sigma * (
        random_generator.standard_normal(
            len(iq_samples)
        )
        + 1j
        * random_generator.standard_normal(
            len(iq_samples)
        )
    )

    return iq_samples + noise


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

def bpsk_demodulate(
    received_iq: np.ndarray,
) -> np.ndarray:
    """
    Преобразовать принятые IQ-сэмплы в биты.

    Правило решения:

        I < 0  → бит 0
        I >= 0 → бит 1
    """

    received_bits = (
        received_iq.real >= 0.0
    ).astype(np.uint8)

    return received_bits


# ============================================================
# Формирование исходного пакета
# ============================================================

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

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

transmitted_bits = bytes_to_bits(
    original_packet
)

transmitted_iq = bpsk_modulate(
    transmitted_bits
)


print(
    "=== Lab013. BPSK в Python ==="
)

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

print("\nРазмер пакета:")
print(
    len(original_packet),
    "байт",
)

print("\nКоличество передаваемых битов:")
print(
    len(transmitted_bits)
)

print("\nПервые 32 бита:")

print(
    " ".join(
        str(bit)
        for bit in transmitted_bits[:32]
    )
)

print("\nПервые 16 BPSK-символов:")

print(
    transmitted_iq[:16]
)


# ============================================================
# Передача при разных уровнях SNR
# ============================================================

experiment_results = []

constellation_samples = {}

for experiment_index, snr_db in enumerate(
    SNR_VALUES_DB
):

    random_generator = np.random.default_rng(
        RANDOM_SEED + experiment_index
    )

    received_iq = add_awgn(
        iq_samples=transmitted_iq,
        snr_db=snr_db,
        random_generator=random_generator,
    )

    received_bits = bpsk_demodulate(
        received_iq
    )

    bit_error_count = int(
        np.count_nonzero(
            transmitted_bits != received_bits
        )
    )

    bit_error_rate = (
        bit_error_count
        / len(transmitted_bits)
    )

    received_packet = bits_to_bytes(
        received_bits
    )

    packet_status = "CRC OK"
    restored_message = None

    try:
        parsed_packet = parse_packet(
            received_packet
        )

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

    except CRCError:
        packet_status = "CRC ERROR"

    except PacketError:
        packet_status = "PACKET ERROR"

    except UnicodeDecodeError:
        packet_status = "UTF-8 ERROR"

    experiment_results.append(
        {
            "snr_db": snr_db,
            "bit_errors": bit_error_count,
            "ber": bit_error_rate,
            "packet_status": packet_status,
            "restored_message": restored_message,
        }
    )

    # Для графика достаточно первых 300 символов.
    constellation_samples[snr_db] = (
        received_iq[:300]
    )


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

print("\nРезультаты передачи:")

print(
    f"{'SNR':>8}"
    f"{'Ошибки битов':>16}"
    f"{'BER':>14}"
    f"{'Пакет':>18}"
)

print("-" * 56)

for result in experiment_results:

    print(
        f"{result['snr_db']:>6.1f} дБ"
        f"{result['bit_errors']:>16}"
        f"{result['ber']:>14.6f}"
        f"{result['packet_status']:>18}"
    )


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

for result in experiment_results:

    print(
        f"SNR {result['snr_db']:>4.1f} дБ: "
        f"{result['restored_message']}"
    )


# ============================================================
# График созвездия
# ============================================================

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

axes = axes.ravel()

for axis, snr_db in zip(
    axes,
    SNR_VALUES_DB,
):

    received_iq = constellation_samples[
        snr_db
    ]

    axis.scatter(
        received_iq.real,
        received_iq.imag,
        s=12,
        alpha=0.6,
    )

    axis.axvline(
        0.0,
        linewidth=1,
    )

    axis.set_title(
        f"BPSK, SNR = {snr_db:.1f} дБ"
    )

    axis.set_xlabel(
        "I — синфазная компонента"
    )

    axis.set_ylabel(
        "Q — квадратурная компонента"
    )

    axis.set_xlim(
        -2.5,
        2.5,
    )

    axis.set_ylim(
        -1.8,
        1.8,
    )

    axis.grid(
        True
    )

figure.tight_layout()

constellation_path = (
    OUTPUT_DIRECTORY
    / "lab013_bpsk_constellation.png"
)

figure.savefig(
    constellation_path,
    dpi=150,
)

plt.close(
    figure
)


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

# Без добавления шума преобразование должно быть обратимым.
ideal_received_bits = bpsk_demodulate(
    transmitted_iq
)

ideal_received_packet = bits_to_bytes(
    ideal_received_bits
)

assert ideal_received_packet == original_packet

ideal_parsed_packet = parse_packet(
    ideal_received_packet
)

assert (
    ideal_parsed_packet.payload.decode(
        "utf-8"
    )
    == MESSAGE
)

assert len(experiment_results) == len(
    SNR_VALUES_DB
)

assert constellation_path.exists()


print("\nГрафик созвездия:")
print(constellation_path)

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