"""Lab053: те же переходы soft Viterbi, циклы компилируются только для CPU."""

from __future__ import annotations

from pathlib import Path
import sys
import time

import numpy as np

if __package__ in (None, ""):
    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from experiments import lab051_streaming_receiver as previous

runtime = Path(__file__).resolve().parents[1] / ".lab053_cpu_runtime/packages"
if runtime.is_dir():
    sys.path.insert(0, str(runtime))
try:
    from numba import njit
except ImportError:
    njit = None

conv = previous.lab045.conv
symbols = conv._BRANCH_SYMBOLS
CODES = (2 * (symbols[:, :, 0] > 0) + (symbols[:, :, 1] > 0)).astype(np.intp)
PREDECESSORS = conv._PREDECESSORS.copy()
CODES.flags.writeable = PREDECESSORS.flags.writeable = False


def trellis_decode(metrics, predecessors, codes):
    """Порядок сложений исходный; строгое < сохраняет tie в ветвь 0.

    Не объединяем несколько шагов: это могло бы изменить решение при
    равных округлённых суммах. Старт/финиш 0, 64 состояния, хвост 6 бит.
    """
    step_count = metrics.shape[0]
    path = np.full(64, 1e18, dtype=np.float64)
    path[0] = 0.0
    next_path = np.empty(64, np.float64)
    choices = np.empty((step_count, 64), np.uint8)
    for step in range(step_count):
        for state in range(64):
            zero = path[predecessors[state, 0]] + metrics[step, codes[state, 0]]
            one = path[predecessors[state, 1]] + metrics[step, codes[state, 1]]
            take_one = one < zero
            choices[step, state] = take_one
            next_path[state] = one if take_one else zero
        path, next_path = next_path, path
    decoded = np.empty(step_count, np.uint8)
    state = 0
    for step in range(step_count - 1, -1, -1):
        decoded[step] = state >> 5
        state = predecessors[state, choices[step, state]]
    assert state == 0
    return decoded[:-6]


compiled_trellis = njit(cache=True, nogil=True, fastmath=False)(trellis_decode) if njit is not None else None


def viterbi_decode_soft(soft_values, information_bit_count):
    if compiled_trellis is None:
        return previous.fast_viterbi(soft_values, information_bit_count)
    if not isinstance(information_bit_count, (int, np.integer)) or information_bit_count <= 0:
        return previous.ORIGINAL_VITERBI(soft_values, information_bit_count)
    values = np.asarray(soft_values, dtype=np.float64)
    if (values.ndim != 1 or values.size != conv.encoded_bit_count(information_bit_count)
            or not np.all(np.isfinite(values))):
        return previous.ORIGINAL_VITERBI(soft_values, information_bit_count)
    received = values.reshape(int(information_bit_count) + 6, 2)
    negative = (received + 1.0) ** 2
    positive = (received - 1.0) ** 2
    metrics = np.column_stack((negative[:, 0] + negative[:, 1], negative[:, 0] + positive[:, 1],
                               positive[:, 0] + negative[:, 1], positive[:, 0] + positive[:, 1]))
    return compiled_trellis(metrics, PREDECESSORS, CODES)


def warmup() -> dict:
    started = time.perf_counter()
    values = np.ones(conv.encoded_bit_count(8), np.float64)
    assert np.array_equal(viterbi_decode_soft(values, 8), previous.fast_viterbi(values, 8))
    return {"backend": "NUMBA_CPU" if compiled_trellis is not None else "NUMPY_LAB051",
            "warmup_seconds": time.perf_counter() - started, "fastmath": False, "gpu_used": False}


def run_functional_tests() -> None:
    values = np.zeros(conv.encoded_bit_count(8), np.float64)
    assert np.array_equal(viterbi_decode_soft(values, 8), previous.fast_viterbi(values, 8))


def main() -> None:
    print(warmup())
    run_functional_tests()


if __name__ == "__main__":
    main()