"""CPU JIT для той же 64-состояний решётки Витерби VIDEO-004.

Режим выбирается вызывающим кодом до приёма пакета. При отсутствии Numba
этот модуль не подменяет результат медленным декодером после ошибки CRC.
"""

from __future__ import annotations

from typing import Any, Callable, cast

import numpy as np

from protocol import convolutional_fec, punctured_fec

try:
    from numba import njit
except ImportError:
    njit = None


PREDECESSORS = convolutional_fec._PREDECESSORS
BRANCH_SYMBOLS = convolutional_fec._BRANCH_SYMBOLS


def _trellis(soft_values: np.ndarray, information_bit_count: int,
             predecessors: np.ndarray, branch_symbols: np.ndarray) -> np.ndarray:
    """Те же метрики, порядок сложения и выбор ветви 0 при равенстве."""
    step_count = information_bit_count + convolutional_fec.TAIL_BIT_COUNT
    received = soft_values.reshape(step_count, 2)
    path_metric = np.full(convolutional_fec.STATE_COUNT, 1e18, dtype=np.float64)
    path_metric[0] = 0.0
    next_metric = np.empty(convolutional_fec.STATE_COUNT, dtype=np.float64)
    choices = np.empty((step_count, convolutional_fec.STATE_COUNT), dtype=np.uint8)
    for step in range(step_count):
        received_0 = received[step, 0]
        received_1 = received[step, 1]
        for state in range(convolutional_fec.STATE_COUNT):
            difference_00 = received_0 - branch_symbols[state, 0, 0]
            difference_01 = received_1 - branch_symbols[state, 0, 1]
            difference_10 = received_0 - branch_symbols[state, 1, 0]
            difference_11 = received_1 - branch_symbols[state, 1, 1]
            first = (path_metric[predecessors[state, 0]]
                     + difference_00 ** 2 + difference_01 ** 2)
            second = (path_metric[predecessors[state, 1]]
                      + difference_10 ** 2 + difference_11 ** 2)
            take_second = second < first
            choices[step, state] = take_second
            next_metric[state] = second if take_second else first
        path_metric, next_metric = next_metric, path_metric
    decoded = np.empty(step_count, dtype=np.uint8)
    current_state = 0
    for step in range(step_count - 1, -1, -1):
        branch = choices[step, current_state]
        decoded[step] = current_state >> (convolutional_fec.MEMORY - 1)
        current_state = predecessors[current_state, branch]
    return decoded[:information_bit_count]


_compiler = cast(Any, njit)  # У Numba нет точной типизации фабрики декоратора.
COMPILED_TRELLIS: Callable[[np.ndarray, int, np.ndarray, np.ndarray], np.ndarray] | None = (
    cast(Callable[[np.ndarray, int, np.ndarray, np.ndarray], np.ndarray],
         _compiler(cache=False, nogil=True, fastmath=False)(_trellis))
    if njit is not None else None
)


def decode_soft(soft_values: np.ndarray, information_bit_count: int,
                rate: str) -> np.ndarray:
    """Проверить вход и декодировать проколотое слово без запасного алгоритма."""
    if COMPILED_TRELLIS is None:
        raise RuntimeError("Numba недоступна для VIDEO-004")
    restored = punctured_fec.depuncture_soft(soft_values, information_bit_count, rate)
    return COMPILED_TRELLIS(restored, int(information_bit_count),
                            PREDECESSORS, BRANCH_SYMBOLS)