// Dynamic-programming decoders on DPX instructions.
//
// Three wavefront DP kernels for speech and alignment pipelines:
//   viterbi           max-plus HMM decoding with backtrace
//   dtw               monotonic dynamic time warping (min-plus, 3-way)
//   ctc_forced_align  CTC trellis Viterbi aligning a transcript to emissions
//
// Each op has two numeric paths. The float32 path matches float reference
// implementations. The int32 path quantizes log-domain values to fixed point
// (scale quanta per unit) and runs on the DPX fused min/max instructions
// (__vimin3_s32 / __vimax3_s32), single hardware instructions on sm_90+ and
// bit-identically emulated by the compiler on sm_80-89. Viterbi packs
// (score << 16 | argument) into one int32 so the 3-way max resolves the
// argmax in the same instruction; per-timestep renormalization keeps packed
// scores inside 16 bits. Both paths are deterministic: fixed iteration
// order, no atomics, bitwise-reproducible outputs.

#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#include <cfloat>
#include <climits>
#include <cstdint>

namespace {

constexpr int kThreads = 256;
constexpr int kNegBig = -(1 << 24);  // safely below any renormalized score

__device__ __forceinline__ int imax3(int a, int b, int c) {
#if __CUDA_ARCH__ >= 800
  return __vimax3_s32(a, b, c);
#else
  return max(a, max(b, c));
#endif
}

__device__ __forceinline__ int imin3(int a, int b, int c) {
#if __CUDA_ARCH__ >= 800
  return __vimin3_s32(a, b, c);
#else
  return min(a, min(b, c));
#endif
}

__device__ int block_imax(int v) {
  __shared__ int smem[kThreads];
  smem[threadIdx.x] = v;
  __syncthreads();
  for (int s = kThreads / 2; s > 0; s >>= 1) {
    if (threadIdx.x < s) smem[threadIdx.x] = max(smem[threadIdx.x], smem[threadIdx.x + s]);
    __syncthreads();
  }
  int r = smem[0];
  __syncthreads();
  return r;
}

// ---------------------------------------------------------------------------
// Viterbi, int32 DPX path. One block per sequence. States parallel over
// threads; time sequential. delta values are renormalized each step so the
// packed (score << 16 | state) representation stays exact.
// ---------------------------------------------------------------------------

__global__ void viterbi_i32_kernel(const int32_t* __restrict__ emit /* [B,T,S] */,
                                   const int32_t* __restrict__ trans /* [S,S] */,
                                   const int32_t* __restrict__ prior /* [S] */, int64_t B,
                                   int64_t T, int64_t S, int32_t* __restrict__ bp /* [B,T,S] */,
                                   int32_t* __restrict__ path /* [B,T] */,
                                   int64_t* __restrict__ score /* [B] */,
                                   int32_t* __restrict__ work /* [B, 2S] */) {
  const int64_t b = blockIdx.x;
  int32_t* prev = work + b * 2 * S;
  int32_t* curr = prev + S;
  __shared__ int64_t norm_acc;
  if (threadIdx.x == 0) norm_acc = 0;

  for (int64_t j = threadIdx.x; j < S; j += blockDim.x) {
    prev[j] = prior[j] + emit[(b * T + 0) * S + j];
  }
  __syncthreads();

  for (int64_t t = 1; t < T; t++) {
    // renormalize prev to <= 0 so packed scores fit 16 bits
    int mx = INT_MIN;
    for (int64_t j = threadIdx.x; j < S; j += blockDim.x) mx = max(mx, prev[j]);
    mx = block_imax(mx);
    if (threadIdx.x == 0) norm_acc += mx;
    for (int64_t j = threadIdx.x; j < S; j += blockDim.x) prev[j] -= mx;
    __syncthreads();

    for (int64_t j = threadIdx.x; j < S; j += blockDim.x) {
      // scores clamp to the int16 floor before packing; only candidates far
      // below the winner clamp, so the argmax is unaffected
      int best = INT_MIN;
      int64_t i = 0;
      for (; i + 2 <= S; i += 2) {
        const int s1 = max(prev[i] + trans[i * S + j], -32768);
        const int s2 = max(prev[i + 1] + trans[(i + 1) * S + j], -32768);
        const int p1 = (s1 << 16) | (int)i;
        const int p2 = (s2 << 16) | (int)(i + 1);
        best = imax3(best, p1, p2);
      }
      if (i < S) {
        const int s1 = max(prev[i] + trans[i * S + j], -32768);
        best = max(best, (s1 << 16) | (int)i);
      }
      curr[j] = (best >> 16) + emit[(b * T + t) * S + j];
      bp[(b * T + t) * S + j] = best & 0xFFFF;
    }
    __syncthreads();
    for (int64_t j = threadIdx.x; j < S; j += blockDim.x) prev[j] = curr[j];
    __syncthreads();
  }

  if (threadIdx.x == 0) {
    int best = INT_MIN;
    int bj = 0;
    for (int64_t j = 0; j < S; j++) {
      if (prev[j] > best) { best = prev[j]; bj = (int)j; }
    }
    score[b] = (int64_t)best + norm_acc;
    path[b * T + T - 1] = bj;
    for (int64_t t = T - 1; t > 0; t--) {
      bj = bp[(b * T + t) * S + bj];
      path[b * T + t - 1] = bj;
    }
  }
}

// Viterbi, float32 path (explicit argmax compares; same structure).
__global__ void viterbi_f32_kernel(const float* __restrict__ emit, const float* __restrict__ trans,
                                   const float* __restrict__ prior, int64_t B, int64_t T,
                                   int64_t S, int32_t* __restrict__ bp,
                                   int32_t* __restrict__ path, float* __restrict__ score,
                                   float* __restrict__ work) {
  const int64_t b = blockIdx.x;
  float* prev = work + b * 2 * S;
  float* curr = prev + S;
  for (int64_t j = threadIdx.x; j < S; j += blockDim.x) {
    prev[j] = prior[j] + emit[(b * T + 0) * S + j];
  }
  __syncthreads();
  for (int64_t t = 1; t < T; t++) {
    for (int64_t j = threadIdx.x; j < S; j += blockDim.x) {
      float best = -FLT_MAX;
      int bi = 0;
      for (int64_t i = 0; i < S; i++) {
        const float c = prev[i] + trans[i * S + j];
        if (c > best) { best = c; bi = (int)i; }
      }
      curr[j] = best + emit[(b * T + t) * S + j];
      bp[(b * T + t) * S + j] = bi;
    }
    __syncthreads();
    for (int64_t j = threadIdx.x; j < S; j += blockDim.x) prev[j] = curr[j];
    __syncthreads();
  }
  if (threadIdx.x == 0) {
    float best = -FLT_MAX;
    int bj = 0;
    for (int64_t j = 0; j < S; j++) {
      if (prev[j] > best) { best = prev[j]; bj = (int)j; }
    }
    score[b] = best;
    path[b * T + T - 1] = bj;
    for (int64_t t = T - 1; t > 0; t--) {
      bj = bp[(b * T + t) * S + bj];
      path[b * T + t - 1] = bj;
    }
  }
}

// ---------------------------------------------------------------------------
// DTW (min-plus, steps: match/insert/delete). Anti-diagonal wavefront, one
// block per sequence, full D matrix stored for the backtrace.
// ---------------------------------------------------------------------------

template <typename T_, bool kIsInt>
__global__ void dtw_kernel(const T_* __restrict__ cost /* [B,N,M] */, int64_t B, int64_t N,
                           int64_t M, T_* __restrict__ D /* [B,N,M] */,
                           int32_t* __restrict__ path /* [B,N+M,2], -1 padded */,
                           int32_t* __restrict__ path_len /* [B] */) {
  const int64_t b = blockIdx.x;
  const T_* c = cost + b * N * M;
  T_* d = D + b * N * M;
  const T_ BIG = kIsInt ? (T_)(1 << 30) : (T_)FLT_MAX;

  for (int64_t diag = 0; diag < N + M - 1; diag++) {
    const int64_t i_lo = diag < M ? 0 : diag - M + 1;
    const int64_t i_hi = diag < N ? diag : N - 1;
    for (int64_t i = i_lo + threadIdx.x; i <= i_hi; i += blockDim.x) {
      const int64_t j = diag - i;
      T_ up = (i > 0) ? d[(i - 1) * M + j] : BIG;
      T_ left = (j > 0) ? d[i * M + j - 1] : BIG;
      T_ ul = (i > 0 && j > 0) ? d[(i - 1) * M + j - 1] : ((i == 0 && j == 0) ? (T_)0 : BIG);
      T_ m;
      if (kIsInt) {
        m = (T_)imin3((int)up, (int)left, (int)ul);
      } else {
        m = fminf((float)up, fminf((float)left, (float)ul));
      }
      d[i * M + j] = c[i * M + j] + m;
    }
    __syncthreads();
  }

  if (threadIdx.x == 0) {
    int64_t i = N - 1, j = M - 1, n = 0;
    int32_t* p = path + b * (N + M) * 2;
    while (true) {
      p[n * 2] = (int32_t)i;
      p[n * 2 + 1] = (int32_t)j;
      n++;
      if (i == 0 && j == 0) break;
      const T_ up = (i > 0) ? d[(i - 1) * M + j] : BIG;
      const T_ left = (j > 0) ? d[i * M + j - 1] : BIG;
      const T_ ul = (i > 0 && j > 0) ? d[(i - 1) * M + j - 1] : BIG;
      if (ul <= up && ul <= left) { i--; j--; }
      else if (up <= left) { i--; }
      else { j--; }
    }
    path_len[b] = (int32_t)n;
    for (int64_t k = n; k < N + M; k++) { p[k * 2] = -1; p[k * 2 + 1] = -1; }
  }
}

// ---------------------------------------------------------------------------
// CTC forced alignment: Viterbi over the 2L+1 blank-interleaved trellis.
// Allowed predecessors of state s at time t: s, s-1, and s-2 when s is a
// non-blank whose label differs from the label at s-2. Backpointers store
// the step (0/1/2). One block per sequence.
// ---------------------------------------------------------------------------

template <typename T_, bool kIsInt>
__global__ void ctc_align_kernel(const T_* __restrict__ logp /* [B,T,C] */,
                                 const int64_t* __restrict__ tgt /* [B,L], -1 pad */, int64_t B,
                                 int64_t T, int64_t C, int64_t L, int64_t blank,
                                 uint8_t* __restrict__ bp /* [B,T,2L+1] */,
                                 int32_t* __restrict__ frames /* [B,T] label per frame */,
                                 T_* __restrict__ score /* [B] */,
                                 T_* __restrict__ work /* [B, 2*(2L+1)] */) {
  const int64_t b = blockIdx.x;
  int64_t Lb = 0;
  for (int64_t l = 0; l < L; l++) {
    if (tgt[b * L + l] >= 0) Lb++;
  }
  const int64_t Sn = 2 * Lb + 1;
  T_* prev = work + b * 2 * (2 * L + 1);
  T_* curr = prev + (2 * L + 1);
  const T_ NEG = kIsInt ? (T_)kNegBig : (T_)(-FLT_MAX / 2);

  auto state_label = [&](int64_t s) -> int64_t {
    return (s % 2 == 0) ? blank : tgt[b * L + s / 2];
  };

  for (int64_t s = threadIdx.x; s < Sn; s += blockDim.x) {
    prev[s] = (s <= 1) ? logp[(b * T) * C + state_label(s)] : NEG;
  }
  __syncthreads();

  for (int64_t t = 1; t < T; t++) {
    for (int64_t s = threadIdx.x; s < Sn; s += blockDim.x) {
      const T_ stay = prev[s];
      const T_ from1 = (s >= 1) ? prev[s - 1] : NEG;
      T_ from2 = NEG;
      if (s >= 2 && s % 2 == 1 && state_label(s) != state_label(s - 2)) from2 = prev[s - 2];
      T_ best;
      int choice;
      if (kIsInt) {
        best = (T_)imax3((int)stay, (int)from1, (int)from2);
      } else {
        best = fmaxf((float)stay, fmaxf((float)from1, (float)from2));
      }
      choice = (best == stay) ? 0 : (best == from1 ? 1 : 2);
      curr[s] = best + logp[(b * T + t) * C + state_label(s)];
      bp[(b * T + t) * (2 * L + 1) + s] = (uint8_t)choice;
    }
    __syncthreads();
    for (int64_t s = threadIdx.x; s < Sn; s += blockDim.x) prev[s] = curr[s];
    __syncthreads();
  }

  if (threadIdx.x == 0) {
    int64_t s = (Sn >= 2 && prev[Sn - 1] < prev[Sn - 2]) ? Sn - 2 : Sn - 1;
    score[b] = prev[s];
    for (int64_t t = T - 1; t >= 0; t--) {
      frames[b * T + t] = (int32_t)state_label(s);
      if (t > 0) s -= bp[(b * T + t) * (2 * L + 1) + s];
    }
  }
}

}  // namespace

void viterbi(torch::Tensor& path, torch::Tensor& score, torch::Tensor& bp,
             torch::Tensor const& emissions, torch::Tensor const& transitions,
             torch::Tensor const& priors) {
  TORCH_CHECK(emissions.is_cuda() && emissions.dim() == 3 && emissions.is_contiguous(),
              "emissions must be contiguous [B, T, S] CUDA");
  const int64_t B = emissions.size(0), T = emissions.size(1), S = emissions.size(2);
  TORCH_CHECK(S <= 32768, "S <= 32768 (packed argmax uses 16 bits)");
  TORCH_CHECK(transitions.sizes() == torch::IntArrayRef({S, S}) && transitions.is_contiguous(),
              "transitions must be [S, S]");
  TORCH_CHECK(priors.numel() == S && priors.is_contiguous(), "priors must be [S]");
  TORCH_CHECK(path.dtype() == torch::kInt32 && path.sizes() == torch::IntArrayRef({B, T}));
  TORCH_CHECK(bp.dtype() == torch::kInt32 && bp.sizes() == torch::IntArrayRef({B, T, S}));
  const at::cuda::CUDAGuard guard(emissions.device());
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  if (emissions.dtype() == torch::kInt32) {
    TORCH_CHECK(transitions.dtype() == torch::kInt32 && priors.dtype() == torch::kInt32,
                "dtype mismatch");
    TORCH_CHECK(score.dtype() == torch::kInt64 && score.numel() == B);
    auto work = torch::empty({B, 2 * S}, emissions.options());
    viterbi_i32_kernel<<<B, kThreads, 0, stream>>>(
        emissions.const_data_ptr<int32_t>(), transitions.const_data_ptr<int32_t>(),
        priors.const_data_ptr<int32_t>(), B, T, S, bp.data_ptr<int32_t>(),
        path.data_ptr<int32_t>(), score.data_ptr<int64_t>(), work.data_ptr<int32_t>());
  } else if (emissions.dtype() == torch::kFloat32) {
    TORCH_CHECK(transitions.dtype() == torch::kFloat32 && priors.dtype() == torch::kFloat32,
                "dtype mismatch");
    TORCH_CHECK(score.dtype() == torch::kFloat32 && score.numel() == B);
    auto work = torch::empty({B, 2 * S}, emissions.options());
    viterbi_f32_kernel<<<B, kThreads, 0, stream>>>(
        emissions.const_data_ptr<float>(), transitions.const_data_ptr<float>(),
        priors.const_data_ptr<float>(), B, T, S, bp.data_ptr<int32_t>(),
        path.data_ptr<int32_t>(), score.data_ptr<float>(), work.data_ptr<float>());
  } else {
    TORCH_CHECK(false, "emissions must be int32 (DPX path) or float32");
  }
}

void dtw(torch::Tensor& path, torch::Tensor& path_len, torch::Tensor& D,
         torch::Tensor const& cost) {
  TORCH_CHECK(cost.is_cuda() && cost.dim() == 3 && cost.is_contiguous(),
              "cost must be contiguous [B, N, M] CUDA");
  const int64_t B = cost.size(0), N = cost.size(1), M = cost.size(2);
  TORCH_CHECK(D.sizes() == cost.sizes() && D.dtype() == cost.dtype());
  TORCH_CHECK(path.dtype() == torch::kInt32 &&
              path.sizes() == torch::IntArrayRef({B, N + M, 2}));
  TORCH_CHECK(path_len.dtype() == torch::kInt32 && path_len.numel() == B);
  const at::cuda::CUDAGuard guard(cost.device());
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  if (cost.dtype() == torch::kInt32) {
    dtw_kernel<int32_t, true><<<B, kThreads, 0, stream>>>(
        cost.const_data_ptr<int32_t>(), B, N, M, D.data_ptr<int32_t>(), path.data_ptr<int32_t>(),
        path_len.data_ptr<int32_t>());
  } else if (cost.dtype() == torch::kFloat32) {
    dtw_kernel<float, false><<<B, kThreads, 0, stream>>>(
        cost.const_data_ptr<float>(), B, N, M, D.data_ptr<float>(), path.data_ptr<int32_t>(),
        path_len.data_ptr<int32_t>());
  } else {
    TORCH_CHECK(false, "cost must be int32 (DPX path) or float32");
  }
}

void ctc_forced_align(torch::Tensor& frames, torch::Tensor& score, torch::Tensor const& log_probs,
                      torch::Tensor const& targets, int64_t blank) {
  TORCH_CHECK(log_probs.is_cuda() && log_probs.dim() == 3 && log_probs.is_contiguous(),
              "log_probs must be contiguous [B, T, C] CUDA");
  const int64_t B = log_probs.size(0), T = log_probs.size(1), C = log_probs.size(2);
  TORCH_CHECK(targets.is_cuda() && targets.dtype() == torch::kInt64 && targets.dim() == 2 &&
                  targets.size(0) == B && targets.is_contiguous(),
              "targets must be contiguous int64 [B, L] (-1 padded)");
  const int64_t L = targets.size(1);
  TORCH_CHECK(frames.dtype() == torch::kInt32 && frames.sizes() == torch::IntArrayRef({B, T}));
  TORCH_CHECK(blank >= 0 && blank < C, "blank out of range");
  const at::cuda::CUDAGuard guard(log_probs.device());
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  auto bp = torch::empty({B, T, 2 * L + 1}, log_probs.options().dtype(torch::kUInt8));
  if (log_probs.dtype() == torch::kInt32) {
    TORCH_CHECK(score.dtype() == torch::kInt32 && score.numel() == B);
    auto work = torch::empty({B, 2 * (2 * L + 1)}, log_probs.options());
    ctc_align_kernel<int32_t, true><<<B, kThreads, 0, stream>>>(
        log_probs.const_data_ptr<int32_t>(), targets.const_data_ptr<int64_t>(), B, T, C, L, blank,
        bp.data_ptr<uint8_t>(), frames.data_ptr<int32_t>(), score.data_ptr<int32_t>(),
        work.data_ptr<int32_t>());
  } else if (log_probs.dtype() == torch::kFloat32) {
    TORCH_CHECK(score.dtype() == torch::kFloat32 && score.numel() == B);
    auto work = torch::empty({B, 2 * (2 * L + 1)}, log_probs.options());
    ctc_align_kernel<float, false><<<B, kThreads, 0, stream>>>(
        log_probs.const_data_ptr<float>(), targets.const_data_ptr<int64_t>(), B, T, C, L, blank,
        bp.data_ptr<uint8_t>(), frames.data_ptr<int32_t>(), score.data_ptr<float>(),
        work.data_ptr<float>());
  } else {
    TORCH_CHECK(false, "log_probs must be int32 (DPX path) or float32");
  }
}