#include "common.cuh"
#include "cache.h"
#include "hash/hash_fucntion.cuh"


namespace gpucache {

    constexpr unsigned int defaultBlockX = 256;
    constexpr unsigned int warpsize = 32;
    constexpr unsigned int defaultNumWarpsPerBlock = defaultBlockX / 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;
    };

    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) {
//                    printf("try get lock\n");
//                }
                while(true){
                    auto i = atomicCAS(&flag, 0, 1);
                    if(i == 0){
                        break;
                    }else{
                        printf("warp %u try get lock for bucket %u, i = %u\n",ctx.global_warp_idx, bucket_id,i);
                    }
                }
                // printf("bucket id: Get Lock\n");

            }
            __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 check lock\n",global_thread_idx);
        if (global_thread_idx < n_bucket) {
            auto mutex = reinterpret_cast<WarpMutex *>(bucketMutexes) + global_thread_idx;
            if(mutex->flag != 0 && mutex->flag != 1){
                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>
//    __global__ void
//    GetInternal(LRUCache<KeyType, ElemType> cache, uint32_t num_query, KeyType *queries,
//                ElemType *results, bool* find_mask);

//    template<typename KeyType, typename ElemType>
//    __global__ void
//    PutWithoutEvictInternal(LRUCache<KeyType, ElemType> &cache, uint32_t num_query, KeyType *put_keys,
//                            ElemType *put_values, uint32_t *n_missing, KeyType *missing_keys,
//                            uint32_t *missing_indices);
//
//    template<typename KeyType, typename ElemType>
//    __global__ void
//    EvictInternal(LRUCache<KeyType, ElemType> &cache, ElemType *put_values, uint32_t n_missing, KeyType *missing_keys,
//                  uint32_t *missing_indices,uint32_t* num_evict, KeyType* evict_keys);

    template<typename KeyType, typename ElemType>
    __device__ __host__ BucketView<KeyType, ElemType>
    setBucketView(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  __global__ void
//        GetInternal<KeyType, ElemType>(LRUCache<KeyType, ElemType> cache, uint32_t num_query, KeyType *queries,
//                    ElemType *results, bool* find_mask);

//        friend  __global__ void
//        PutWithoutEvictInternal<KeyType, ElemType>(LRUCache<KeyType, ElemType> &cache, uint32_t num_query, KeyType *put_keys,
//                                ElemType *put_values, uint32_t *n_missing, KeyType *missing_keys,
//                                uint32_t *missing_indices);
//
//        friend __global__ void
//        EvictInternal<KeyType, ElemType>(LRUCache<KeyType, ElemType> &cache, ElemType *put_values, uint32_t n_missing,
//                      KeyType *missing_keys,
//                      uint32_t *missing_indices, uint32_t* n_evict, KeyType* evict_keys);


        friend BucketView<KeyType, ElemType>
        setBucketView<KeyType, ElemType>(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) : keySize(cfg.keySize),
                                                    valueSize(cfg.valueSize),
                                                    capacity(cfg.capacity), maxQueryNum(cfg.maxQueryNum) {
            numElemPerValue = valueSize / sizeof(ElemType);
            nbucket = (capacity + warpsize - 1) / warpsize;
            printf("LRUCache: keySize: %u, valueSize: %u, capacity: %llu, nbucket: %u, maxquerymum: %u\n",
                   sizeof(KeyType), valueSize, capacity, nbucket, maxQueryNum);

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

        }

        ~LRUCache() {
            // TODO need to add CudaDeviceGuard
            CHECK(cudaFree(keys));
            CHECK(cudaFree(values));
            CHECK(cudaFree(timestamps));
            CHECK(cudaFree(bucketMutexes));
            CHECK(cudaFree(queryKeyBuffer));
            CHECK(cudaFree(queryIndiceBuffer))
        }


        uint32_t KeySize() override { return keySize; }

        uint32_t ValueSize() override { return valueSize; }

        uint64_t Capacity() override { return capacity; }

        uint32_t NumElemsPerValue() override { return numElemPerValue; }

        uint32_t MaxQueryNum() override { return maxQueryNum; }

        uint32_t NBucket() { return nbucket; }

        void* Mutex() {return bucketMutexes;}

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

        void Clear() override {
            CHECK(cudaMemset(keys, 0, capacity * sizeof(KeyType)));
            CHECK(cudaMemset(timestamps, 0, capacity * sizeof(uint32_t)));
            dim3 block(defaultBlockX);
            dim3 grid((nbucket + defaultBlockX - 1) / defaultBlockX);
            initLocks<<<grid, block>>>(nbucket, bucketMutexes);
        }


        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 *putkeys, ElemType *putvalues, uint32_t *n_evict,
                 KeyType *evict_keys) 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

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


    };

    // 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 max_ts = __ffs(static_cast<int>(__ballot_sync(0xFFFFFFFF,ts != 0)));
                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;
                    printf("warp %u lane_id %u update key %u\n", ctx.global_warp_idx,ctx.lane_id,key);
                }
                __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;
                        printf("warp %u lane_id %u try put key %u success\n",ctx.global_warp_idx,ctx.lane_id,key);
                    }
                }
                __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;
            uint32_t slot_ts = __shfl_sync(0xFFFFFFFF, ts, slot_num);
            if (ts > slot_ts) {
                ts--;
            } else if (ctx.lane_id == slot_num) {

                *evict_key = lane_key;
                bkeys[ctx.lane_id] = 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(KeyType *cache_keys, ElemType *cache_values, uint8_t *cache_timestamps, void *cache_mutexes,
                  uint32_t num_elem_per_value, uint32_t bucket_id) {
        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) {
//            if(ctx.lane_id == 0){
//                printf("warp %u handle data block %u\n",ctx.global_warp_idx,offset / 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;
                // printf("key %u is in bucket %u\n", key, bucket_id);
                blockQueryKeys[ctx.block_warp_idx][ctx.lane_id] = key;
                blockBucketIds[ctx.block_warp_idx][ctx.lane_id] = bucket_id;
            }
            __syncwarp();

//            if (ctx.lane_id == 0){
//                printf("warp %u hash collect query %u keys of data block %u\n", ctx.global_warp_idx, n_query, offset / warpsize);
//            }
//             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>(cache_keys, cache_values, cache_timestamps,
                                                               cache_timestamps, num_elem_per_value, bucket_id);

                if(ctx.global_warp_idx == 0 && ctx.lane_id == 0){
                    printf("thread %u try get lock for bucket %u\n", ctx.lane_id, bucket_id);
                }
//                if(ctx.lane_id == 0){
//                    printf("bucket %u lock flag %u\n",bucket_id,bucket.mutex->flag);
//                }
                bucket.mutex->Lock(ctx, bucket_id);
                if(ctx.global_warp_idx == 0 && ctx.lane_id == 0){
                    printf("thread %u get lock for bucket %u\n", ctx.lane_id, bucket_id);
                }

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

//                if(ctx.global_warp_idx == 0 && ctx.lane_id == 0){
//                    printf("thread %u release lock for bucket %u\n", ctx.lane_id, bucket_id);
//                }
//                if(ctx.lane_id == 0){
//                    printf("warp %u release lock for bucket %u\n", ctx.global_warp_idx, bucket_id);
//                }
//                if(ctx.global_warp_idx == 0 && ctx.lane_id == 0){
//                    printf("finish key %u\n",i);
//                }
            }
//            if (ctx.lane_id == 0){
//                printf("warp %u get %u keys of data block %u\n", ctx.global_warp_idx, n_query, offset / warpsize);
//            }
//            if(ctx.global_warp_idx == 0){
//                printf("thread %u finish\n",ctx.lane_id);
//            }
            __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>(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
                    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) {
                    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 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;
                blockPutKeys[ctx.block_warp_idx][ctx.lane_id] = missing_keys[idx];
                blockBucketIds[ctx.block_warp_idx][ctx.lane_id] = missing_indices[idx];
            }
            __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>(cache_keys, cache_values, cache_timestamps,
                                                               cache_mutexes, num_elem_per_value, bucket_id);
                bucket.mutex->Lock(ctx, bucket_id);
                KeyType evict_key;
                int slot_num = bucket.Evict(ctx, key, &evict_key);
                uint32_t evict_keys_idx = atomicAdd(num_evict, 1);
                evict_keys[evict_keys_idx] = evict_key;
                bucket.WriteOneValue(ctx, slot_num, put_values + missing_indices[idx] * num_elem_per_value);
                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 = 0;

        PutWithoutEvictInternal<KeyType, ElemType><<<grid, block>>>(keys, values, timestamps, bucketMutexes, nbucket,
                                                                    numElemPerValue, num_query, put_keys, put_values,
                                                                    &n_missing, queryKeyBuffer,
                                                                    queryIndiceBuffer);
        EvictInternal<KeyType, ElemType><<<grid, block>>>(keys, values, timestamps, bucketMutexes, numElemPerValue,
                                                          put_values, n_missing, queryKeyBuffer, queryIndiceBuffer,
                                                          n_evict, evict_keys);

    }

    template<typename KeyType>
    __global__ void
    CollectMissingKeys(uint32_t num_query, KeyType *keys, bool *find_mask, uint32_t *n_missing, KeyType *missing_keys) {
        ThreadCtx ctx{};

        __shared__ KeyType blockMissingKeys[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);
            uint32_t warp_n_missing = 0;
            uint32_t base_missing_idx = 0;
            uint32_t warp_missing_idx = 0;

            if (ctx.lane_id < n_query) {
                uint32_t idx = offset + ctx.lane_id;
                uint32_t warp_missing_mask = __ballot_sync(0xFFFFFFFF, !find_mask[idx]);
                warp_n_missing = __popc(warp_missing_mask);
                uint32_t lane_mask = 1 << ctx.lane_id;
                warp_missing_idx = __popc(~((lane_mask - 1) | lane_mask) & warp_missing_mask);
                blockMissingKeys[ctx.block_warp_idx][warp_missing_idx] = keys[idx];
            }
            __syncwarp();
            if (ctx.lane_id == 0) {
                printf("warp %u add %u to n_missing, before %u\n", ctx.global_warp_idx, warp_n_missing, *n_missing);
                base_missing_idx = atomicAdd(n_missing, warp_n_missing);
            }
            if (ctx.lane_id < warp_n_missing) {
                uint32_t write_idx = base_missing_idx + ctx.lane_id;
                missing_keys[write_idx] = blockMissingKeys[ctx.block_warp_idx][ctx.lane_id];
            }
            __syncwarp();
        }
    }

    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];
            }
        }
    }

    // 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) {
        assert(num_query <= maxQueryNum);
//        bool* find_mask;
//        CHECK(cudaMalloc(&find_mask,num_query));
        if (num_query == 0) { return; }
        dim3 block(defaultBlockX);
        dim3 grid((num_query + defaultBlockX - 1) / defaultBlockX);

        GetInternal<KeyType, ElemType><<<grid, block>>>(keys, values, timestamps, bucketMutexes, nbucket,
                                                        numElemPerValue, num_query, queries, results, find_mask);
        CHECK(cudaDeviceSynchronize());
        // CollectMissingKeys<><<<grid,block>>>(num_query,queries,find_mask,n_missing,missing_keys);
//       CHECK(cudaFree(find_mask));
    }
}

