#include "hash/hash_fucntion.cuh"
#include "lru_cache.h"
#include <vector>
#include <string_view>
#include <memory>

//template <typename T>
//constexpr auto type_name() {
//    std::string_view name, prefix, suffix;
//#ifdef __clang__
//    name = __PRETTY_FUNCTION__;
//  prefix = "auto type_name() [T = ";
//  suffix = "]";
//#elif defined(__GNUC__)
//    name = __PRETTY_FUNCTION__;
//    prefix = "constexpr auto type_name() [with T = ";
//    suffix = "]";
//#elif defined(_MSC_VER)
//    name = __FUNCSIG__;
//  prefix = "auto __cdecl type_name<";
//  suffix = ">(void)";
//#endif
//    name.remove_prefix(prefix.size());
//    name.remove_suffix(suffix.size());
//    return name;
//}
namespace gpucache {
    constexpr unsigned int defaultBlockX = 256;
    constexpr unsigned int warpsize = 32;
    constexpr unsigned int defaultNumWarpsPerBlock = defaultBlockX / warpsize;
    // bucket_id + key
    constexpr unsigned int uint32SharedMemorySize = 2 * sizeof(uint32_t) * defaultNumWarpsPerBlock * warpsize;
    constexpr unsigned int uint64SharedMemorySize =
            (sizeof(uint64_t) + sizeof(uint32_t)) * defaultNumWarpsPerBlock * warpsize;

    struct ThreadCtx {
        __device__ ThreadCtx() {
            auto global_thread_id = blockIdx.x * blockDim.x + threadIdx.x;
            global_warp_idx = global_thread_id / warpsize;
            block_warp_idx = threadIdx.x / warpsize;
            lane_id = threadIdx.x % warpsize;
            num_warps = blockDim.x * gridDim.x / warpsize;
        }

        uint32_t global_warp_idx;
        uint32_t block_warp_idx;
        uint32_t num_warps;
        uint32_t lane_id;
    };

    // for test
    template<typename KeyType>
    __global__ void CollectMissingKeysNew(uint32_t num_query, KeyType *keys,
                                          bool *find_mask, uint32_t *n_missing,
                                          KeyType *missing_keys) {

        ThreadCtx ctx{};

        for (auto offset = ctx.global_warp_idx * warpsize; offset < num_query;
             offset += ctx.num_warps * warpsize) {
            auto idx = offset + ctx.lane_id;
            if (!find_mask[idx]) {
                uint32_t base_missing_idx = atomicAdd(n_missing, 1);
                missing_keys[base_missing_idx] = keys[idx];
            }
        }
    }

    struct WarpMutex {
    public:
        __device__ WarpMutex() : flag(0) {}

        ~WarpMutex() = default;

        WarpMutex(const WarpMutex &) = delete;

        WarpMutex &operator=(const WarpMutex &) = delete;

        WarpMutex(WarpMutex &&) = delete;

        WarpMutex &operator=(WarpMutex &&) = delete;

        __device__ void Lock(ThreadCtx &ctx, uint32_t bucket_id) {
            if (ctx.lane_id == 0) {
                while (atomicCAS(&flag, 0, 1) != 0) {}
            }
            __threadfence();
            __syncwarp();
        }

        __device__ void UnLock(ThreadCtx &ctx) {
            __syncwarp();
            __threadfence();
            if (ctx.lane_id == 0) {
                auto i = atomicExch(&flag, 0);
                // printf("bucket id: Release Lock, flag = %u\n", i);
            }
        }

        //    private:
        uint32_t flag;
    };

    __global__ void initLocks(uint32_t n_bucket, void *bucketMutexes) {
        uint32_t global_thread_idx = blockDim.x * blockIdx.x + threadIdx.x;
        if (global_thread_idx < n_bucket) {
            new(reinterpret_cast<WarpMutex *>(bucketMutexes) + global_thread_idx)
                    WarpMutex();
        }
    }

    __global__ void checkLocks(uint32_t n_bucket, void *bucketMutexes) {
        uint32_t global_thread_idx = blockDim.x * blockIdx.x + threadIdx.x;
        // printf("thread %u CUDA_CHECK lock\n",global_thread_idx);
        if (global_thread_idx < n_bucket) {
            auto mutex =
                    reinterpret_cast<WarpMutex *>(bucketMutexes) + global_thread_idx;
            if (mutex->flag != 0u && mutex->flag != 1u) {
                printf("bucket id %u not equal 0 or 1, is %u\n", global_thread_idx,
                       mutex->flag);
            }
        }
    }

    template<typename KeyType, typename ElemType>
    class LRUCache;

    template<typename KeyType, typename ElemType>
    struct BucketView;

    template<typename KeyType, typename ElemType>
    __device__ __host__ BucketView<KeyType, ElemType>
    setBucketView(ThreadCtx ctx, KeyType *cache_keys, ElemType *cache_values,
                  uint8_t *cache_timestamps, void *cache_mutexes,
                  uint32_t num_elem_per_value, uint32_t bucket_id);

    template<typename KeyType, typename ElemType>
    class LRUCache : public Cache<KeyType, ElemType> {

        friend BucketView<KeyType, ElemType> __device__ __host__ setBucketView<KeyType, ElemType>(
                ThreadCtx ctx, KeyType *cache_keys, ElemType *cache_values,
                uint8_t *cache_timestamps, void *cache_mutexes,
                uint32_t num_elem_per_value, uint32_t bucket_id);

    public:
        explicit LRUCache(const CacheConfig &cfg);


        ~LRUCache();


        uint32_t KeySize() override;

        uint32_t ValueSize() override;

        uint64_t Capacity() override;

        uint32_t NumElemsPerValue() override;

        uint32_t MaxQueryNum() override;

        uint32_t NBucket();

        int8_t DeviceId() override;

        uint32_t Dim() override;

        void Test();

        // for test
        // void *Mutex() { return bucketMutexes; }

        CacheConfig::CacheEvictStrategy Strategy() override;

        void Clear() override;

        void Get(cudaStream_t stream, uint32_t num_query, KeyType *queries,
                 ElemType *results, bool *find_mask) override;

        void Put(cudaStream_t stream, uint32_t num_query, KeyType *put_keys,
                 ElemType *put_values, uint32_t *n_evict = nullptr,
                 KeyType *evict_keys = nullptr) override;

    // private:
        KeyType *keys;
        ElemType *values;
        uint8_t *timestamps{};
        uint32_t nbucket; // 32 values for one bucket
        void *bucketMutexes{};

        // CacheConfig::CacheEvictStrategy strategy;
        uint64_t capacity;
        uint32_t keySize;
        uint32_t valueSize;
        uint32_t numElemPerValue; // embedding dim
        int8_t device_id;
        uint32_t dim;

        // store missing keys and indices for Evict
        KeyType *queryKeyBuffer{};
        uint32_t *queryIndiceBuffer{};
        uint32_t maxQueryNum;
    };



    template<typename KeyType, typename ElemType>
    LRUCache<KeyType, ElemType>::LRUCache(const CacheConfig &cfg)
            : keySize(cfg.keySize), valueSize(cfg.valueSize), capacity(cfg.capacity),
              maxQueryNum(cfg.maxQueryNum), device_id(cfg.deviceId), dim(cfg.dim) {
//        std::cout << "LRUCache constructor" << std::endl;
        numElemPerValue = valueSize / sizeof(ElemType);
        nbucket = (capacity + warpsize - 1) / warpsize;
        printf("LRUCache: keySize: %lu, valueSize: %u, dim: %u, capacity: %lu, "
               "maxQueryNum: %u, deviceId: %u\n",
               sizeof(KeyType), valueSize, dim, capacity, maxQueryNum, device_id);

        CUDA_CHECK(cudaMalloc((void **) &keys, capacity * sizeof(KeyType)));
        CUDA_CHECK(cudaMalloc((void **) &values, capacity * valueSize));
        CUDA_CHECK(cudaMalloc((void **) &timestamps, capacity * sizeof(uint8_t)));
        CUDA_CHECK(cudaMalloc((void **) &bucketMutexes, nbucket * sizeof(WarpMutex)));
        dim3 block(defaultBlockX);
        dim3 grid((nbucket + defaultBlockX - 1) / defaultBlockX);
        initLocks<<<grid, block>>>(nbucket, bucketMutexes);
        CUDA_CHECK(cudaMalloc((void **) &queryKeyBuffer, maxQueryNum * sizeof(KeyType)));
        CUDA_CHECK(cudaMalloc((void **) &queryIndiceBuffer,
                         maxQueryNum * sizeof(uint32_t)));
    }

    template<typename KeyType, typename ElemType>
    void LRUCache<KeyType, ElemType>::Clear() {
        CUDA_CHECK(cudaMemset(keys, 0, capacity * sizeof(KeyType)));
        CUDA_CHECK(cudaMemset(timestamps, 0, capacity * sizeof(uint32_t)));
        dim3 block(defaultBlockX);
        dim3 grid((nbucket + defaultBlockX - 1) / defaultBlockX);
        initLocks<<<grid, block>>>(nbucket, bucketMutexes
        );
    }

    template<typename KeyType, typename ElemType>
    LRUCache<KeyType, ElemType>::~LRUCache() {
        // TODO need to add CudaDeviceGuard
        CUDA_CHECK(cudaFree(keys));
        CUDA_CHECK(cudaFree(values));
        CUDA_CHECK(cudaFree(timestamps));
        CUDA_CHECK(cudaFree(bucketMutexes));
        CUDA_CHECK(cudaFree(queryKeyBuffer));
        CUDA_CHECK(cudaFree(queryIndiceBuffer));
        // std::cout << "LRUCache destructed!" << std::endl;
    }

    template<typename KeyType, typename ElemType>
    uint32_t LRUCache<KeyType, ElemType>::KeySize() { return keySize; }

    template<typename KeyType, typename ElemType>
    uint32_t LRUCache<KeyType, ElemType>::ValueSize() { return valueSize; }

    template<typename KeyType, typename ElemType>
    uint64_t LRUCache<KeyType, ElemType>::Capacity() { return capacity; }

    template<typename KeyType, typename ElemType>
    uint32_t LRUCache<KeyType, ElemType>::NumElemsPerValue() { return numElemPerValue; }

    template<typename KeyType, typename ElemType>
    uint32_t LRUCache<KeyType, ElemType>::MaxQueryNum() { return maxQueryNum; }

    template<typename KeyType, typename ElemType>
    uint32_t LRUCache<KeyType, ElemType>::NBucket() { return nbucket; }

    template<typename KeyType, typename ElemType>
    int8_t LRUCache<KeyType, ElemType>::DeviceId() { return device_id; }

    template<typename KeyType, typename ElemType>
    uint32_t LRUCache<KeyType, ElemType>::Dim() { return dim; }


    template<typename KeyType, typename ElemType>
    CacheConfig::CacheEvictStrategy LRUCache<KeyType, ElemType>::Strategy() {
        return CacheConfig::LRU;
    }


// bucket view for warp put/get
    template<typename KeyType, typename ElemType>
    struct BucketView {
        __device__ BucketView(KeyType *k, ElemType *v, uint8_t *ts, WarpMutex *m,
                              uint32_t num_elems_per_value)
                : mutex(m), bkeys(k), bvalues(v), btimestamps(ts),
                  num_elems_per_value(num_elems_per_value) {}

        __device__ int Get(const ThreadCtx &ctx, const KeyType key) {
            KeyType lane_key = bkeys[ctx.lane_id];
            uint8_t ts = btimestamps[ctx.lane_id];
            bool exist = (ts > 0 && lane_key == key);
            uint32_t exist_mask = __ballot_sync(0xFFFFFFFF, exist);
            int slot_num = __ffs(static_cast<int>(exist_mask)) - 1;
            if (exist_mask == 0) { // not found
                return -1;
            } else {
                auto slot_ts = __shfl_sync(0xFFFFFFFF, ts, slot_num);
                if (ts > slot_ts) {
                    btimestamps[ctx.lane_id]--;
                }
                if (ctx.lane_id == slot_num) {
                    btimestamps[ctx.lane_id] = warpsize;
                }
                return slot_num; // return exist slot num
            };
        }

        __device__ int TryPut(const ThreadCtx &ctx, const KeyType key) {
            KeyType lane_key = bkeys[ctx.lane_id];
            uint8_t ts = btimestamps[ctx.lane_id];
            bool exist = (ts > 0 && lane_key == key);
            uint32_t exist_mask = __ballot_sync(0xFFFFFFFF, exist);
            int slot_num = __ffs(static_cast<int>(exist_mask)) - 1;
            if (exist_mask != 0) { // find key just update value
                uint32_t slot_ts = __shfl_sync(0xFFFFFFFF, ts, slot_num);
                if (ts > slot_ts) {
                    ts--;
                } else if (ctx.lane_id == slot_num) {
                    ts = warpsize;
                }
                __syncwarp();
            }
            if (slot_num == -1) { // key don't exist
                uint32_t emptyMask = __ballot_sync(0xFFFFFFFF, ts != 0);
                if (emptyMask != 0xFFFFFFFF) { // bucket not full
                    slot_num = __popc(static_cast<int>(emptyMask));
                    if (ts > 0) {
                        ts--;
                    } else if (ctx.lane_id == slot_num) {
                        bkeys[slot_num] = key; // insert key
                        ts = warpsize;
                    }
                }
                __syncwarp();
            }

            if (slot_num != -1) {
                btimestamps[ctx.lane_id] = ts;
            }
            return slot_num;
        }

        __device__ int Evict(const ThreadCtx &ctx, const KeyType key,
                             KeyType *evict_key) {
            KeyType lane_key = bkeys[ctx.lane_id];
            uint8_t ts = btimestamps[ctx.lane_id];
            uint32_t evict_mask = __ballot_sync(0xFFFFFFFF, ts == 1);
            int slot_num = __ffs(static_cast<int>(evict_mask)) - 1;
            if (ts > 1) {
                ts--;
            } else if (ctx.lane_id == slot_num) {
                // printf("evict key %u\n",lane_key);
                *evict_key = lane_key;
                // printf("evict evict key %u\n",*evict_key);
                bkeys[slot_num] = key;
                ts = warpsize;
            }
            btimestamps[ctx.lane_id] = ts;
            return slot_num;
        }

        __device__ void ReadOneValue(const ThreadCtx &ctx, uint8_t slot_num,
                                     ElemType *out) {
            for (size_t i = ctx.lane_id; i < num_elems_per_value; i += warpsize) {
                out[i] = bvalues[slot_num * num_elems_per_value + i];
            }
        }

        __device__ void WriteOneValue(const ThreadCtx &ctx, uint8_t slot_num,
                                      ElemType *v) {
            for (size_t i = ctx.lane_id; i < num_elems_per_value; i += warpsize) {
                bvalues[slot_num * num_elems_per_value + i] = v[i];
            }
        }

        WarpMutex *mutex;
        KeyType *bkeys;
        ElemType *bvalues;
        uint8_t *btimestamps;
        uint32_t num_elems_per_value;
    };

    template<typename KeyType, typename ElemType>
    __device__ __host__ BucketView<KeyType, ElemType>
    setBucketView(ThreadCtx ctx, KeyType *cache_keys, ElemType *cache_values,
                  uint8_t *cache_timestamps, void *cache_mutexes,
                  uint32_t num_elem_per_value, uint32_t bucket_id) {
        // auto p = reinterpret_cast<WarpMutex *>(cache_mutexes) + bucket_id;
        // if (p->flag != 0 && p->flag != 1) {
        //     if (ctx.lane_id == 0) {
        //         printf("warp %u bucket %u flag not equal 0 or 1 is
        //         %u\n",ctx.global_warp_idx,bucket_id,p->flag);
        //     }
        // }

        return BucketView<KeyType, ElemType>(
                cache_keys + bucket_id * warpsize,
                cache_values + bucket_id * warpsize * num_elem_per_value,
                cache_timestamps + bucket_id * warpsize,
                reinterpret_cast<WarpMutex *>(cache_mutexes) + bucket_id,
                num_elem_per_value);
    }

    template<typename KeyType, typename ElemType>
    __global__ void GetInternal(KeyType *cache_keys, ElemType *cache_values,
                                uint8_t *cache_timestamps, void *cache_mutexes,
                                uint32_t nbucket, uint32_t num_elem_per_value,
                                uint32_t num_query, KeyType *queries,
                                ElemType *results, bool *find_mask) {
        ThreadCtx ctx{};
        __shared__ KeyType blockQueryKeys[defaultNumWarpsPerBlock][warpsize];
        __shared__ uint32_t blockBucketIds[defaultNumWarpsPerBlock][warpsize];
        for (uint32_t offset = ctx.global_warp_idx * warpsize; offset < num_query;
             offset += ctx.num_warps * warpsize) {
            uint32_t n_query = min(warpsize, num_query - offset);
            if (ctx.lane_id < n_query) {
                uint32_t idx = offset + ctx.lane_id;
                KeyType key = queries[idx];
                size_t hash = getHash<KeyType>(key);
                uint32_t bucket_id = hash % nbucket;
                blockQueryKeys[ctx.block_warp_idx][ctx.lane_id] = key;
                blockBucketIds[ctx.block_warp_idx][ctx.lane_id] = bucket_id;
            }
            __syncwarp();

            //          32 threads compare it own slot with key
            //          if find parallel write to result
            for (uint32_t i = 0; i < n_query; i++) {
                uint32_t idx = offset + i;
                KeyType key = blockQueryKeys[ctx.block_warp_idx][i];
                uint32_t bucket_id = blockBucketIds[ctx.block_warp_idx][i];

                auto bucket = setBucketView<KeyType, ElemType>(
                        ctx, cache_keys, cache_values, cache_timestamps, cache_mutexes,
                        num_elem_per_value, bucket_id);

                bucket.mutex->Lock(ctx, bucket_id);

                int slot_num = bucket.Get(ctx, key);
                if (slot_num != -1) {
                    find_mask[idx] = true;
                    bucket.ReadOneValue(ctx, slot_num, &results[idx * num_elem_per_value]);
                }
                bucket.mutex->UnLock(ctx);
            }
            __syncwarp();
        }
    }

    template<typename KeyType, typename ElemType>
    __global__ void

    PutWithoutEvictInternal(KeyType *cache_keys, ElemType *cache_values,
                            uint8_t *cache_timestamps, void *cache_mutexes,
                            uint32_t nbucket, uint32_t num_elem_per_value,
                            uint32_t num_query, KeyType *put_keys,
                            ElemType *put_values, uint32_t *n_missing,
                            KeyType *missing_keys, uint32_t *missing_indices) {
        ThreadCtx ctx{};
        __shared__ KeyType blockPutKeys[defaultNumWarpsPerBlock][warpsize];
        __shared__ uint32_t blockBucketIds[defaultNumWarpsPerBlock][warpsize];
        for (uint32_t offset = ctx.global_warp_idx * warpsize; offset < num_query;
             offset += ctx.num_warps * warpsize) {
            uint32_t n_query = min(warpsize, num_query - offset);
            if (ctx.lane_id < n_query) {
                uint32_t idx = offset + ctx.lane_id;
                KeyType key = put_keys[idx];
                size_t hash = getHash(key);
                uint32_t bucket_id = hash % nbucket;
                blockPutKeys[ctx.block_warp_idx][ctx.lane_id] = key;
                blockBucketIds[ctx.block_warp_idx][ctx.lane_id] = bucket_id;
            }
            __syncwarp();

            uint32_t n_warp_missing = 0;
            KeyType warp_missing_key;
            uint32_t warp_missing_index = 0;

            // 32 threads handle n_query keys together instead of 1 thread for 1 key
            for (uint32_t i = 0; i < n_query; i++) {
                uint32_t idx = offset + i;
                KeyType key = blockPutKeys[ctx.block_warp_idx][i];
                // ElemType* Value = &put_values[idx];
                uint32_t bucket_id = blockBucketIds[ctx.block_warp_idx][i];

                auto bucket = setBucketView<KeyType, ElemType>(
                        ctx, cache_keys, cache_values, cache_timestamps, cache_mutexes,
                        num_elem_per_value, bucket_id);
                bucket.mutex->Lock(ctx, bucket_id);
                int slot_num = bucket.TryPut(ctx, key);
                if (slot_num != -1) { // insert value
                    // if(ctx.lane_id == 0){
                    //     printf("warp %u put key %u value %u\n",ctx.block_warp_idx, key,
                    //     put_values[idx*num_elem_per_value]);
                    // }
                    bucket.WriteOneValue(ctx, slot_num,
                                         &put_values[idx * num_elem_per_value]);
                } else { // bucket full record missing_key + idx for Evict
                    if (ctx.lane_id ==
                        n_warp_missing) { // i-th thread keep i-th missing key
                        // printf("warp %u lane_id %u try put key %u fail\n",
                        // ctx.global_warp_idx, ctx.lane_id, key);
                        warp_missing_key = key;
                        warp_missing_index = idx;
                    }
                    n_warp_missing += 1;
                }
                bucket.mutex->UnLock(ctx);
            }
            // reduce missing_key & idx
            if (n_warp_missing > 0) {
                uint32_t base_missing_idx = 0;
                if (ctx.lane_id == 0) {
                    base_missing_idx = atomicAdd(n_missing, n_warp_missing);
                }
                base_missing_idx = __shfl_sync(0xFFFFFFFF, base_missing_idx, 0);
                if (ctx.lane_id < n_warp_missing) {
                    // printf("warp %u thread %u add missing_key %u index %u in idx
                    // %u\n",ctx.global_warp_idx, ctx.lane_id,warp_missing_key,
                    // warp_missing_index,base_missing_idx + ctx.lane_id);
                    missing_keys[base_missing_idx + ctx.lane_id] = warp_missing_key;
                    missing_indices[base_missing_idx + ctx.lane_id] = warp_missing_index;
                }
            }
            __syncwarp();
        }
    }

    template<typename KeyType, typename ElemType>
    __global__ void

    EvictInternal(KeyType *cache_keys, ElemType *cache_values,
                  uint8_t *cache_timestamps, void *cache_mutexes, uint32_t nbucket,
                  uint32_t num_elem_per_value, ElemType *put_values,
                  uint32_t n_missing, KeyType *missing_keys,
                  uint32_t *missing_indices, uint32_t *num_evict,
                  KeyType *evict_keys) {
        ThreadCtx ctx{};
        __shared__ KeyType blockPutKeys[defaultNumWarpsPerBlock][warpsize];
        __shared__ uint32_t blockBucketIds[defaultNumWarpsPerBlock][warpsize];
        for (uint32_t offset = ctx.global_warp_idx * warpsize; offset < n_missing;
             offset += ctx.num_warps * warpsize) {
            uint32_t n_evict = min(warpsize, n_missing - offset);
            if (ctx.lane_id < n_evict) {
                uint32_t idx = offset + ctx.lane_id;
                KeyType key = missing_keys[idx];
                size_t hash = getHash(missing_keys[idx]);
                uint32_t bucket_id = hash % nbucket;
                blockPutKeys[ctx.block_warp_idx][ctx.lane_id] = key;
                blockBucketIds[ctx.block_warp_idx][ctx.lane_id] = bucket_id;
            }
            __syncwarp();
            for (uint32_t i = 0; i < n_evict; i++) {
                uint32_t idx = offset + i;
                KeyType key = blockPutKeys[ctx.block_warp_idx][i];
                uint32_t bucket_id = blockBucketIds[ctx.block_warp_idx][i];

                auto bucket = setBucketView<KeyType, ElemType>(
                        ctx, cache_keys, cache_values, cache_timestamps, cache_mutexes,
                        num_elem_per_value, bucket_id);
                bucket.mutex->Lock(ctx, bucket_id);
                KeyType evict_key = 0;
                int slot_num = bucket.Evict(ctx, key, &evict_key);
                // if(ctx.lane_id == 0){
                //     printf("evict key %u\n",evict_key);
                // }
                uint32_t evict_keys_idx = 0;
                if (num_evict && ctx.lane_id == slot_num) {
                    evict_keys_idx = atomicAdd(num_evict, 1);
                    evict_keys[evict_keys_idx] = evict_key;
                    // printf("evict_keys[%u] = %u\n",evict_keys_idx,evict_key);
                }
                bucket.WriteOneValue(ctx, slot_num,
                                     put_values +
                                     missing_indices[idx] * num_elem_per_value);
                // if(ctx.lane_id == 0){
                //     printf("warp %u put key %u value %u and evict key %u in bucket
                //     %u\n",ctx.block_warp_idx, key, put_values[idx*num_elem_per_value],
                //     evict_key, bucket_id);
                // }
                bucket.mutex->UnLock(ctx);
            }
        }
    }

// TODO switch to cuda stream
    template<typename KeyType, typename ElemType>
    void LRUCache<KeyType, ElemType>::Put(cudaStream_t stream, uint32_t num_query,
                                          KeyType *put_keys,

                                          ElemType *put_values, uint32_t *n_evict,
                                          KeyType *evict_keys) {

        assert(num_query <= maxQueryNum);
        if (num_query == 0) {
            return;
        }
        dim3 block(defaultBlockX);
        dim3 grid((num_query + defaultBlockX - 1) / defaultBlockX);
        uint32_t *n_missing;
        CUDA_CHECK(cudaMallocManaged(&n_missing, sizeof(uint32_t)));
        CUDA_CHECK(cudaMemset(n_missing, 0, sizeof(uint32_t)));

        if (std::is_same<KeyType, int32_t>::value || std::is_same<KeyType, uint32_t>::value) {
            PutWithoutEvictInternal<KeyType, ElemType>
            <<<grid, block, uint32SharedMemorySize, stream>>>(keys, values, timestamps, bucketMutexes, nbucket,
                                                              numElemPerValue, num_query, put_keys, put_values,
                                                              n_missing, queryKeyBuffer, queryIndiceBuffer);
            CUDA_CHECK(cudaStreamSynchronize(stream));
            EvictInternal<KeyType, ElemType>
            <<<grid, block, uint32SharedMemorySize, stream>>>(keys, values, timestamps, bucketMutexes, nbucket,
                                                              numElemPerValue, put_values, *n_missing, queryKeyBuffer,
                                                              queryIndiceBuffer, n_evict, evict_keys);
        } else if (std::is_same<KeyType, int64_t>::value || std::is_same<KeyType, uint64_t>::value) {
            PutWithoutEvictInternal<KeyType, ElemType>
            <<<grid, block, uint64SharedMemorySize, stream>>>(keys, values, timestamps, bucketMutexes, nbucket,
                                                              numElemPerValue, num_query, put_keys, put_values,
                                                              n_missing, queryKeyBuffer, queryIndiceBuffer);
            CUDA_CHECK(cudaStreamSynchronize(stream));
            EvictInternal<KeyType, ElemType>
            <<<grid, block, uint64SharedMemorySize, stream>>>(keys, values, timestamps, bucketMutexes, nbucket,
                                                              numElemPerValue, put_values, *n_missing, queryKeyBuffer,
                                                              queryIndiceBuffer, n_evict, evict_keys);
        } else {
            // std::cout << "unsupported keytype" << std::endl;
            throw "Unsupported KeyType";
        }

        // printf("n_missing = %u\n",*n_missing);
        CUDA_CHECK(cudaFree(n_missing));
    }
    template<typename KeyType, typename ElemType>
    void LRUCache<KeyType, ElemType>::Test(){
        std::cout << "LRCache Test" << std::endl;
    }

// TODO switch to cuda stream
    template<typename KeyType, typename ElemType>
    void LRUCache<KeyType, ElemType>::Get(cudaStream_t stream, uint32_t num_query,
                                          KeyType *queries, ElemType *results,
                                          bool *find_mask) {
//        std::cout << "num_query" << num_query << "\n";
//        std::cout <<"type" << type_name<KeyType>() << "\n";
        assert(num_query <= maxQueryNum && "num_query should less than maxQueryNum");
        if (num_query == 0) {
            return;
        }
        dim3 block(defaultBlockX);
        dim3 grid((num_query + defaultBlockX - 1) / defaultBlockX);
        if (std::is_same<KeyType, int32_t>::value) {
            //std::cout << "lru_get_internal_int32" << "\n";
            GetInternal<KeyType, ElemType><<<grid, block, uint32SharedMemorySize, stream>>>(
                    keys, values, timestamps, bucketMutexes, nbucket, numElemPerValue,
                    num_query, queries, results, find_mask);
        } else if (std::is_same<KeyType, int64_t>::value) {
            //std::cout << "lru_get_internal_int64" << "\n";
            GetInternal<KeyType, ElemType><<<grid, block, uint64SharedMemorySize, stream>>>(
                    keys, values, timestamps, bucketMutexes, nbucket, numElemPerValue,
                    num_query, queries, results, find_mask);
        } else {
            //std::cout << "unsupported keytype" << std::endl;
            throw "Unsupported KeyType!";
        }
    }
} // namespace gpucache

namespace gpucache {
    template<typename ElemType>
    void NewInt32LRUCache(void **cache, CacheConfig cfg) {
        auto c = new LRUCache<int32_t, ElemType>(cfg);
        //std::cout << "NewInt32LRUCache LRUCache keysize " << c->keySize << std::endl;
        *cache = reinterpret_cast<void *>(c);
    }

    template<typename ElemType>
    void NewInt64LRUCache(void **cache, CacheConfig cfg) {
        auto c = new LRUCache<int64_t, ElemType>(cfg);
        *cache = reinterpret_cast<void *>(c);
    }

    LRUCacheWrapper::LRUCacheWrapper(torch::Tensor t, CacheConfig cfg) {
        //std::cout << "LRUCacheWrapper constructor" << std::endl;
        dtype = t.scalar_type();
        cache_cfg = cfg;
        key_is_int32 = cfg.capacity < INT32_MAX;
        if (key_is_int32) {
            kdtype = torch::kInt32;
            AT_DISPATCH_ALL_TYPES(t.scalar_type(), "int32_lru_cache", [&] {
                NewInt32LRUCache<scalar_t>(&lru_cache, cfg);
            });
        } else {
            kdtype = torch::kInt64;
            AT_DISPATCH_ALL_TYPES(t.scalar_type(), "int64_lru_cache", [&] {
                NewInt64LRUCache<scalar_t>(&lru_cache, cfg);
            });
        }
//        AT_DISPATCH_ALL_TYPES(dtype,"lru_test",[&]{
//            auto c = reinterpret_cast<LRUCache<int32_t,scalar_t>*>(lru_cache);
//            std::cout << "LRUCacheWrapper constructor LRUCache dim: " << c->Dim() << " keySize: " << c->KeySize() << "\n";
//        });

        // std::cout << "LRUCacheWrapper constructor dtype: " << dtype << " kdtype: " << kdtype << "\n";
    }

    LRUCacheWrapper::~LRUCacheWrapper() {
        if (key_is_int32) {
            AT_DISPATCH_ALL_TYPES(dtype, "int32_lru_destructor", [&] {
                auto cache = reinterpret_cast<LRUCache<int32_t, scalar_t> *>(lru_cache);
                delete cache;
            });
        } else {
            AT_DISPATCH_ALL_TYPES(dtype, "int64_lru_destructor", [&] {
                auto cache = reinterpret_cast<LRUCache<int64_t, scalar_t> *>(lru_cache);
                delete cache;
            });
        }
    //std::cout << "~LRUCacheWrapper()" << std::endl;
    }

    std::pair<torch::Tensor, torch::Tensor> LRUCacheWrapper::Get(uint32_t num_query, const torch::Tensor queries) {
        assert(queries.device().index() == cache_cfg.deviceId && "query keys and cache should be in the same device");
        assert(queries.size(0) == num_query && "num_query should equal queries.size(0)");
        assert(num_query <= cache_cfg.maxQueryNum && "num_query must less than max_query_num");
        assert(queries.dtype().toScalarType() == kdtype && "key type doesn't match");
        //std::cout << "Get queries: " << queries << std::endl;
//        printf("hello!!!\n");
        //std::cout << "Get num_query: " << num_query <<"\n";
        //std::cout << "Get queries.sizes() " << queries.sizes() << "\n";
        auto keys = queries.data_ptr();
        auto stream = at::cuda::getCurrentCUDAStream(cache_cfg.deviceId).stream();
        auto result = torch::empty({num_query, cache_cfg.dim},
                                   torch::dtype(dtype).device(torch::kCUDA, cache_cfg.deviceId));
        auto find_mask = torch::empty({num_query}, torch::dtype(torch::kBool).device(torch::kCUDA, cache_cfg.deviceId));

        if (key_is_int32) {
            AT_DISPATCH_ALL_TYPES(dtype, "int32_lru_get", [&] {
                LRUCache<int32_t, scalar_t> * cache = reinterpret_cast<LRUCache<int32_t, scalar_t> *>(lru_cache);
//                cache->Test();
                //std::cout <<  "keysize: " << cache->KeySize() << "\n";
                //std::cout << "dim: " << cache->Dim() << "\n";
                cache->Get(stream, num_query, reinterpret_cast<int32_t *>(keys),
                           reinterpret_cast<scalar_t *>(result.data_ptr()),
                           reinterpret_cast<bool *>(find_mask.data_ptr()));
            });
        } else {
            //std::cout << "int64_lru_get" << "\n";
            AT_DISPATCH_ALL_TYPES(dtype, "int64_lru_get", [&] {
                LRUCache<int64_t, scalar_t>* cache = reinterpret_cast<LRUCache<int64_t, scalar_t> *>(lru_cache);
                //std::cout << "cast success!" << std::endl;
                cache->Get(stream, num_query, reinterpret_cast<int64_t *>(keys),
                           reinterpret_cast<scalar_t *>(result.data_ptr()),
                           reinterpret_cast<bool *>(find_mask.data_ptr()));
            });
        }
        return std::make_pair(result, find_mask);
    }

    void LRUCacheWrapper::Put(uint32_t num_query, const torch::Tensor keys, const torch::Tensor values) {
        assert(keys.device().index() == cache_cfg.deviceId && cache_cfg.deviceId == values.device().index() &&
               "put keys, put values and cache should be in the same device");
        assert(num_query == keys.size(0) && keys.size(0) == values.size(0) && "dim 0 of keys,values and num_query should be equal");
        assert(num_query <= cache_cfg.maxQueryNum && "num_query must less than max_query_num");
        assert(values.dtype() == dtype);
        auto stream = at::cuda::getCurrentCUDAStream(cache_cfg.deviceId).stream();
        if (key_is_int32) {
            //std::cout << "int32_lru_put" << "\n";
            AT_DISPATCH_ALL_TYPES(dtype, "int32_lru_put", [&] {
                auto cache = reinterpret_cast<LRUCache<int32_t, scalar_t> *>(lru_cache);
                //std::cout << "put_lru_dim" << cache->Dim() << std::endl;
                cache->Put(stream, num_query, reinterpret_cast<int32_t *>(keys.data_ptr()),
                           reinterpret_cast<scalar_t *>(values.data_ptr()));
            });
        } else {
            //std::cout << "int64_lru_put" << "\n";
            AT_DISPATCH_ALL_TYPES(dtype, "int64_lru_put", [&] {
                auto cache = reinterpret_cast<LRUCache<int64_t, scalar_t> *>(lru_cache);
                cache->Put(stream, num_query, reinterpret_cast<int64_t *>(keys.data_ptr()),
                           reinterpret_cast<scalar_t *>(values.data_ptr()));
            });
        }
    }

    CacheConfig::CacheEvictStrategy LRUCacheWrapper::Strategy() {
        return CacheConfig::CacheEvictStrategy::LRU;
    }

    uint64_t LRUCacheWrapper::Capacity() { return cache_cfg.capacity; }

    uint32_t LRUCacheWrapper::KeySize() { return cache_cfg.keySize; }

    uint32_t LRUCacheWrapper::ValueSize() { return cache_cfg.valueSize; }

    uint32_t LRUCacheWrapper::MaxQueryNum() { return cache_cfg.maxQueryNum; }

    uint64_t LRUCacheWrapper::DeviceId() { return cache_cfg.deviceId; }

    uint32_t LRUCacheWrapper::Dim() { return cache_cfg.dim; }

    void LRUCacheWrapper::Clear() {
        if (key_is_int32) {
            AT_DISPATCH_ALL_TYPES(dtype, "int32_lru_clear", [&] {
                auto cache = reinterpret_cast<LRUCache<int32_t, scalar_t> *>(lru_cache);
                cache->Clear();
            });
        } else {
            AT_DISPATCH_ALL_TYPES(dtype, "int64_lru_clear", [&] {
                auto cache = reinterpret_cast<LRUCache<int64_t, scalar_t> *>(lru_cache);
                cache->Clear();
            });
        }
    }

    std::unique_ptr<LRUCacheWrapper> NewLRUCache(at::Tensor t, CacheConfig cfg) {
        return std::make_unique<LRUCacheWrapper>(t,cfg);
    }

} // namespace gpucache
