// LP probability dispatch kernel: collapse the ~7 torch ops in
// _topk_ids_logical_to_physical_probability into one launch.
//
// Python equivalent:
//
//     topk_probs = log2phy_prob[topk_ids]                           # gather
//     row_sums   = topk_probs.sum(dim=-1)                            # reduce
//     fallback   = (log2phy_map[topk_ids] >= 0).float()              # gather + cast
//     topk_probs = where(row_sums > 0, topk_probs, fallback)         # cmp + select
//     chosen     = multinomial(topk_probs, 1).flatten()              # sample
//     out        = log2phy_map[topk_ids, chosen]                     # gather
//
// Each thread handles one (token, slot) in the flattened topk_ids.
// `random_vals` is pre-generated by the caller via torch.rand (one kernel
// launch, ~5 µs); we sample the multinomial via prefix-sum + comparison
// against `random_vals[i] * row_sum`.
//
// Templated on (MAX_COPIES, BLOCK_DIM). MAX_COPIES is small (typically 2-3)
// so the per-row prefix sum unrolls into a few instructions per thread.

#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>

#include <sgl_kernel/utils.cuh>

#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>

#include <cstdint>

namespace {

template <int MAX_COPIES, int BLOCK_DIM>
__global__ void dispatch_probability_kernel(
    int32_t* __restrict__ out_topk_ids,       // (N,)
    const int32_t* __restrict__ in_topk_ids,  // (N,)
    const float* __restrict__ log2phy_prob,   // (NUM_LOGICAL, MAX_COPIES)
    const int32_t* __restrict__ log2phy_map,  // (NUM_LOGICAL, MAX_COPIES)
    const float* __restrict__ random_vals,    // (N,)
    int N) {
  const int idx = blockIdx.x * BLOCK_DIM + threadIdx.x;
  if (idx >= N) return;

  const int32_t logical_id = in_topk_ids[idx];
  const int32_t* row_map = log2phy_map + logical_id * MAX_COPIES;
  const float* row_prob = log2phy_prob + logical_id * MAX_COPIES;

  float probs[MAX_COPIES];
  int32_t maps[MAX_COPIES];
  float row_sum = 0.f;

#pragma unroll
  for (int c = 0; c < MAX_COPIES; c++) {
    maps[c] = row_map[c];
    probs[c] = row_prob[c];
    row_sum += probs[c];
  }

  // Fallback: if all LP probs for this row are 0, sample uniformly from the
  // valid physical copies (map != -1).
  if (row_sum <= 0.f) {
    row_sum = 0.f;
#pragma unroll
    for (int c = 0; c < MAX_COPIES; c++) {
      probs[c] = (maps[c] >= 0) ? 1.0f : 0.0f;
      row_sum += probs[c];
    }
  }

  // Multinomial sample: smallest c such that cumsum[0..c] > u * row_sum.
  // Implemented branch-free for unroll friendliness: chosen accumulates the
  // largest index where cumsum[..c] is still <= u, then we add 1 (clamped to
  // MAX_COPIES-1 for the all-cumsum-<=-u edge case from float rounding).
  const float u = random_vals[idx] * row_sum;
  float cum = 0.f;
  int chosen = 0;
#pragma unroll
  for (int c = 0; c < MAX_COPIES; c++) {
    cum += probs[c];
    if (u >= cum) chosen = c + 1;
  }
  if (chosen >= MAX_COPIES) chosen = MAX_COPIES - 1;

  out_topk_ids[idx] = maps[chosen];
}

template <int MAX_COPIES, int BLOCK_DIM>
void dispatch_probability(
    tvm::ffi::TensorView out_topk_ids,   // (N,) int32
    tvm::ffi::TensorView in_topk_ids,    // (N,) int32
    tvm::ffi::TensorView log2phy_prob,   // (NUM_LOGICAL, MAX_COPIES) float32
    tvm::ffi::TensorView log2phy_map,    // (NUM_LOGICAL, MAX_COPIES) int32
    tvm::ffi::TensorView random_vals) {  // (N,) float32
  using namespace host;

  SymbolicSize N{"num_topk_entries"};
  SymbolicSize NUM_LOGICAL{"num_logical"};
  SymbolicDevice device_;

  TensorMatcher({N}).with_dtype<int32_t>().with_device<kDLCUDA>(device_).verify(out_topk_ids).verify(in_topk_ids);
  TensorMatcher({NUM_LOGICAL, MAX_COPIES}).with_dtype<float>().with_device<kDLCUDA>(device_).verify(log2phy_prob);
  TensorMatcher({NUM_LOGICAL, MAX_COPIES}).with_dtype<int32_t>().with_device<kDLCUDA>(device_).verify(log2phy_map);
  TensorMatcher({N}).with_dtype<float>().with_device<kDLCUDA>(device_).verify(random_vals);

  const int n = static_cast<int>(N.unwrap());
  const int grid = (n + BLOCK_DIM - 1) / BLOCK_DIM;
  const DLDevice device = device_.unwrap();

  using KernelT = void (*)(int32_t*, const int32_t*, const float*, const int32_t*, const float*, int);
  KernelT kernel = dispatch_probability_kernel<MAX_COPIES, BLOCK_DIM>;

  LaunchKernel(grid, BLOCK_DIM, device)(
      kernel,
      static_cast<int32_t*>(out_topk_ids.data_ptr()),
      static_cast<const int32_t*>(in_topk_ids.data_ptr()),
      static_cast<const float*>(log2phy_prob.data_ptr()),
      static_cast<const int32_t*>(log2phy_map.data_ptr()),
      static_cast<const float*>(random_vals.data_ptr()),
      n);
}

}  // namespace