// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "contrib_ops/cpu/quantization/matmul_nbits_impl.h"

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <cstring>
#include <type_traits>

#include "core/common/common.h"
#include "core/common/float16.h"
#include "core/providers/common.h"
#include "core/platform/threadpool.h"

namespace onnxruntime {
namespace contrib {

template <class T, class zeroT>
void Dequantize4BitsKernelReOrder(
    T* output, const uint8_t* quant_data, const T* scale_data,
    const zeroT* zero_points, const int32_t* reorder_idx, int block_size,
    int groups_per_threadblock, int total_groups, int out_rows, int out_cols,
    int blockIdx_x, int threadIdx_x) {
  const int group_id = blockIdx_x * groups_per_threadblock + ((threadIdx_x * 8) / block_size);
  if (group_id >= total_groups) {
    return;
  }
  const int scales_shape_x = (out_cols + block_size - 1) / block_size;
  const int zero_point_shape_x = (scales_shape_x + 1) / 2;

  int n_idx = group_id / scales_shape_x;
  int kb_idx = group_id % scales_shape_x;
  int element_offset = group_id * block_size + ((threadIdx_x * 8) & (block_size - 1));

  const int out_x = element_offset % (scales_shape_x * block_size);
  const int out_y = element_offset / (scales_shape_x * block_size);
  if (out_y >= out_rows || out_x >= out_cols) {
    return;
  }
  T* output_i = output + out_y * out_cols + out_x;
  uint32_t quant_value = *(reinterpret_cast<const uint32_t*>(quant_data + element_offset / 2));
  if constexpr (onnxruntime::endian::native == onnxruntime::endian::big) {
    const uint8_t* c = reinterpret_cast<const uint8_t*>(&quant_value);
    quant_value = static_cast<uint32_t>(c[0]) |
                  static_cast<uint32_t>(c[1]) << 8 |
                  static_cast<uint32_t>(c[2]) << 16 |
                  static_cast<uint32_t>(c[3]) << 24;
  }
  const int remain_x = std::min(8, out_cols - out_x);
  const int32_t* reorder_idx_with_off = reorder_idx + kb_idx * block_size + ((threadIdx_x * 8) & (block_size - 1));
  for (int i = 0; i < remain_x; i++) {
    int32_t rid = reorder_idx ? reorder_idx_with_off[i] : kb_idx;
    T scale = *(scale_data + n_idx * scales_shape_x + rid);
    float zp_f = 8;
    if (zero_points) {
      if constexpr (std::is_same_v<zeroT, uint8_t>) {
        uint8_t zp = 8;
        zp = zero_points[n_idx * zero_point_shape_x + rid / 2];
        zp = (rid & 0x01) ? (zp >> 4) : (zp & 0x0f);
      } else {
        zp_f = *(zero_points + static_cast<uint64_t>(n_idx) * static_cast<uint64_t>(scales_shape_x) + static_cast<uint64_t>(rid));
      }
    }

    if constexpr (std::is_same_v<T, MLFloat16>) {
      T zp_adjust = -scale * MLFloat16(zp_f);
      output_i[i] = static_cast<float>((quant_value >> (4 * i)) & 0xF) * scale + zp_adjust;
    } else {
      T zp_adjust = -scale * zp_f;
      output_i[i] = T((quant_value >> (4 * i)) & 0xF) * scale + zp_adjust;
    }
  }
}

template <typename inputT, typename zeroT, int qbits>
void DequantizeBlockwise(
    inputT* output,              // dequantized output
    const uint8_t* quant_data,   // quantized input
    const inputT* scales_data,   // quantization scales
    const zeroT* zero_points,    // quantization zero points
    const int32_t* reorder_idx,  // reorder_idx for groupwise quantization
    int32_t block_size,          // quantization block size
    bool,                        // columnwise quantization or row-wise
    int32_t K,                   // number of rows in quantized input
    int32_t N,                   // number of columns in quantized input
    onnxruntime::concurrency::ThreadPool* pool) {
  auto ceildiv = [](int a, int b) { return (a + b - 1) / b; };
  constexpr int element_per_thread = 8;
  int groups_per_threadblock = 256 * element_per_thread / block_size;
  int groups_per_K = ceildiv(K, block_size);
  int total_groups = N * groups_per_K;  // total elemenets in quant_data
  int blocks_per_grid = static_cast<int>(ceildiv(total_groups, groups_per_threadblock));
  concurrency::ThreadPool::TrySimpleParallelFor(
      pool, static_cast<std::ptrdiff_t>(blocks_per_grid),
      [&](std::ptrdiff_t block_id) {
        for (int j = 0; j < 256; j++) {
          Dequantize4BitsKernelReOrder(output, quant_data, scales_data, zero_points,
                                       reorder_idx, block_size, groups_per_threadblock,
                                       total_groups, N, K, static_cast<int>(block_id), j);
        }
      });
}

template void DequantizeBlockwise<float, uint8_t, 4>(
    float* output, const uint8_t* quant_data, const float* scales_data,
    const uint8_t* zero_points, const int32_t* reorder_idx, int32_t block_size,
    bool columnwise, int32_t K, int32_t N, onnxruntime::concurrency::ThreadPool* thread_pool);

template void DequantizeBlockwise<float, float, 4>(
    float* output, const uint8_t* quant_data, const float* scales_data,
    const float* zero_points, const int32_t* reorder_idx, int32_t block_size,
    bool columnwise, int32_t K, int32_t N, onnxruntime::concurrency::ThreadPool* thread_pool);

template void DequantizeBlockwise<float, MLFloat16, 4>(
    float* output, const uint8_t* quant_data, const float* scales_data,
    const MLFloat16* zero_points, const int32_t* reorder_idx, int32_t block_size,
    bool columnwise, int32_t K, int32_t N, onnxruntime::concurrency::ThreadPool* thread_pool);

// 2-bit dequantization kernel for float/MLFloat16 zero points.
// Processes 16 elements at a time (16 x 2-bit = 32 bits = one uint32_t).
// Layout: columnwise packing — elements within a column are packed consecutively,
//         output[n * K + k] = (quant_value - zp) * scale
template <class T, class zeroT>
void Dequantize2BitsKernel(
    T* output, const uint8_t* quant_data, const T* scale_data,
    const zeroT* zero_points, int block_size,
    int groups_per_threadblock, int total_groups, int N, int K,
    int blockIdx_x, int threadIdx_x) {
  // Each "thread" handles 16 elements (one uint32 of packed 2-bit values)
  constexpr int elements_per_thread = 16;
  const int group_id = blockIdx_x * groups_per_threadblock + ((threadIdx_x * elements_per_thread) / block_size);
  if (group_id >= total_groups) {
    return;
  }
  const int k_blocks = (K + block_size - 1) / block_size;

  int n_idx = group_id / k_blocks;
  int kb_idx = group_id % k_blocks;
  int element_offset = group_id * block_size + ((threadIdx_x * elements_per_thread) & (block_size - 1));

  const int k_offset = element_offset % (k_blocks * block_size);
  const int n_offset = element_offset / (k_blocks * block_size);
  if (n_offset >= N || k_offset >= K) {
    return;
  }

  T* output_i = output + n_offset * K + k_offset;
  // 16 elements × 2 bits = 4 bytes. Use memcpy to avoid alignment UB.
  uint32_t quant_value = 0;
  std::memcpy(&quant_value, quant_data + element_offset / 4, sizeof(uint32_t));
  if constexpr (onnxruntime::endian::native == onnxruntime::endian::big) {
    const uint8_t* c = reinterpret_cast<const uint8_t*>(&quant_value);
    quant_value = static_cast<uint32_t>(c[0]) |
                  static_cast<uint32_t>(c[1]) << 8 |
                  static_cast<uint32_t>(c[2]) << 16 |
                  static_cast<uint32_t>(c[3]) << 24;
  }
  const int remain_k = std::min(elements_per_thread, K - k_offset);

  float scale_f = static_cast<float>(*(scale_data + static_cast<uint64_t>(n_idx) * static_cast<uint64_t>(k_blocks) + static_cast<uint64_t>(kb_idx)));
  float zp_f = 0.0f;
  if (zero_points) {
    if constexpr (std::is_same_v<zeroT, MLFloat16>) {
      zp_f = (*(zero_points + static_cast<uint64_t>(n_idx) * static_cast<uint64_t>(k_blocks) + static_cast<uint64_t>(kb_idx))).ToFloat();
    } else {
      zp_f = static_cast<float>(*(zero_points + static_cast<uint64_t>(n_idx) * static_cast<uint64_t>(k_blocks) + static_cast<uint64_t>(kb_idx)));
    }
  }

  float zp_adjust = -scale_f * zp_f;
  for (int i = 0; i < remain_k; i++) {
    float q = static_cast<float>((quant_value >> (2 * i)) & 0x3);
    output_i[i] = static_cast<T>(q * scale_f + zp_adjust);
  }
}

template <class T, class zeroT>
void Dequantize2BitsFallback(
    T* output, const uint8_t* quant_data, const T* scale_data,
    const zeroT* zero_points, int block_size, int N, int K) {
  const int k_blocks = (K + block_size - 1) / block_size;

  for (int n = 0; n < N; ++n) {
    for (int kb = 0; kb < k_blocks; ++kb) {
      const int group_offset = (n * k_blocks + kb) * block_size;
      const int k_start = kb * block_size;
      const int k_count = std::min(block_size, K - k_start);

      const float scale = static_cast<float>(scale_data[static_cast<uint64_t>(n) * static_cast<uint64_t>(k_blocks) + static_cast<uint64_t>(kb)]);
      float zp_f = 0.0f;
      if (zero_points) {
        if constexpr (std::is_same_v<zeroT, MLFloat16>) {
          zp_f = zero_points[static_cast<uint64_t>(n) * static_cast<uint64_t>(k_blocks) + static_cast<uint64_t>(kb)].ToFloat();
        } else {
          zp_f = static_cast<float>(zero_points[static_cast<uint64_t>(n) * static_cast<uint64_t>(k_blocks) + static_cast<uint64_t>(kb)]);
        }
      }
      const float zp_adjust = -scale * zp_f;
      T* output_i = output + static_cast<uint64_t>(n) * static_cast<uint64_t>(K) + static_cast<uint64_t>(k_start);

      for (int i = 0; i < k_count; ++i) {
        const int element_offset = group_offset + i;
        const uint8_t packed = quant_data[element_offset / 4];
        const uint8_t q = (packed >> (2 * (element_offset & 0x3))) & 0x3;
        output_i[i] = static_cast<T>(static_cast<float>(q) * scale + zp_adjust);
      }
    }
  }
}

// Specialization of DequantizeBlockwise for qbits=2
template <typename inputT, typename zeroT>
void DequantizeBlockwise2Bits(
    inputT* output,
    const uint8_t* quant_data,
    const inputT* scales_data,
    const zeroT* zero_points,
    int32_t block_size,
    bool columnwise,
    int32_t K,
    int32_t N,
    onnxruntime::concurrency::ThreadPool* pool) {
  auto ceildiv = [](int a, int b) { return (a + b - 1) / b; };
  constexpr int elements_per_thread = 16;
  ORT_ENFORCE(columnwise, "Row-wise quantization is not supported");
  ORT_ENFORCE(block_size > 0, "block_size must be positive, got: ", block_size);
  ORT_ENFORCE((block_size & (block_size - 1)) == 0, "block_size must be a power of two, got: ", block_size);
  if (block_size > 256 * elements_per_thread || block_size % elements_per_thread != 0) {
    Dequantize2BitsFallback(output, quant_data, scales_data, zero_points, block_size, N, K);
    return;
  }

  int groups_per_threadblock = 256 * elements_per_thread / block_size;
  int groups_per_K = ceildiv(K, block_size);
  int total_groups = N * groups_per_K;
  int blocks_per_grid = static_cast<int>(ceildiv(total_groups, groups_per_threadblock));
  concurrency::ThreadPool::TrySimpleParallelFor(
      pool, static_cast<std::ptrdiff_t>(blocks_per_grid),
      [&](std::ptrdiff_t block_id) {
        for (int j = 0; j < 256; j++) {
          Dequantize2BitsKernel(output, quant_data, scales_data, zero_points,
                                block_size, groups_per_threadblock,
                                total_groups, N, K, static_cast<int>(block_id), j);
        }
      });
}

// Explicit instantiations for 2-bit dequantization
template void DequantizeBlockwise2Bits<float, float>(
    float* output, const uint8_t* quant_data, const float* scales_data,
    const float* zero_points, int32_t block_size,
    bool columnwise, int32_t K, int32_t N, onnxruntime::concurrency::ThreadPool* thread_pool);

template void DequantizeBlockwise2Bits<float, MLFloat16>(
    float* output, const uint8_t* quant_data, const float* scales_data,
    const MLFloat16* zero_points, int32_t block_size,
    bool columnwise, int32_t K, int32_t N, onnxruntime::concurrency::ThreadPool* thread_pool);

}  // namespace contrib
}  // namespace onnxruntime
