// Copyright (c) 2025 PaddlePaddle Authors. 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 <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <utility>
#include <vector>
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/core/enforce.h"
#include "paddle/phi/core/sparse_coo_tensor.h"
#include "paddle/phi/kernels/empty_kernel.h"
#include "paddle/phi/kernels/funcs/math_function_impl.h"
#include "paddle/phi/kernels/funcs/transpose_function.cuh"

/** Reserved value for indicating "empty". */
#define EMPTY_CELL (0)
/** CUDA naive thread block size. */
#define BLOCK_SIZE (256)

__inline__ __device__ int8_t atomicCAS(int8_t* address,
                                       int8_t compare,
                                       int8_t val) {
  auto address_value = reinterpret_cast<uintptr_t>(address);
  int32_t* base_address = reinterpret_cast<int32_t*>(
      reinterpret_cast<char*>(address) - (address_value & 3));
  int32_t int_val = static_cast<int32_t>(val) << ((address_value & 3) * 8);
  int32_t int_comp = static_cast<int32_t>(compare) << ((address_value & 3) * 8);
  return static_cast<int8_t>(atomicCAS(base_address, int_comp, int_val));
}

// TODO(ShigureNyako): can we do this more efficiently?
__inline__ __device__ int16_t atomicCAS(int16_t* address,
                                        int16_t compare,
                                        int16_t val) {
  auto address_value = reinterpret_cast<uintptr_t>(address);
  int32_t* base_address = reinterpret_cast<int32_t*>(
      reinterpret_cast<char*>(address) - (address_value & 2));
  int32_t int_val = static_cast<int32_t>(val) << ((address_value & 2) * 8);
  int32_t int_comp = static_cast<int32_t>(compare) << ((address_value & 2) * 8);
  return static_cast<int16_t>(atomicCAS(base_address, int_comp, int_val));
}

__inline__ __device__ int64_t atomicCAS(int64_t* address,
                                        int64_t compare,
                                        int64_t val) {
  using AtomicCAS64Type = unsigned long long;  // NOLINT(runtime/int)
  return static_cast<int64_t>(
      atomicCAS(reinterpret_cast<AtomicCAS64Type*>(address),
                static_cast<AtomicCAS64Type>(compare),
                static_cast<AtomicCAS64Type>(val)));
}

namespace phi {
namespace sparse {

template <typename dtype = int>
__device__ uint64_t hash_func_64b(dtype* data, int n = 4) {
  uint64_t hash = UINT64_C(14695981039346656037);
  for (int j = 0; j < n; j++) {
    hash ^= static_cast<unsigned int>(data[j]);
    hash *= UINT64_C(1099511628211);
  }
  // hash = (hash >> 60) ^ (hash & 0xFFFFFFFFFFFFFFF);
  return hash;
}

template <typename key_type>
__device__ int hash(key_type key, int _capacity) {
  return static_cast<uint64_t>(key) % _capacity;
}

template <typename key_type, typename val_type>
class GPUHashTable {
 private:
  // public:
  bool free_pointers;
  const int _capacity;
  const int _divisor;
  const int _width;
  key_type* table_keys;
  val_type* table_vals;
  void insert_many_coords(const GPUContext& dev_ctx,
                          const int* coords,
                          const int n);
  void lookup_many_coords(const GPUContext& dev_ctx,
                          const int* coords,
                          val_type* results,
                          const int* kernel_sizes,
                          const int* tensor_strides,
                          const int n,
                          const int kernel_volume);

 public:
  GPUHashTable(DenseTensor* table_keys,
               DenseTensor* table_vals,
               const int divisor,
               const int width)
      : _capacity(table_keys->dims()[0]),
        free_pointers(false),
        table_keys(table_keys->data<key_type>()),
        table_vals(table_vals->data<val_type>()),
        _divisor(divisor),
        _width(width) {}
  ~GPUHashTable() {
    if (free_pointers) {
      cudaFree(table_keys);
      cudaFree(table_vals);
    }
  }
  void insert_coords(const GPUContext& dev_ctx, const DenseTensor& coords);
  void lookup_coords(const GPUContext& dev_ctx,
                     const DenseTensor& coords,
                     const int* kernel_sizes,
                     const int* tensor_strides,
                     int kernel_volume,
                     DenseTensor* results);
  int get_divisor() { return _divisor; }
  int get_capacity() { return _capacity; }
};

using hashtable = GPUHashTable<int64_t, int>;
using hashtable32 = GPUHashTable<int, int>;

template <typename key_type = int64_t, typename val_type = int>
__global__ void insert_coords_kernel(key_type* table_keys,
                                     val_type* table_vals,
                                     const int* coords,
                                     int n,
                                     int _capacity,
                                     int _width) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < n) {
    key_type key = (key_type)(hash_func_64b(coords + idx * _width, _width));
    int value = idx + 1;
    int slot = hash(key, _capacity);
    while (true) {
      key_type prev = atomicCAS(&table_keys[slot], EMPTY_CELL, key);
      if (prev == EMPTY_CELL || prev == key) {
        table_vals[slot] = value;
        return;
      }
      slot = (slot + 1) % _capacity;
    }
  }
}

template <typename key_type = int64_t, typename val_type = int, bool odd>
__global__ void lookup_coords_kernel(key_type* table_keys,
                                     val_type* table_vals,
                                     const int* coords,
                                     val_type* vals,
                                     const int* kernel_sizes,
                                     const int* strides,
                                     int n,
                                     int _capacity,
                                     int kernel_volume,
                                     int _width) {
  int tidx = blockIdx.x * blockDim.x + threadIdx.x;
  int idx = tidx / kernel_volume;
  if (idx >= n) return;
  int _kernel_idx = tidx % kernel_volume;
  int kernel_idx = _kernel_idx;
  const int* in_coords = coords + _width * idx;
  int coords_out[4];
  // coords_out[2] = in_coords[2];
  // coords_out[3] = in_coords[3];
  coords_out[0] = in_coords[0];

  if constexpr (odd) {
#pragma unroll
    for (int i = 0; i <= _width - 2; i++) {
      int cur_offset = _kernel_idx % kernel_sizes[i];
      cur_offset -= (kernel_sizes[i] - 1) / 2;
      coords_out[i + 1] = in_coords[i + 1] * strides[i] + cur_offset;
      _kernel_idx /= kernel_sizes[i];
    }
  } else {
#pragma unroll
    for (int i = _width - 2; i >= 0; i--) {
      int cur_offset = _kernel_idx % kernel_sizes[i];
      cur_offset -= (kernel_sizes[i] - 1) / 2;
      coords_out[i + 1] = in_coords[i + 1] * strides[i] + cur_offset;
      _kernel_idx /= kernel_sizes[i];
    }
  }

  if (idx < n) {
    key_type key = (key_type)(hash_func_64b(coords_out, _width));
    int slot = hash(key, _capacity);

    while (true) {
      key_type cur_key = table_keys[slot];
      if (key == cur_key) {
        vals[idx * kernel_volume + kernel_idx] =
            table_vals[slot] -
            1;  // need to subtract 1 to avoid extra operations in python
      }
      if (table_keys[slot] == EMPTY_CELL) {
        return;
      }
      slot = (slot + 1) % _capacity;
    }
  }
}

__global__ void set_kernel_sizes_and_strides_tensor(int* kernel_size_tensor,
                                                    int* strides_tensor,
                                                    int kernel_size0,
                                                    int kernel_size1,
                                                    int kernel_size2,
                                                    int stride0,
                                                    int stride1,
                                                    int stride2) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < 6) {
    switch (idx) {
      case 0:
        kernel_size_tensor[idx] = kernel_size0;
        break;
      case 1:
        kernel_size_tensor[idx] = kernel_size1;
        break;
      case 2:
        kernel_size_tensor[idx] = kernel_size2;
        break;
      case 3:
        strides_tensor[idx - 3] = stride0;
        break;
      case 4:
        strides_tensor[idx - 3] = stride1;
        break;
      case 5:
        strides_tensor[idx - 3] = stride2;
        break;
      default:
        break;
    }
  }
}

template <typename key_type, typename val_type>
void GPUHashTable<key_type, val_type>::insert_many_coords(
    const GPUContext& dev_ctx, const int* coords, const int n) {
  insert_coords_kernel<key_type, val_type>
      <<<(n + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, dev_ctx.stream()>>>(
          table_keys, table_vals, coords, n, _capacity, _width);
}

template <typename key_type, typename val_type>
void GPUHashTable<key_type, val_type>::insert_coords(
    const GPUContext& dev_ctx, const DenseTensor& coords) {
  insert_many_coords(dev_ctx, coords.data<int>(), coords.dims()[0]);
}

template <typename key_type, typename val_type>
void GPUHashTable<key_type, val_type>::lookup_many_coords(
    const GPUContext& dev_ctx,
    const int* coords,
    val_type* results,
    const int* kernel_sizes,
    const int* strides,
    const int n,
    const int kernel_volume) {
  if (kernel_volume % 2)
    lookup_coords_kernel<key_type, val_type, true>
        <<<(n * kernel_volume + BLOCK_SIZE - 1) / BLOCK_SIZE,
           BLOCK_SIZE,
           0,
           dev_ctx.stream()>>>(table_keys,
                               table_vals,
                               coords,
                               results,
                               kernel_sizes,
                               strides,
                               n,
                               _capacity,
                               kernel_volume,
                               _width);
  else
    lookup_coords_kernel<key_type, val_type, false>
        <<<(n * kernel_volume + BLOCK_SIZE - 1) / BLOCK_SIZE,
           BLOCK_SIZE,
           0,
           dev_ctx.stream()>>>(table_keys,
                               table_vals,
                               coords,
                               results,
                               kernel_sizes,
                               strides,
                               n,
                               _capacity,
                               kernel_volume,
                               _width);
}

template <typename key_type, typename val_type>
void GPUHashTable<key_type, val_type>::lookup_coords(const GPUContext& dev_ctx,
                                                     const DenseTensor& coords,
                                                     const int* kernel_sizes,
                                                     const int* strides,
                                                     const int kernel_volume,
                                                     DenseTensor* results) {
  int32_t* results_data = results->data<int32_t>();
  lookup_many_coords(dev_ctx,
                     coords.data<int>(),
                     results_data,
                     kernel_sizes,
                     strides,
                     coords.dims()[0],
                     kernel_volume);
}

template <typename IntT>
void build_sparse_conv_kmap(const GPUContext& dev_ctx,
                            const SparseCooTensor& x,
                            const std::string& key,
                            const std::vector<int>& kernel_sizes,
                            const std::vector<int>& strides,
                            const int kernel_volume,
                            const bool is2D,
                            SparseCooTensor* out) {
  int nnz = x.nnz();
  const KmapCache* in_kmap_cache_ptr = x.GetKmapCache(key);
  out->ClearKmaps();
  KmapCache* out_kmap_cache_ptr = nullptr;
  bool to_insert = false;
  if (in_kmap_cache_ptr == nullptr) {
    KmapCache kmap_cache;
    out_kmap_cache_ptr = out->SetKmapCache(key, kmap_cache);
    if (out_kmap_cache_ptr->hashmap_keys == nullptr) {
      DenseTensor* tmp_hashmap_keys = new DenseTensor();
      tmp_hashmap_keys->Resize({2 * x.nnz()});
      dev_ctx.template Alloc<IntT>(tmp_hashmap_keys);
      funcs::SetConstant<GPUContext, IntT> set_zero;
      set_zero(dev_ctx, tmp_hashmap_keys, static_cast<IntT>(0));
      out_kmap_cache_ptr->hashmap_keys = tmp_hashmap_keys;
      to_insert = true;
    }
    if (out_kmap_cache_ptr->hashmap_values == nullptr) {
      DenseTensor* tmp_hashmap_values = new DenseTensor();
      tmp_hashmap_values->Resize({2 * x.nnz()});
      dev_ctx.template Alloc<int32_t>(tmp_hashmap_values);
      funcs::SetConstant<GPUContext, int32_t> set_zero;
      set_zero(dev_ctx, tmp_hashmap_values, static_cast<int32_t>(0));
      out_kmap_cache_ptr->hashmap_values = tmp_hashmap_values;
    }

    if (out_kmap_cache_ptr->coords == nullptr) {
      DenseTensor* tmp_indices = new DenseTensor();
      tmp_indices->Resize({x.indices().dims()[1], x.indices().dims()[0]});
      dev_ctx.template Alloc<int32_t>(tmp_indices);
      // transpose indices
      std::vector<int> perm = {1, 0};
      funcs::TransposeGPUKernelDriver<int32_t>(
          dev_ctx, x.indices(), perm, tmp_indices);
      out_kmap_cache_ptr->coords = tmp_indices;
    }

    const int divisor = 128;
    const int width = is2D ? 3 : 4;
    auto hashmap =
        GPUHashTable<IntT, int32_t>(out_kmap_cache_ptr->hashmap_keys,
                                    out_kmap_cache_ptr->hashmap_values,
                                    divisor,
                                    width);
    if (to_insert) {
      hashmap.insert_coords(dev_ctx, *(out_kmap_cache_ptr->coords));
    }

    DenseTensor* tmp_out_in_map = new DenseTensor();
    tmp_out_in_map->Resize(
        {(x.nnz() + divisor - 1) / divisor * divisor, kernel_volume});
    dev_ctx.template Alloc<int32_t>(tmp_out_in_map);
    out_kmap_cache_ptr->out_in_map = tmp_out_in_map;
    funcs::SetConstant<GPUContext, int32_t> set_neg_one;
    set_neg_one(
        dev_ctx, out_kmap_cache_ptr->out_in_map, static_cast<int32_t>(-1));

    // need to put kernel_sizes and strides to GPU
    auto kernel_sizes_tensor = phi::Empty<int32_t>(dev_ctx, {3});
    auto strides_tensor = phi::Empty<int32_t>(dev_ctx, {3});
    set_kernel_sizes_and_strides_tensor<<<1, 32, 0, dev_ctx.stream()>>>(
        kernel_sizes_tensor.data<int32_t>(),
        strides_tensor.data<int32_t>(),
        kernel_sizes[0],
        kernel_sizes[1],
        kernel_sizes[2],
        strides[0],
        strides[1],
        strides[2]);
    hashmap.lookup_coords(dev_ctx,
                          *(out_kmap_cache_ptr->coords),
                          kernel_sizes_tensor.data<int32_t>(),
                          strides_tensor.data<int32_t>(),
                          kernel_volume,
                          out_kmap_cache_ptr->out_in_map);

  } else {
    // out tensor takes the kmaps from x
    out->SetKmaps(x.GetKmaps());
    // force clear the kmaps of x
    const_cast<SparseCooTensor&>(x).ClearKmaps();
  }
  const KmapCache* new_out_kmap_cache_ptr = out->GetKmapCache(key);
  assert(new_out_kmap_cache_ptr != nullptr);
  assert(new_out_kmap_cache_ptr->hashmap_keys != nullptr);
  assert(new_out_kmap_cache_ptr->hashmap_values != nullptr);
  assert(new_out_kmap_cache_ptr->coords != nullptr);
  assert(new_out_kmap_cache_ptr->out_in_map != nullptr);
  return;
}
}  // namespace sparse
}  // namespace phi