"""
Lab015. Подбор оптимального размера фрагмента изображения.

Программа сравнивает фрагменты размером:

- 64 байта;
- 128 байт;
- 256 байт;
- 512 байт;
- 1024 байта.

Для каждого размера рассчитываются:

1. Количество фрагментов изображения.
2. Служебные расходы протокола.
3. Теоретический BER BPSK.
4. Вероятность успешной доставки DATA-пакета.
5. Вероятность успешной доставки ACK.
6. Среднее число передач одного фрагмента.
7. Ожидаемый полный радиообмен с ARQ.
8. Время передачи фотографии при 20 кбит/с.
9. Эффективная полезная скорость.

Модель канала:

- когерентная BPSK;
- AWGN;
- ошибки битов независимы;
- повреждённый пакет отбрасывается по CRC;
- применяется Stop-and-Wait ARQ;
- DATA и ACK работают при одинаковом Eb/N0.

Не учитываются:

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

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

import matplotlib.pyplot as plt

from protocol.image_fragments import (
    encode_image_fragment,
    split_image_bytes,
)

from protocol.packet import (
    MESSAGE_TYPE_ACK,
    MESSAGE_TYPE_IMAGE_FRAGMENT,
    build_packet,
)


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

SOURCE_PATH = Path(
    "data/raw/lab009_source.jpg"
)

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

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

FRAGMENT_SIZES = [
    64,
    128,
    256,
    512,
    1024,
]

EB_N0_VALUES_DB = [
    6.0,
    8.0,
    10.0,
    12.0,
]

CHANNEL_BITRATE_BPS = 20_000

IMAGE_ID_BASE = 2026071600

CSV_PATH = (
    OUTPUT_DIRECTORY
    / "lab015_fragment_results.csv"
)

TIME_GRAPH_PATH = (
    OUTPUT_DIRECTORY
    / "lab015_transfer_time.png"
)

GOODPUT_GRAPH_PATH = (
    OUTPUT_DIRECTORY
    / "lab015_effective_goodput.png"
)


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

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 packet_success_probability(
    ber: float,
    packet_bit_count: int,
) -> float:
    """
    Вероятность того, что весь пакет будет принят
    без единой битовой ошибки.

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

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

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

    if ber == 0.0:
        return 1.0

    if ber == 1.0:
        return 0.0

    # Численно устойчивый вариант выражения:
    #
    # (1 - BER) ** N
    return exp(
        packet_bit_count
        * log1p(-ber)
    )


def format_bytes(
    byte_count: float,
) -> str:
    """
    Представить объём в удобном виде.
    """

    if byte_count < 1024:
        return f"{byte_count:.0f} байт"

    kibibytes = byte_count / 1024

    if kibibytes < 1024:
        return f"{kibibytes:.2f} КиБ"

    mebibytes = kibibytes / 1024

    return f"{mebibytes:.2f} МиБ"


def format_duration(
    seconds: float,
) -> str:
    """
    Представить длительность в удобном виде.
    """

    if seconds < 1:
        return f"{seconds * 1000:.0f} мс"

    if seconds < 60:
        return f"{seconds:.2f} с"

    if seconds < 3600:
        minutes = int(seconds // 60)
        remaining_seconds = seconds % 60

        return (
            f"{minutes} мин "
            f"{remaining_seconds:.1f} с"
        )

    hours = int(seconds // 3600)
    remaining_minutes = (
        seconds % 3600
    ) / 60

    return (
        f"{hours} ч "
        f"{remaining_minutes:.1f} мин"
    )


# ============================================================
# Проверка исходного JPEG
# ============================================================

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

source_bytes = SOURCE_PATH.read_bytes()

if not source_bytes:
    raise ValueError(
        "Исходный JPEG пуст"
    )

source_size = len(source_bytes)


# ============================================================
# Размер ACK
# ============================================================

ack_packet = build_packet(
    payload=b"",
    message_type=MESSAGE_TYPE_ACK,
    sequence_number=0,
)

ack_packet_size_bytes = len(
    ack_packet
)

ack_packet_size_bits = (
    ack_packet_size_bytes * 8
)


# ============================================================
# Основной расчёт
# ============================================================

results = []

for fragment_size_index, fragment_size in enumerate(
    FRAGMENT_SIZES
):

    fragments = split_image_bytes(
        image_bytes=source_bytes,
        image_id=(
            IMAGE_ID_BASE
            + fragment_size_index
        ),
        fragment_data_size=fragment_size,
    )

    data_packet_sizes_bytes = []

    for fragment in fragments:

        fragment_payload = encode_image_fragment(
            fragment
        )

        data_packet = build_packet(
            payload=fragment_payload,
            message_type=(
                MESSAGE_TYPE_IMAGE_FRAGMENT
            ),
            sequence_number=(
                fragment.fragment_index
            ),
        )

        data_packet_sizes_bytes.append(
            len(data_packet)
        )

    ideal_data_bytes = sum(
        data_packet_sizes_bytes
    )

    ideal_ack_bytes = (
        len(fragments)
        * ack_packet_size_bytes
    )

    ideal_total_bytes = (
        ideal_data_bytes
        + ideal_ack_bytes
    )

    ideal_efficiency_percent = (
        source_size
        / ideal_total_bytes
        * 100
    )

    for eb_n0_db in EB_N0_VALUES_DB:

        ber = theoretical_bpsk_ber(
            eb_n0_db
        )

        ack_success_probability = (
            packet_success_probability(
                ber=ber,
                packet_bit_count=(
                    ack_packet_size_bits
                ),
            )
        )

        expected_data_bytes = 0.0
        expected_ack_bytes = 0.0

        expected_data_transmissions = 0.0

        data_success_probabilities = []

        for data_packet_size_bytes in (
            data_packet_sizes_bytes
        ):

            data_packet_size_bits = (
                data_packet_size_bytes * 8
            )

            data_success_probability = (
                packet_success_probability(
                    ber=ber,
                    packet_bit_count=(
                        data_packet_size_bits
                    ),
                )
            )

            data_success_probabilities.append(
                data_success_probability
            )

            # Для успешного завершения попытки должны
            # одновременно правильно пройти DATA и ACK.
            confirmed_attempt_probability = (
                data_success_probability
                * ack_success_probability
            )

            if confirmed_attempt_probability == 0.0:
                raise RuntimeError(
                    "Вероятность подтверждения "
                    "оказалась равной нулю"
                )

            # Число DATA-передач до успешного ACK
            # подчиняется геометрическому распределению.
            expected_attempt_count = (
                1.0
                / confirmed_attempt_probability
            )

            expected_data_transmissions += (
                expected_attempt_count
            )

            expected_data_bytes += (
                data_packet_size_bytes
                * expected_attempt_count
            )

            # ACK передаётся только после правильно
            # принятого DATA.
            #
            # Среднее количество передач ACK до
            # успешного ACK равно 1 / P_ACK.
            expected_ack_bytes += (
                ack_packet_size_bytes
                / ack_success_probability
            )

        expected_total_bytes = (
            expected_data_bytes
            + expected_ack_bytes
        )

        expected_transfer_seconds = (
            expected_total_bytes
            * 8
            / CHANNEL_BITRATE_BPS
        )

        expected_efficiency_percent = (
            source_size
            / expected_total_bytes
            * 100
        )

        effective_goodput_bps = (
            source_size
            * 8
            / expected_transfer_seconds
        )

        average_data_transmissions = (
            expected_data_transmissions
            / len(fragments)
        )

        minimum_data_success_probability = min(
            data_success_probabilities
        )

        maximum_data_success_probability = max(
            data_success_probabilities
        )

        results.append(
            {
                "fragment_size": fragment_size,
                "fragment_count": len(fragments),
                "eb_n0_db": eb_n0_db,
                "ber": ber,
                "ideal_total_bytes": (
                    ideal_total_bytes
                ),
                "ideal_efficiency_percent": (
                    ideal_efficiency_percent
                ),
                "ack_success_probability": (
                    ack_success_probability
                ),
                "minimum_data_success_probability": (
                    minimum_data_success_probability
                ),
                "maximum_data_success_probability": (
                    maximum_data_success_probability
                ),
                "average_data_transmissions": (
                    average_data_transmissions
                ),
                "expected_total_bytes": (
                    expected_total_bytes
                ),
                "expected_transfer_seconds": (
                    expected_transfer_seconds
                ),
                "expected_efficiency_percent": (
                    expected_efficiency_percent
                ),
                "effective_goodput_bps": (
                    effective_goodput_bps
                ),
            }
        )


# ============================================================
# Вывод общих данных
# ============================================================

print(
    "=== Lab015. Оптимальный размер фрагмента ==="
)

print("\nИсходный JPEG:")

print(SOURCE_PATH)

print("\nРазмер JPEG:")

print(
    format_bytes(source_size)
)

print("\nСкорость физического канала:")

print(
    CHANNEL_BITRATE_BPS / 1000,
    "кбит/с",
)

print("\nРазмер ACK-пакета:")

print(
    ack_packet_size_bytes,
    "байт",
)


# ============================================================
# Вывод результатов отдельно для каждого Eb/N0
# ============================================================

for eb_n0_db in EB_N0_VALUES_DB:

    selected_results = [
        result
        for result in results
        if result["eb_n0_db"] == eb_n0_db
    ]

    print(
        f"\n--- Eb/N0 = {eb_n0_db:.1f} дБ ---"
    )

    print(
        f"{'Фрагмент':>10}"
        f"{'Кол-во':>9}"
        f"{'P DATA':>12}"
        f"{'Попыток':>11}"
        f"{'Объём с ARQ':>16}"
        f"{'Время':>16}"
        f"{'Goodput':>13}"
    )

    print("-" * 87)

    for result in selected_results:

        print(
            f"{result['fragment_size']:>8} Б"
            f"{result['fragment_count']:>9}"
            f"{result['minimum_data_success_probability']:>12.6f}"
            f"{result['average_data_transmissions']:>11.2f}"
            f"{format_bytes(result['expected_total_bytes']):>16}"
            f"{format_duration(result['expected_transfer_seconds']):>16}"
            f"{result['effective_goodput_bps'] / 1000:>10.2f} кбит/с"
        )

    optimum_result = min(
        selected_results,
        key=lambda item: (
            item["expected_transfer_seconds"]
        ),
    )

    print(
        "\nОптимальный размер фрагмента:"
    )

    print(
        optimum_result["fragment_size"],
        "байт",
    )

    print(
        "Ожидаемое время:",
        format_duration(
            optimum_result[
                "expected_transfer_seconds"
            ]
        ),
    )

    print(
        "Полезная скорость:",
        f"{optimum_result['effective_goodput_bps'] / 1000:.2f}",
        "кбит/с",
    )


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

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

    fieldnames = [
        "fragment_size",
        "fragment_count",
        "eb_n0_db",
        "ber",
        "ideal_total_bytes",
        "ideal_efficiency_percent",
        "ack_success_probability",
        "minimum_data_success_probability",
        "maximum_data_success_probability",
        "average_data_transmissions",
        "expected_total_bytes",
        "expected_transfer_seconds",
        "expected_efficiency_percent",
        "effective_goodput_bps",
    ]

    writer = DictWriter(
        csv_file,
        fieldnames=fieldnames,
    )

    writer.writeheader()

    writer.writerows(results)


# ============================================================
# График ожидаемого времени передачи
# ============================================================

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

for eb_n0_db in EB_N0_VALUES_DB:

    selected_results = [
        result
        for result in results
        if result["eb_n0_db"] == eb_n0_db
    ]

    fragment_sizes = [
        result["fragment_size"]
        for result in selected_results
    ]

    transfer_times = [
        result["expected_transfer_seconds"]
        for result in selected_results
    ]

    plt.plot(
        fragment_sizes,
        transfer_times,
        marker="o",
        label=f"Eb/N0 = {eb_n0_db:.0f} дБ",
    )

plt.xscale(
    "log",
    base=2,
)

plt.yscale(
    "log",
)

plt.xlabel(
    "Размер данных фрагмента, байт"
)

plt.ylabel(
    "Ожидаемое время передачи, с"
)

plt.title(
    "Влияние размера фрагмента "
    "на время передачи JPEG"
)

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

plt.legend()

plt.tight_layout()

plt.savefig(
    TIME_GRAPH_PATH,
    dpi=160,
)

plt.close()


# ============================================================
# График эффективной полезной скорости
# ============================================================

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

for eb_n0_db in EB_N0_VALUES_DB:

    selected_results = [
        result
        for result in results
        if result["eb_n0_db"] == eb_n0_db
    ]

    fragment_sizes = [
        result["fragment_size"]
        for result in selected_results
    ]

    goodput_values = [
        result["effective_goodput_bps"] / 1000
        for result in selected_results
    ]

    plt.plot(
        fragment_sizes,
        goodput_values,
        marker="o",
        label=f"Eb/N0 = {eb_n0_db:.0f} дБ",
    )

plt.xscale(
    "log",
    base=2,
)

plt.xlabel(
    "Размер данных фрагмента, байт"
)

plt.ylabel(
    "Полезная скорость, кбит/с"
)

plt.title(
    "Полезная скорость JPEG "
    "с учётом CRC и ARQ"
)

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

plt.legend()

plt.tight_layout()

plt.savefig(
    GOODPUT_GRAPH_PATH,
    dpi=160,
)

plt.close()


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

assert results

assert CSV_PATH.exists()

assert TIME_GRAPH_PATH.exists()

assert GOODPUT_GRAPH_PATH.exists()

assert all(
    result["expected_transfer_seconds"] > 0
    for result in results
)

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


print("\nCSV:")

print(CSV_PATH)

print("\nГрафик времени:")

print(TIME_GRAPH_PATH)

print("\nГрафик полезной скорости:")

print(GOODPUT_GRAPH_PATH)

print(
    "\nПроверка пройдена: "
    "оптимальный размер фрагмента рассчитан."
)