"""
Lab016. Адаптивный выбор размера фрагмента изображения.

Программа имитирует изменение качества канала во времени.

Для каждой оценки Eb/N0 передатчик решает:

- отключить изображения;
- использовать 128 байт;
- использовать 512 байт;
- использовать 1024 байта.

Дополнительно рассчитывается ожидаемое время передачи
реального JPEG-файла.
"""

from csv import DictWriter
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

from protocol.image_fragments import (
    encode_image_fragment,
    split_image_bytes,
)

from protocol.link_adaptation import (
    choose_image_mode,
    packet_success_probability,
)

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


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

CHANNEL_BITRATE_BPS = 20_000

MAX_ATTEMPTS = 5

CANDIDATE_FRAGMENT_SIZES = (
    128,
    512,
    1024,
)

# Изменение качества канала во времени.
EB_N0_PROFILE_DB = [
    12.0,
    10.0,
    9.0,
    8.0,
    7.0,
    6.0,
    7.0,
    8.0,
    9.0,
    10.0,
    12.0,
]

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

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

CSV_PATH = (
    OUTPUT_DIRECTORY
    / "lab016_adaptation_results.csv"
)

MODE_GRAPH_PATH = (
    OUTPUT_DIRECTORY
    / "lab016_selected_mode.png"
)

TIME_GRAPH_PATH = (
    OUTPUT_DIRECTORY
    / "lab016_expected_transfer_time.png"
)


# ============================================================
# Выбор исходного кадра
# ============================================================

source_candidates = [
    Path(
        "data/processed/lab012/"
        "03_color_320_q15.jpg"
    ),
    Path(
        "data/raw/lab009_source.jpg"
    ),
]

SOURCE_PATH = next(
    (
        path
        for path in source_candidates
        if path.exists()
    ),
    None,
)

if SOURCE_PATH is None:
    raise FileNotFoundError(
        "Не найден кадр для передачи. "
        "Необходимо выполнить Lab009 или Lab012."
    )

source_bytes = SOURCE_PATH.read_bytes()

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


# ============================================================
# Оценка передачи всего изображения
# ============================================================

def estimate_image_transfer(
    image_bytes: bytes,
    fragment_size: int,
    ber: float,
    image_id: int,
) -> dict:
    """
    Оценить передачу полного изображения с бесконечным ARQ.

    Расчёт учитывает фактический размер последнего фрагмента.
    """

    fragments = split_image_bytes(
        image_bytes=image_bytes,
        image_id=image_id,
        fragment_data_size=fragment_size,
    )

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

    ack_size_bytes = len(
        ack_packet
    )

    ack_success_probability = (
        packet_success_probability(
            ber=ber,
            packet_bit_count=(
                ack_size_bytes * 8
            ),
        )
    )

    expected_total_bytes = 0.0

    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_success_probability = (
            packet_success_probability(
                ber=ber,
                packet_bit_count=(
                    len(data_packet) * 8
                ),
            )
        )

        confirmed_probability = (
            data_success_probability
            * ack_success_probability
        )

        expected_total_bytes += (
            len(data_packet)
            / confirmed_probability
        )

        expected_total_bytes += (
            ack_size_bytes
            / ack_success_probability
        )

    expected_seconds = (
        expected_total_bytes
        * 8
        / CHANNEL_BITRATE_BPS
    )

    effective_goodput_bps = (
        len(image_bytes)
        * 8
        / expected_seconds
    )

    return {
        "fragment_count": len(fragments),
        "expected_total_bytes": expected_total_bytes,
        "expected_seconds": expected_seconds,
        "effective_goodput_bps": (
            effective_goodput_bps
        ),
    }


# ============================================================
# Адаптация по профилю канала
# ============================================================

results = []

for step_index, eb_n0_db in enumerate(
    EB_N0_PROFILE_DB
):

    decision = choose_image_mode(
        eb_n0_db=eb_n0_db,
        candidate_fragment_sizes=(
            CANDIDATE_FRAGMENT_SIZES
        ),
        channel_bitrate_bps=(
            CHANNEL_BITRATE_BPS
        ),
        max_attempts=MAX_ATTEMPTS,
        minimum_success_probability=0.85,
        minimum_goodput_bps=2_000.0,
    )

    if decision.images_enabled:

        selected_estimate = next(
            estimate
            for estimate in decision.estimates
            if (
                estimate.fragment_size
                == decision.selected_fragment_size
            )
        )

        image_result = estimate_image_transfer(
            image_bytes=source_bytes,
            fragment_size=(
                decision.selected_fragment_size
            ),
            ber=selected_estimate.ber,
            image_id=2026071700 + step_index,
        )

        fragment_size = (
            decision.selected_fragment_size
        )

        fragment_count = (
            image_result["fragment_count"]
        )

        expected_seconds = (
            image_result["expected_seconds"]
        )

        effective_goodput_bps = (
            image_result[
                "effective_goodput_bps"
            ]
        )

        success_with_retries = (
            selected_estimate
            .success_probability_with_retries
        )

        expected_attempts = (
            selected_estimate.expected_attempts
        )

        mode_name = (
            f"{fragment_size} B"
        )

    else:
        fragment_size = 0
        fragment_count = 0
        expected_seconds = None
        effective_goodput_bps = 0.0
        success_with_retries = 0.0
        expected_attempts = 0.0
        mode_name = "IMAGE OFF"

    results.append(
        {
            "step": step_index,
            "eb_n0_db": eb_n0_db,
            "mode": mode_name,
            "fragment_size": fragment_size,
            "fragment_count": fragment_count,
            "expected_seconds": expected_seconds,
            "effective_goodput_bps": (
                effective_goodput_bps
            ),
            "success_with_retries": (
                success_with_retries
            ),
            "expected_attempts": (
                expected_attempts
            ),
            "reason": decision.reason,
        }
    )


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

print(
    "=== Lab016. Адаптация размера фрагмента ==="
)

print("\nИсходный кадр:")

print(SOURCE_PATH)

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

print(
    len(source_bytes),
    "байт",
)

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

print(
    f"{'Шаг':>5}"
    f"{'Eb/N0':>10}"
    f"{'Режим':>14}"
    f"{'Фрагм.':>9}"
    f"{'Попыток':>11}"
    f"{'Успех x5':>12}"
    f"{'Время кадра':>15}"
    f"{'Goodput':>13}"
)

print("-" * 89)

for result in results:

    if result["expected_seconds"] is None:
        time_text = "—"
    else:
        time_text = (
            f"{result['expected_seconds']:.2f} с"
        )

    print(
        f"{result['step']:>5}"
        f"{result['eb_n0_db']:>7.1f} дБ"
        f"{result['mode']:>14}"
        f"{result['fragment_count']:>9}"
        f"{result['expected_attempts']:>11.2f}"
        f"{result['success_with_retries'] * 100:>10.1f} %"
        f"{time_text:>15}"
        f"{result['effective_goodput_bps'] / 1000:>10.2f} кбит/с"
    )


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

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

    fieldnames = list(
        results[0].keys()
    )

    writer = DictWriter(
        csv_file,
        fieldnames=fieldnames,
    )

    writer.writeheader()
    writer.writerows(results)


# ============================================================
# График выбранного режима
# ============================================================

steps = [
    result["step"]
    for result in results
]

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

plt.figure(
    figsize=(11, 6)
)

plt.step(
    steps,
    fragment_sizes,
    where="mid",
    marker="o",
)

plt.yticks(
    [
        0,
        128,
        512,
        1024,
    ],
    [
        "IMAGE OFF",
        "128 B",
        "512 B",
        "1024 B",
    ],
)

plt.xlabel(
    "Шаг времени"
)

plt.ylabel(
    "Выбранный режим"
)

plt.title(
    "Автоматический выбор размера фрагмента"
)

plt.grid(
    True
)

plt.tight_layout()

plt.savefig(
    MODE_GRAPH_PATH,
    dpi=160,
)

plt.close()


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

transfer_times = np.array(
    [
        (
            result["expected_seconds"]
            if result["expected_seconds"] is not None
            else np.nan
        )
        for result in results
    ],
    dtype=np.float64,
)

plt.figure(
    figsize=(11, 6)
)

plt.plot(
    steps,
    transfer_times,
    marker="o",
)

plt.xlabel(
    "Шаг времени"
)

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

plt.title(
    "Время передачи кадра при адаптации канала"
)

plt.grid(
    True
)

plt.tight_layout()

plt.savefig(
    TIME_GRAPH_PATH,
    dpi=160,
)

plt.close()


# ============================================================
# Проверки ожидаемой логики
# ============================================================

decision_at_6_db = next(
    result
    for result in results
    if result["eb_n0_db"] == 6.0
)

decision_at_8_db = next(
    result
    for result in results
    if result["eb_n0_db"] == 8.0
)

decision_at_9_db = next(
    result
    for result in results
    if result["eb_n0_db"] == 9.0
)

decision_at_10_db = next(
    result
    for result in results
    if result["eb_n0_db"] == 10.0
)

assert (
    decision_at_6_db["fragment_size"]
    == 0
)

assert (
    decision_at_8_db["fragment_size"]
    == 128
)

assert (
    decision_at_9_db["fragment_size"]
    == 512
)

assert (
    decision_at_10_db["fragment_size"]
    == 1024
)

assert CSV_PATH.exists()
assert MODE_GRAPH_PATH.exists()
assert TIME_GRAPH_PATH.exists()


print("\nCSV:")

print(CSV_PATH)

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

print(MODE_GRAPH_PATH)

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

print(TIME_GRAPH_PATH)

print(
    "\nПроверка пройдена: "
    "адаптивный выбор режима работает."
)