/*
 * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>

#include <cub/cub.cuh>

namespace moe::dev::routing {

namespace topk {

////////////////////////////////////////////////////////////////////////////////////////////////////
namespace cg = cooperative_groups;

////////////////////////////////////////////////////////////////////////////////////////////////////

static constexpr int WarpSize = 32;
static constexpr int MaxNumExpertsUnit = 128;
static constexpr int MaxSupportedTopExperts = 32;

////////////////////////////////////////////////////////////////////////////////////////////////////

template <typename TypeExpW_>
struct TopKRedType {
  using TypeExpW = TypeExpW_;
  static_assert(std::is_same_v<TypeExpW, float> || std::is_same_v<TypeExpW, half> ||
                    std::is_same_v<TypeExpW, __nv_bfloat16>,
                "Top K reduction only implemented for float, float16 and bfloat16");

  using TypeCmp = std::conditional_t<sizeof(TypeExpW) == 4, uint64_t, uint32_t>;
  using IdxT = std::conditional_t<sizeof(TypeExpW) == 4, int32_t, int16_t>;
  static constexpr int moveBits = (sizeof(TypeExpW) == 4) ? 32 : 16;
  static constexpr int maxIdx = 65535;

  TypeCmp compVal;

  static __host__ __device__ inline TypeCmp makeCmpVal(TypeExpW val, int32_t idx = 0) {
    auto valueBits = cub::Traits<TypeExpW>::TwiddleIn(
        reinterpret_cast<typename cub::Traits<TypeExpW>::UnsignedBits&>(val));
    TypeCmp compactTmp = valueBits;
    compactTmp = (compactTmp << moveBits) | (0xFFFF & (maxIdx - idx));
    // Use 65535 minus idx to give higher priority to elements with smaller indices.
    return compactTmp;
  }

  static __host__ __device__ inline void unpack(TypeExpW& value, int32_t& index, TypeCmp cmp) {
    // Since idx is always smaller than 65536 and positive, we can directly use it as the lower 16
    // bits
    index = maxIdx - static_cast<int32_t>(cmp & 0xFFFF);

    auto compactTmp = cmp >> moveBits;
    auto valueBits = cub::Traits<TypeExpW>::TwiddleOut(
        reinterpret_cast<typename cub::Traits<TypeExpW>::UnsignedBits&>(compactTmp));
    value = reinterpret_cast<TypeExpW&>(valueBits);
  }

  __host__ __device__ TopKRedType() = default;

  __host__ __device__ TopKRedType(TypeExpW val, int32_t idx) : compVal(makeCmpVal(val, idx)) {}

  __host__ __device__ operator TypeCmp() const noexcept { return compVal; }

  __device__ inline TypeCmp reduce(cg::thread_block_tile<WarpSize> const& warp) {
    // 32-bit keys (bf16/fp16) reduce with redux.sync.max.u32 directly. 64-bit keys
    // (float) split into hi=value / lo=idx-complement and take two
    // redux.sync.max.u32 (masking non-winning lanes' lo to 0), re-packed to the
    // lexicographic 64-bit max -- bit-identical to cg::reduce, used as the fallback
    // when fast redux is unavailable.
#ifdef __CUDA_ARCH__
    static constexpr bool hasFastRedux = (__CUDA_ARCH__ / 100) >= 10;
#else
    static constexpr bool hasFastRedux = false;
#endif
    if constexpr (!hasFastRedux) {
      return cg::reduce(warp, compVal, cg::greater<TypeCmp>{});
    } else if constexpr (sizeof(TypeCmp) == 8) {
      uint32_t hi = static_cast<uint32_t>(compVal >> 32);
      uint32_t lo = static_cast<uint32_t>(compVal & 0xffffffffu);
      uint32_t maxHi;
      asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" : "=r"(maxHi) : "r"(hi));
      uint32_t loContrib = (hi == maxHi) ? lo : 0u;
      uint32_t maxLo;
      asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" : "=r"(maxLo) : "r"(loContrib));
      return (static_cast<TypeCmp>(maxHi) << 32) | static_cast<TypeCmp>(maxLo);
    } else {
      TypeCmp result;
      asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" : "=r"(result) : "r"(compVal));
      return result;
    }
  }
};

////////////////////////////////////////////////////////////////////////////////////////////////////

#define TOPK_SWAP(I, J)                                   \
  {                                                       \
    auto pairMin = min(topK[I].compVal, topK[J].compVal); \
    auto pairMax = max(topK[I].compVal, topK[J].compVal); \
    topK[I].compVal = pairMax;                            \
    topK[J].compVal = pairMin;                            \
  }

// Helper to check if N is a power of 2
template <int N>
struct IsPowerOf2 {
  static constexpr bool value = (N > 0) && ((N & (N - 1)) == 0);
};

// Helper to compute the next power of 2 (>= N).
template <int N>
struct NextPow2 {
 private:
  static constexpr unsigned u = static_cast<unsigned>(N - 1);
  static constexpr unsigned s1 = u | (u >> 1);
  static constexpr unsigned s2 = s1 | (s1 >> 2);
  static constexpr unsigned s3 = s2 | (s2 >> 4);
  static constexpr unsigned s4 = s3 | (s3 >> 8);
  static constexpr unsigned s5 = s4 | (s4 >> 16);

 public:
  static constexpr int value = (N <= 1) ? 1 : static_cast<int>(s5 + 1);
};

// Batcher's odd-even mergesort as compile-time template recursion (straight-line
// SASS). Built on the span P = NextPow2<Size>; any comparator touching the padded
// tail [Size, P) is dropped at compile time, so every N (incl. non-pow-2) sorts
// the real elements with no runtime padding or storage.
template <int A, int B, int Size, typename T>
__device__ __forceinline__ void topkCompareSwap(T* a) {
  if constexpr (A < Size && B < Size) {
    if (a[A] < a[B]) {
      T tmp = a[A];
      a[A] = a[B];
      a[B] = tmp;
    }
  } else {
    (void)a;
  }
}

// Inner pair-merge loop unrolled as template recursion.
template <int I, int End, int Step, int PairStride, int Size, typename T>
__device__ __forceinline__ void topkMergePairs(T* a) {
  if constexpr (I + Step < End) {
    topkCompareSwap<I, I + Step, Size, T>(a);
    topkMergePairs<I + PairStride, End, Step, PairStride, Size, T>(a);
  } else {
    (void)a;
  }
}

// Batcher's odd-even merge of two sorted halves spanning [Lo, Lo+N) at stride R.
template <int Lo, int N, int R, int Size, typename T>
__device__ __forceinline__ void topkOEM(T* a) {
  constexpr int M = R * 2;
  if constexpr (M < N) {
    topkOEM<Lo, N, M, Size, T>(a);
    topkOEM<Lo + R, N - R, M, Size, T>(a);
    topkMergePairs<Lo + R, Lo + N, R, M, Size, T>(a);
  } else if constexpr (R < N) {
    topkCompareSwap<Lo, Lo + R, Size, T>(a);
  } else {
    (void)a;
  }
}

// Recursive Batcher's odd-even mergesort over the span [Lo, Lo+N). Size is the
// number of real (non-sentinel) elements used to drop padded comparators.
template <int Lo, int N, int Size, typename T>
__device__ __forceinline__ void topkSortBatcher(T* a) {
  if constexpr (N > 1) {
    constexpr int Half = N / 2;
    topkSortBatcher<Lo, Half, Size, T>(a);
    topkSortBatcher<Lo + Half, N - Half, Size, T>(a);
    topkOEM<Lo, N, 1, Size, T>(a);
  } else {
    (void)a;
  }
}

////////////////////////////////////////////////////////////////////////////////////////////////////
template <int N, typename RedType>
struct Sort {
  static_assert(N > 0 && N <= 64, "Sort only supports N in range [1, 64]");

  static __device__ void run(RedType* topK) {
    if constexpr (IsPowerOf2<N>::value) {
// Bitonic sort for power-of-2 sizes - more efficient
#pragma unroll
      for (int k = 2; k <= N; k *= 2) {
#pragma unroll
        for (int j = k / 2; j > 0; j /= 2) {
#pragma unroll
          for (int i = 0; i < N; ++i) {
            int ixj = i ^ j;
            if (ixj > i) {
              if ((i & k) == 0) {
                if (topK[i].compVal < topK[ixj].compVal) {
                  auto tmp = topK[i].compVal;
                  topK[i].compVal = topK[ixj].compVal;
                  topK[ixj].compVal = tmp;
                }
              } else {
                if (topK[i].compVal > topK[ixj].compVal) {
                  auto tmp = topK[i].compVal;
                  topK[i].compVal = topK[ixj].compVal;
                  topK[ixj].compVal = tmp;
                }
              }
            }
          }
        }
      }
    } else {
      // Non-power-of-2: Batcher OEM over NextPow2<N>, preserving the same
      // descending order as the prior odd-even transposition sort.
      constexpr int P = NextPow2<N>::value;
      topkSortBatcher<0, P, N, RedType>(topK);
    }
  }
};

template <typename RedType>
struct Sort<1, RedType> {
  static __device__ void run(RedType* topK) {}
};

template <typename RedType>
struct Sort<2, RedType> {
  static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); }
};

template <typename RedType>
struct Sort<3, RedType> {
  static __device__ void run(RedType* topK) {
    TOPK_SWAP(0, 1);
    TOPK_SWAP(1, 2);
    TOPK_SWAP(0, 1);
  }
};

template <typename RedType>
struct Sort<4, RedType> {
  static __device__ void run(RedType* topK) {
    TOPK_SWAP(0, 2);
    TOPK_SWAP(1, 3);
    TOPK_SWAP(0, 1);
    TOPK_SWAP(2, 3);
    TOPK_SWAP(1, 2);
  }
};

////////////////////////////////////////////////////////////////////////////////////////////////////

template <int K, typename Type>
__forceinline__ __device__ void reduceTopK(cg::thread_block_tile<WarpSize> const& warp,
                                           Type (&out)[K], int32_t (&outIdx)[K], Type value,
                                           int32_t idx, Type const minValue, int actualK = K) {
  static_assert(K > 0, "Top K must have K > 0");
  static_assert(K < WarpSize, "Top K must have K < WarpSize");
  using RedType = TopKRedType<Type>;
  RedType topK{value, idx};
  typename RedType::TypeCmp packedMax{};
#pragma unroll
  for (int kk = 0; kk < actualK; ++kk)  //@todo: check if actualK is correct
  {
    topK = kk > 0 && packedMax == topK.compVal ? RedType{minValue, idx} : topK;
    // get the next largest value
    packedMax = topK.reduce(warp);
    RedType::unpack(out[kk], outIdx[kk], packedMax);
  }
};

template <int K, typename Type, int N>
__forceinline__ __device__ void reduceTopK(cg::thread_block_tile<WarpSize> const& warp,
                                           Type (&out)[K], int32_t (&outIdx)[K], Type (&value)[N],
                                           int32_t (&idx)[N], Type const minValue,
                                           int actualK = K) {
  static_assert(K > 0, "Top K must have K > 0");
  static_assert(K <= WarpSize, "Top K must have K <= WarpSize");
  static_assert(N > 0, "Top K must have N > 0");
  static_assert(N <= 64, "Only support candidates number less than or equal to 64*32=2048");
  using RedType = TopKRedType<Type>;
  RedType topK[N];
#pragma unroll
  for (int nn = 0; nn < N; ++nn) {
    topK[nn] = RedType{value[nn], idx[nn]};
  }

  Sort<N, RedType>::run(topK);

  typename RedType::TypeCmp packedMax{};
  for (int kk = 0; kk < actualK; ++kk) {
    bool update = kk > 0 && packedMax == topK[0].compVal;
#pragma unroll
    for (int nn = 0; nn < N; ++nn) {
      topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]}
                 : update              ? topK[nn + 1]
                                       : topK[nn];
    }
    // get the next largest value
    packedMax = topK[0].reduce(warp);
    RedType::unpack(out[kk], outIdx[kk], packedMax);
  }
};

/// Capability envelope for the high-expert lane-owned routing specialization.
/// This predicate is shared by compile-time tier selection and runtime dispatch;
/// dispatch tables still decide which concrete tiers are instantiated.
static constexpr int HighExpertLaneOwnedTopKMinExperts = 512;
static constexpr int HighExpertLaneOwnedTopKMaxExperts = 1024;
static constexpr int HighExpertLaneOwnedTopKMinTopExperts = 9;
static constexpr int HighExpertLaneOwnedTopKMaxTopExperts = 16;

__host__ __device__ constexpr bool isInHighExpertLaneOwnedTopKRange(int numExperts,
                                                                    int numTopExperts) {
  return numExperts >= HighExpertLaneOwnedTopKMinExperts &&
         numExperts <= HighExpertLaneOwnedTopKMaxExperts &&
         numTopExperts >= HighExpertLaneOwnedTopKMinTopExperts &&
         numTopExperts <= HighExpertLaneOwnedTopKMaxTopExperts;
}

// Exact-K variant that returns only the result owned by this lane. The generic
// reduceTopK interface materializes K scores and K indices per thread even though
// routing kernels ultimately consume only element laneIdx. For high-K tiers that
// dynamic lane-indexed array is placed in local memory and can spill the sorted
// candidates as well. Keeping a single packed result per lane lets nvcc scalarize
// the output while preserving the same comparison and tie-breaking order.
template <int K, typename Type, int N>
__forceinline__ __device__ void reduceTopKForLane(cg::thread_block_tile<WarpSize> const& warp,
                                                  Type& out, int32_t& outIdx, Type (&value)[N],
                                                  int32_t (&idx)[N], Type const minValue,
                                                  int32_t laneIdx) {
  static_assert(K > 0, "Top K must have K > 0");
  static_assert(K <= WarpSize, "Top K must have K <= WarpSize");
  static_assert(N > 0, "Top K must have N > 0");
  static_assert(N <= 64, "Only support candidates number less than or equal to 64*32=2048");
  using RedType = TopKRedType<Type>;
  RedType topK[N];
#pragma unroll
  for (int nn = 0; nn < N; ++nn) {
    topK[nn] = RedType{value[nn], idx[nn]};
  }

  Sort<N, RedType>::run(topK);

  typename RedType::TypeCmp packedMax{};
  typename RedType::TypeCmp lanePacked{};
#pragma unroll
  for (int kk = 0; kk < K; ++kk) {
    bool update = kk > 0 && packedMax == topK[0].compVal;
#pragma unroll
    for (int nn = 0; nn < N; ++nn) {
      topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]}
                 : update              ? topK[nn + 1]
                                       : topK[nn];
    }
    packedMax = topK[0].reduce(warp);
    if (laneIdx == kk) {
      lanePacked = packedMax;
    }
  }

  if (laneIdx < K) {
    RedType::unpack(out, outIdx, lanePacked);
  } else {
    out = minValue;
    outIdx = -1;
  }
}
#undef TOPK_SWAP
}  // namespace topk
}  // namespace moe::dev::routing
