"""dpx-decode: dynamic-programming decoders on DPX instructions.

Wavefront DP kernels for speech and alignment pipelines: Viterbi decoding,
monotonic DTW, and CTC forced alignment. Each op has a float32 path and an
int32 fixed-point path running on the DPX fused min/max instructions,
hardware on Hopper and newer (sm_90+), compiler-emulated bit-identically on
Ampere and Ada. Both paths are deterministic and bitwise reproducible:
fixed iteration order, documented tie-breaking, no atomics.

The int32 path quantizes log-domain inputs to fixed point (`quantize`,
default 256 quanta per unit). Viterbi packs (score << 16 | state) so the
3-way max resolves the argmax in the same DPX instruction; scores clamp to
the int16 floor after per-step renormalization, which cannot affect the
argmax.
"""
from typing import Optional, Tuple

import torch

from ._ops import ops

DEFAULT_SCALE = 256.0


def quantize(x: torch.Tensor, scale: float = DEFAULT_SCALE) -> torch.Tensor:
    """Fixed-point quantization for the int32 DPX paths."""
    return torch.round(x * scale).to(torch.int32)


def viterbi(emissions: torch.Tensor, transitions: torch.Tensor,
            priors: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
    """Max-plus Viterbi decoding.

    Args:
        emissions: [B, T, S] float32 or int32 (quantized) CUDA log-scores.
        transitions: [S, S] log transition scores, transitions[i, j] = i -> j.
        priors: [S] initial log-scores (zeros when omitted).
    Returns:
        (path [B, T] int32, score [B]); score dtype is float32 on the float
        path and int64 (in quanta) on the int path. Ties select the larger
        predecessor index on the int path and the smaller on the float path;
        both are deterministic.
    """
    B, T, S = emissions.shape
    if priors is None:
        priors = torch.zeros(S, dtype=emissions.dtype, device=emissions.device)
    path = torch.empty(B, T, dtype=torch.int32, device=emissions.device)
    bp = torch.empty(B, T, S, dtype=torch.int32, device=emissions.device)
    sdt = torch.int64 if emissions.dtype == torch.int32 else torch.float32
    score = torch.empty(B, dtype=sdt, device=emissions.device)
    ops.viterbi(path, score, bp, emissions.contiguous(), transitions.contiguous(),
                priors.contiguous())
    return path, score


def dtw(cost: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Monotonic DTW over a cost matrix (steps: diagonal, down, right).

    Args:
        cost: [B, N, M] float32 or int32 CUDA costs (lower is better).
    Returns:
        (path [B, N+M, 2] int32, start-to-end, -1 padded; path_len [B] int32;
        D [B, N, M] accumulated-cost matrix). Backtrace prefers diagonal,
        then vertical, deterministically.
    """
    B, N, M = cost.shape
    D = torch.empty_like(cost.contiguous())
    path = torch.empty(B, N + M, 2, dtype=torch.int32, device=cost.device)
    path_len = torch.empty(B, dtype=torch.int32, device=cost.device)
    ops.dtw(path, path_len, D, cost.contiguous())
    return path.flip(1), path_len, D  # emitted end-to-start by the backtrace


def ctc_forced_align(log_probs: torch.Tensor, targets: torch.Tensor,
                     blank: int = 0) -> Tuple[torch.Tensor, torch.Tensor]:
    """CTC forced alignment: per-frame labels for a known transcript.

    Args:
        log_probs: [B, T, C] float32 or int32 (quantized) CUDA log-probs.
        targets: [B, L] int64 label ids, -1 padded on the right.
        blank: blank label id.
    Returns:
        (frames [B, T] int32 per-frame label ids including blanks,
        score [B] best-path log-score; int32 in quanta on the int path).
    """
    B, T, C = log_probs.shape
    frames = torch.empty(B, T, dtype=torch.int32, device=log_probs.device)
    sdt = torch.int32 if log_probs.dtype == torch.int32 else torch.float32
    score = torch.empty(B, dtype=sdt, device=log_probs.device)
    ops.ctc_forced_align(frames, score, log_probs.contiguous(), targets.contiguous(), blank)
    return frames, score


__all__ = ["viterbi", "dtw", "ctc_forced_align", "quantize", "DEFAULT_SCALE"]