"""
Адаптация параметров SDR Rover Link к качеству канала.

Модуль оценивает несколько размеров фрагмента изображения
и выбирает режим с максимальной полезной скоростью.

Используемая модель:

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

from dataclasses import dataclass
from math import (
    erfc,
    exp,
    log1p,
    sqrt,
)

from protocol.image_fragments import (
    FRAGMENT_HEADER_SIZE,
)

from protocol.packet import (
    CRC_SIZE,
    HEADER_SIZE,
)


# ============================================================
# Результат оценки одного размера фрагмента
# ============================================================

@dataclass(frozen=True)
class FragmentModeEstimate:
    """
    Оценка одного возможного размера фрагмента.
    """

    fragment_size: int

    ber: float

    data_packet_size_bytes: int
    ack_packet_size_bytes: int

    data_success_probability: float
    ack_success_probability: float

    confirmed_attempt_probability: float

    expected_attempts: float

    success_probability_with_retries: float

    expected_goodput_bps: float


# ============================================================
# Итоговое решение адаптера
# ============================================================

@dataclass(frozen=True)
class LinkDecision:
    """
    Решение о режиме передачи изображения.
    """

    eb_n0_db: float

    images_enabled: bool

    selected_fragment_size: int | None

    reason: str

    estimates: tuple[FragmentModeEstimate, ...]


# ============================================================
# Теоретический BER
# ============================================================

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 = (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 1.0

    if ber == 1.0:
        return 0.0

    return exp(
        packet_bit_count
        * log1p(-ber)
    )


# ============================================================
# Оценка отдельного режима
# ============================================================

def evaluate_fragment_mode(
    eb_n0_db: float,
    fragment_size: int,
    channel_bitrate_bps: int,
    max_attempts: int,
) -> FragmentModeEstimate:
    """
    Оценить передачу полного фрагмента заданного размера.
    """

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

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

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

    ber = theoretical_bpsk_ber(
        eb_n0_db
    )

    # DATA содержит:
    #
    # заголовок протокола;
    # заголовок фрагмента изображения;
    # данные JPEG;
    # CRC-32.
    data_packet_size_bytes = (
        HEADER_SIZE
        + FRAGMENT_HEADER_SIZE
        + fragment_size
        + CRC_SIZE
    )

    # ACK содержит только заголовок протокола и CRC.
    ack_packet_size_bytes = (
        HEADER_SIZE
        + CRC_SIZE
    )

    data_success_probability = (
        packet_success_probability(
            ber=ber,
            packet_bit_count=(
                data_packet_size_bytes * 8
            ),
        )
    )

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

    # Одна попытка считается завершённой,
    # только если правильно прошли DATA и ACK.
    confirmed_attempt_probability = (
        data_success_probability
        * ack_success_probability
    )

    if confirmed_attempt_probability == 0.0:
        expected_attempts = float("inf")
        expected_goodput_bps = 0.0

    else:
        expected_attempts = (
            1.0
            / confirmed_attempt_probability
        )

        # Ожидаемый объём DATA до подтверждения.
        expected_data_bytes = (
            data_packet_size_bytes
            / confirmed_attempt_probability
        )

        # ACK отправляется только после исправного DATA.
        expected_ack_bytes = (
            ack_packet_size_bytes
            / ack_success_probability
        )

        expected_total_bytes = (
            expected_data_bytes
            + expected_ack_bytes
        )

        expected_goodput_bps = (
            channel_bitrate_bps
            * fragment_size
            / expected_total_bytes
        )

    success_probability_with_retries = (
        1.0
        - (
            1.0
            - confirmed_attempt_probability
        ) ** max_attempts
    )

    return FragmentModeEstimate(
        fragment_size=fragment_size,
        ber=ber,
        data_packet_size_bytes=data_packet_size_bytes,
        ack_packet_size_bytes=ack_packet_size_bytes,
        data_success_probability=(
            data_success_probability
        ),
        ack_success_probability=(
            ack_success_probability
        ),
        confirmed_attempt_probability=(
            confirmed_attempt_probability
        ),
        expected_attempts=expected_attempts,
        success_probability_with_retries=(
            success_probability_with_retries
        ),
        expected_goodput_bps=(
            expected_goodput_bps
        ),
    )


# ============================================================
# Выбор режима
# ============================================================

def choose_image_mode(
    eb_n0_db: float,
    candidate_fragment_sizes: tuple[int, ...] = (
        128,
        512,
        1024,
    ),
    channel_bitrate_bps: int = 20_000,
    max_attempts: int = 5,
    minimum_success_probability: float = 0.85,
    minimum_goodput_bps: float = 2_000.0,
) -> LinkDecision:
    """
    Выбрать режим передачи изображения.

    Режим допускается, если:

    1. вероятность подтверждения за max_attempts
       не ниже minimum_success_probability;

    2. ожидаемый goodput не ниже minimum_goodput_bps.

    Из подходящих режимов выбирается режим
    с максимальной полезной скоростью.
    """

    if not candidate_fragment_sizes:
        raise ValueError(
            "Список размеров фрагментов пуст"
        )

    if not 0.0 <= minimum_success_probability <= 1.0:
        raise ValueError(
            "minimum_success_probability должен быть 0...1"
        )

    if minimum_goodput_bps < 0.0:
        raise ValueError(
            "minimum_goodput_bps не может быть отрицательным"
        )

    estimates = tuple(
        evaluate_fragment_mode(
            eb_n0_db=eb_n0_db,
            fragment_size=fragment_size,
            channel_bitrate_bps=channel_bitrate_bps,
            max_attempts=max_attempts,
        )
        for fragment_size in candidate_fragment_sizes
    )

    eligible_estimates = [
        estimate
        for estimate in estimates
        if (
            estimate.success_probability_with_retries
            >= minimum_success_probability
            and estimate.expected_goodput_bps
            >= minimum_goodput_bps
        )
    ]

    if not eligible_estimates:
        return LinkDecision(
            eb_n0_db=eb_n0_db,
            images_enabled=False,
            selected_fragment_size=None,
            reason=(
                "Канал недостаточен для надёжной "
                "передачи изображения"
            ),
            estimates=estimates,
        )

    selected_estimate = max(
        eligible_estimates,
        key=lambda estimate: (
            estimate.expected_goodput_bps
        ),
    )

    return LinkDecision(
        eb_n0_db=eb_n0_db,
        images_enabled=True,
        selected_fragment_size=(
            selected_estimate.fragment_size
        ),
        reason=(
            "Выбран режим с максимальной "
            "ожидаемой полезной скоростью"
        ),
        estimates=estimates,
    )