"""
Lab014. Измерение BER для BPSK в канале AWGN.

Программа:

1. Генерирует 1 000 000 случайных битов.
2. Преобразует их в BPSK-символы.
3. Добавляет комплексный гауссов шум.
4. Демодулирует принятый сигнал.
5. Измеряет экспериментальный BER.
6. Рассчитывает теоретический BER.
7. Оценивает вероятность повреждения пакетов разной длины.
8. Строит BER-кривую.
9. Сохраняет результаты в CSV.

Упрощения модели:

- один бит передаётся одним BPSK-символом;
- один символ представлен одним IQ-сэмплом;
- частотная и фазовая синхронизация идеальны;
- межсимвольные искажения отсутствуют;
- канал содержит только AWGN.
"""

from csv import DictWriter
from math import (
    erfc,
    expm1,
    log1p,
    sqrt,
)
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


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

BIT_COUNT = 1_000_000

EB_N0_VALUES_DB = [
    -4.0,
    -2.0,
    0.0,
    2.0,
    4.0,
    6.0,
    8.0,
    10.0,
    12.0,
]

RANDOM_SEED = 2026

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

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

GRAPH_PATH = (
    OUTPUT_DIRECTORY
    / "lab014_bpsk_ber_curve.png"
)

CSV_PATH = (
    OUTPUT_DIRECTORY
    / "lab014_bpsk_ber_results.csv"
)


# ============================================================
# Размеры пакетов для оценки PER
# ============================================================

# Пакет из Lab013:
#
# заголовок     8 байт
# PAYLOAD      16 байт
# CRC           4 байта
#
# Итого        28 байт = 224 бита
SHORT_PACKET_BITS = 224

# Типовой пакет изображения:
#
# заголовок протокола       8 байт
# заголовок фрагмента      12 байт
# данные JPEG             512 байт
# CRC                       4 байта
#
# Итого                    536 байт
IMAGE_PACKET_BITS = 536 * 8


# ============================================================
# Вспомогательные функции
# ============================================================

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
    )


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

    Энергия BPSK-символа равна единице.

    Для одного бита на символ:

        Es/N0 = Eb/N0
    """

    eb_n0_linear = 10.0 ** (
        eb_n0_db / 10.0
    )

    # Для комплексного AWGN каждая компонента
    # I и Q получает половину полной мощности шума.
    noise_sigma = sqrt(
        1.0
        / (
            2.0
            * eb_n0_linear
        )
    )

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

    return iq_samples + noise


def bpsk_demodulate(
    received_iq: np.ndarray,
) -> np.ndarray:
    """
    Демодулировать BPSK по знаку компоненты I.

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

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


def theoretical_bpsk_ber(
    eb_n0_db: float,
) -> float:
    """
    Рассчитать теоретический BER когерентной BPSK
    в канале AWGN.

        BER = 0.5 * erfc(sqrt(Eb/N0))
    """

    eb_n0_linear = 10.0 ** (
        eb_n0_db / 10.0
    )

    return 0.5 * erfc(
        sqrt(eb_n0_linear)
    )


def ber_to_per(
    ber: float,
    packet_bit_count: int,
) -> float:
    """
    Оценить Packet Error Rate из BER.

    Предполагается:

    - ошибки отдельных битов независимы;
    - пакет считается повреждённым,
      если ошибся хотя бы один бит.

        PER = 1 - (1 - BER) ** N
    """

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

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

    if ber == 0.0:
        return 0.0

    if ber == 1.0:
        return 1.0

    # Такая запись численно устойчивее,
    # чем прямое возведение в степень
    # для очень маленьких BER.
    return -expm1(
        packet_bit_count
        * log1p(-ber)
    )


# ============================================================
# Генерация исходной последовательности
# ============================================================

bit_generator = np.random.default_rng(
    RANDOM_SEED
)

transmitted_bits = bit_generator.integers(
    low=0,
    high=2,
    size=BIT_COUNT,
    dtype=np.uint8,
)

transmitted_iq = bpsk_modulate(
    transmitted_bits
)


# ============================================================
# Эксперимент при разных Eb/N0
# ============================================================

results = []

for experiment_index, eb_n0_db in enumerate(
    EB_N0_VALUES_DB
):

    noise_generator = np.random.default_rng(
        RANDOM_SEED
        + 1000
        + experiment_index
    )

    received_iq = add_awgn(
        iq_samples=transmitted_iq,
        eb_n0_db=eb_n0_db,
        random_generator=noise_generator,
    )

    received_bits = bpsk_demodulate(
        received_iq
    )

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

    experimental_ber = (
        bit_error_count
        / BIT_COUNT
    )

    theoretical_ber = theoretical_bpsk_ber(
        eb_n0_db
    )

    short_packet_per = ber_to_per(
        theoretical_ber,
        SHORT_PACKET_BITS,
    )

    image_packet_per = ber_to_per(
        theoretical_ber,
        IMAGE_PACKET_BITS,
    )

    results.append(
        {
            "eb_n0_db": eb_n0_db,
            "bit_errors": bit_error_count,
            "experimental_ber": experimental_ber,
            "theoretical_ber": theoretical_ber,
            "short_packet_per": short_packet_per,
            "image_packet_per": image_packet_per,
        }
    )


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

print(
    "=== Lab014. BER-кривая BPSK ==="
)

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

print(
    f"{BIT_COUNT:,}".replace(",", " ")
)

print("\nРезультаты:")

print(
    f"{'Eb/N0':>9}"
    f"{'Ошибки':>12}"
    f"{'BER эксперимент':>19}"
    f"{'BER теория':>16}"
    f"{'PER 224 бит':>16}"
    f"{'PER 4288 бит':>17}"
)

print("-" * 89)

for result in results:

    print(
        f"{result['eb_n0_db']:>6.1f} дБ"
        f"{result['bit_errors']:>12}"
        f"{result['experimental_ber']:>19.6e}"
        f"{result['theoretical_ber']:>16.6e}"
        f"{result['short_packet_per']:>16.6f}"
        f"{result['image_packet_per']:>17.6f}"
    )


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

with CSV_PATH.open(
    "w",
    newline="",
    encoding="utf-8-sig",
) as csv_file:

    fieldnames = [
        "eb_n0_db",
        "bit_errors",
        "experimental_ber",
        "theoretical_ber",
        "short_packet_per_224_bits",
        "image_packet_per_4288_bits",
    ]

    writer = DictWriter(
        csv_file,
        fieldnames=fieldnames,
    )

    writer.writeheader()

    for result in results:

        writer.writerow(
            {
                "eb_n0_db": result["eb_n0_db"],
                "bit_errors": result["bit_errors"],
                "experimental_ber": (
                    result["experimental_ber"]
                ),
                "theoretical_ber": (
                    result["theoretical_ber"]
                ),
                "short_packet_per_224_bits": (
                    result["short_packet_per"]
                ),
                "image_packet_per_4288_bits": (
                    result["image_packet_per"]
                ),
            }
        )


# ============================================================
# Подготовка данных для графика
# ============================================================

eb_n0_plot_values = np.array(
    [
        result["eb_n0_db"]
        for result in results
    ],
    dtype=np.float64,
)

experimental_ber_values = np.array(
    [
        result["experimental_ber"]
        for result in results
    ],
    dtype=np.float64,
)

theoretical_ber_values = np.array(
    [
        result["theoretical_ber"]
        for result in results
    ],
    dtype=np.float64,
)


# Нулевой измеренный BER невозможно показать
# на логарифмической шкале.
#
# Поэтому для графика ставим такую точку
# на уровень половины одного наблюдаемого события.
measurement_floor = (
    0.5 / BIT_COUNT
)

experimental_ber_for_plot = np.maximum(
    experimental_ber_values,
    measurement_floor,
)


# ============================================================
# Построение BER-графика
# ============================================================

figure = plt.figure(
    figsize=(10, 7)
)

plt.semilogy(
    eb_n0_plot_values,
    theoretical_ber_values,
    marker="o",
    label="Теоретический BER BPSK",
)

plt.semilogy(
    eb_n0_plot_values,
    experimental_ber_for_plot,
    marker="s",
    linestyle="--",
    label="Экспериментальный BER",
)

plt.axhline(
    measurement_floor,
    linestyle=":",
    label=(
        "Предел измерения "
        f"{measurement_floor:.1e}"
    ),
)

plt.xlabel(
    "Eb/N0, дБ"
)

plt.ylabel(
    "BER"
)

plt.title(
    "BPSK в канале AWGN: "
    "эксперимент и теория"
)

plt.grid(
    True,
    which="both",
)

plt.legend()

plt.tight_layout()

figure.savefig(
    GRAPH_PATH,
    dpi=160,
)

plt.close(
    figure
)


# ============================================================
# Инженерные контрольные точки
# ============================================================

print("\nИнженерные контрольные точки:")

for target_db in [
    6.0,
    8.0,
    10.0,
]:

    result = next(
        item
        for item in results
        if item["eb_n0_db"] == target_db
    )

    print(
        f"\nEb/N0 = {target_db:.1f} дБ"
    )

    print(
        "Теоретический BER:",
        f"{result['theoretical_ber']:.6e}",
    )

    print(
        "PER короткого пакета 224 бита:",
        f"{result['short_packet_per'] * 100:.2f} %",
    )

    print(
        "PER JPEG-пакета 4288 бит:",
        f"{result['image_packet_per'] * 100:.2f} %",
    )


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

assert len(transmitted_bits) == BIT_COUNT

assert len(received_bits) == BIT_COUNT

assert all(
    0.0
    <= result["experimental_ber"]
    <= 1.0
    for result in results
)

assert all(
    0.0
    <= result["theoretical_ber"]
    <= 1.0
    for result in results
)

assert GRAPH_PATH.exists()

assert CSV_PATH.exists()


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

print(GRAPH_PATH)

print("\nТаблица CSV:")

print(CSV_PATH)

print(
    "\nПроверка пройдена: "
    "экспериментальная BER-кривая построена."
)