"""Soft/hard Viterbi decoder for Galileo rate-1/2, K=7 convolutional FEC."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

import numpy as np
import numpy.typing as npt


@dataclass(frozen=True, slots=True)
class FecDecodeResult:
    decoded_bits: list[int]
    path_metric: float
    corrected_error_metric: float
    valid: bool


def _reverse_seven_bits(value: int) -> int:
    return int(f"{value:07b}"[::-1], 2)


class InavViterbiDecoder:
    """ICD 2.2 Table 26 decoder with deterministic predecessor tie-breaking."""

    constraint_length = 7
    generator_polynomials_octal = (0o171, 0o133)

    def __init__(self, traceback_depth: int = 35) -> None:
        if traceback_depth < 1:
            raise ValueError("traceback_depth must be positive")
        self.traceback_depth = int(traceback_depth)
        self._polynomials = tuple(_reverse_seven_bits(item) for item in self.generator_polynomials_octal)
        self._next = np.empty((64, 2), dtype=np.int16)
        self._outputs = np.empty((64, 2, 2), dtype=np.int8)
        for state in range(64):
            for bit in (0, 1):
                register = ((state << 1) | bit) & 0x7F
                self._next[state, bit] = register & 0x3F
                self._outputs[state, bit, 0] = (register & self._polynomials[0]).bit_count() & 1
                # Figure 13 explicitly inverts the second encoder branch.
                self._outputs[state, bit, 1] = ((register & self._polynomials[1]).bit_count() & 1) ^ 1

    def decode(
        self,
        received: npt.ArrayLike,
        *,
        mode: Literal["soft", "hard"] = "soft",
        initial_state: int = 0,
        terminal_state: int | None = 0,
    ) -> FecDecodeResult:
        values = np.asarray(received)
        if values.ndim != 1 or values.size == 0 or values.size % 2:
            raise ValueError("Viterbi input must be a nonempty even-length symbol vector")
        if initial_state not in range(64) or terminal_state is not None and terminal_state not in range(64):
            raise ValueError("Viterbi states must be in 0..63")
        if mode == "hard":
            hard = values.astype(np.int8, copy=False)
            if not np.all((hard == 0) | (hard == 1)):
                raise ValueError("hard Viterbi symbols must be 0 or 1")
            observations = hard.reshape(-1, 2)
        elif mode == "soft":
            observations = values.astype(np.float64, copy=False).reshape(-1, 2)
            if not np.all(np.isfinite(observations)):
                raise ValueError("soft Viterbi symbols must be finite")
        else:
            raise ValueError("Viterbi mode must be 'soft' or 'hard'")

        steps = observations.shape[0]
        infinity = np.float64(np.inf)
        metrics = np.full(64, infinity, dtype=np.float64)
        metrics[initial_state] = 0.0
        predecessors = np.full((steps, 64), -1, dtype=np.int16)
        decisions = np.full((steps, 64), -1, dtype=np.int8)
        for index, observation in enumerate(observations):
            updated = np.full(64, infinity, dtype=np.float64)
            for state in range(64):
                if not np.isfinite(metrics[state]):
                    continue
                for bit in (0, 1):
                    target = int(self._next[state, bit])
                    expected = self._outputs[state, bit]
                    if mode == "hard":
                        branch = np.float64(np.count_nonzero(observation != expected))
                    else:
                        expected_sign = np.where(expected == 1, 1.0, -1.0)
                        branch = np.float64(-np.dot(observation, expected_sign))
                    candidate = metrics[state] + branch
                    # Strict comparison preserves the first (lowest state, then
                    # input 0) predecessor for deterministic ties.
                    if candidate < updated[target]:
                        updated[target] = candidate
                        predecessors[index, target] = state
                        decisions[index, target] = bit
            metrics = updated

        final_state = int(np.argmin(metrics)) if terminal_state is None else terminal_state
        if not np.isfinite(metrics[final_state]):
            return FecDecodeResult([], float("inf"), float("inf"), False)
        decoded = np.empty(steps, dtype=np.int8)
        state = final_state
        for index in range(steps - 1, -1, -1):
            decoded[index] = decisions[index, state]
            state = int(predecessors[index, state])
            if state < 0 and index:
                return FecDecodeResult([], float(metrics[final_state]), float("inf"), False)
        encoded = self.encode(decoded, initial_state=initial_state)
        received_hard = (
            values.astype(np.int8, copy=False)
            if mode == "hard"
            else np.where(values.astype(np.float64) >= 0.0, 1, 0).astype(np.int8)
        )
        corrected = float(np.count_nonzero(encoded != received_hard))
        tail_valid = decoded.size >= 6 and bool(np.all(decoded[-6:] == 0))
        state_valid = terminal_state is None or final_state == terminal_state
        return FecDecodeResult(
            decoded_bits=decoded.astype(int).tolist(),
            path_metric=float(metrics[final_state]),
            corrected_error_metric=corrected,
            valid=bool(tail_valid and state_valid),
        )

    def encode(self, bits: npt.ArrayLike, *, initial_state: int = 0) -> npt.NDArray[np.int8]:
        """Diagnostic encoder; decoder tests use the independent ICD Annex D vector."""
        source = np.asarray(bits, dtype=np.int8).reshape(-1)
        if not np.all((source == 0) | (source == 1)):
            raise ValueError("encoder bits must be 0 or 1")
        state = initial_state
        output = np.empty(source.size * 2, dtype=np.int8)
        for index, bit_value in enumerate(source):
            bit = int(bit_value)
            output[2 * index : 2 * index + 2] = self._outputs[state, bit]
            state = int(self._next[state, bit])
        return output