"""Программные варианты R=2/3 и 3/4 на базе штатного кода R=1/2.

Прокалывание удаляет известные позиции кодового слова. При декодировании
эти позиции имеют нулевую мягкую метрику (стирание), после чего работает
штатный терминированный декодер Витерби.
"""

from __future__ import annotations

from functools import lru_cache
import numpy as np

from protocol import convolutional_fec


PUNCTURE_PATTERNS: dict[str, tuple[int, ...]] = {
    "1/2": (1, 1),
    "2/3": (1, 1, 1, 0),
    "3/4": (1, 1, 1, 0, 0, 1),
}


@lru_cache(maxsize=32)
def _mask(mother_bit_count: int, rate: str) -> np.ndarray:
    if rate not in PUNCTURE_PATTERNS:
        raise ValueError("неизвестная скорость кода")
    if mother_bit_count < 0:
        raise ValueError("отрицательная длина кодового слова")
    pattern = np.asarray(PUNCTURE_PATTERNS[rate], dtype=bool)
    mask = np.resize(pattern, mother_bit_count)
    mask.flags.writeable = False
    return mask


def coded_bit_count(information_bit_count: int, rate: str) -> int:
    """Точное число передаваемых битов, включая шесть хвостовых битов кода."""
    mother_count = convolutional_fec.encoded_bit_count(information_bit_count)
    if rate not in PUNCTURE_PATTERNS:
        raise ValueError("неизвестная скорость кода")
    pattern = PUNCTURE_PATTERNS[rate]
    periods, remaining = divmod(mother_count, len(pattern))
    return periods * sum(pattern) + sum(pattern[:remaining])


def effective_rate(information_bit_count: int, rate: str) -> float:
    """Отношение полезных битов к реально передаваемым кодовым битам."""
    if information_bit_count <= 0:
        raise ValueError("число полезных битов должно быть положительным")
    return information_bit_count / coded_bit_count(information_bit_count, rate)


def encode(information_bits: np.ndarray, rate: str) -> np.ndarray:
    """Сначала применить штатный код, затем удалить позиции по маске."""
    mother_bits = convolutional_fec.encode(information_bits)
    return mother_bits[_mask(len(mother_bits), rate)]


def depuncture_soft(
    soft_values: np.ndarray, information_bit_count: int, rate: str
) -> np.ndarray:
    """Вернуть стёртые позиции с нулевой мягкой метрикой."""
    values = np.asarray(soft_values, dtype=np.float64)
    if values.ndim != 1 or not np.all(np.isfinite(values)):
        raise ValueError("мягкие значения должны быть конечным одномерным массивом")
    mother_count = convolutional_fec.encoded_bit_count(information_bit_count)
    mask = _mask(mother_count, rate)
    if len(values) != int(np.count_nonzero(mask)):
        raise ValueError("неверная длина проколотого кодового слова")
    restored = np.zeros(mother_count, dtype=np.float64)
    restored[mask] = values
    return restored


def decode_soft(
    soft_values: np.ndarray, information_bit_count: int, rate: str
) -> np.ndarray:
    """Декодировать через существующий 64-состояний декодер Витерби."""
    return convolutional_fec.viterbi_decode_soft(
        depuncture_soft(soft_values, information_bit_count, rate),
        information_bit_count,
    )