#pragma once
#include "cache.cuh"

namespace gpucache{
    template<typename KeyType, typename ElemType, CacheConfig::CacheEvictStrategy Strategy>
    struct BucketView;

    template<typename KeyType, typename ElemType>
    struct BucketView<KeyType,ElemType,CacheConfig::CacheEvictStrategy::LFU> {
        __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(warpFullMask, exist);
            int slot_num = __ffs(exist_mask) - 1;
            if (exist_mask == 0) { // not found
                return -1;
            } else {
                if (ctx.lane_id != slot_num && ts > 1) { // decrease the counter
                    btimestamps[ctx.lane_id] = ts - 1;
                }
                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(warpFullMask, exist);
            int slot_num = __ffs(exist_mask) - 1;
            if (exist_mask != 0) {
                if (ctx.lane_id == slot_num) {
                    ts = warpSize;
                }
                __syncwarp();
            }
            if (slot_num == -1) { // key don't exist
                uint32_t emptyMask = __ballot_sync(warpFullMask, ts != 0);
                if (emptyMask != warpFullMask) { // bucket not full
                    slot_num = __popc(emptyMask);
                    if (ctx.lane_id == slot_num) {
                        bkeys[slot_num] = key;
                        ts = warpSize;
                    }
                }
                __syncwarp();
            }

            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(warpFullMask, ts == 1);
            int slot_num = __ffs(evict_mask) - 1;
            if(slot_num == -1){ // no ts equals 1, do warp reduction find min ts > 0
                uint8_t min_ts = ts;
                uint32_t min_ts_lane_id = ctx.lane_id;
                uint32_t reduce_mask = __ballot_sync(warpFullMask, ts != 0);
                for(int offset = 16; offset > 0; offset /= 2){
                    uint8_t temp_ts = __shfl_down_sync(reduce_mask,ts,offset);
                    if(temp_ts != 0 && temp_ts < min_ts){
                        min_ts = temp_ts;
                        min_ts_lane_id = ctx.lane_id + offset;
                    }
                }
                slot_num = __shfl_sync(warpFullMask, min_ts_lane_id, 0);
//                if(ctx.lane_id == 0){
//                    printf("evict slot_num %d, ts = %u",slot_num,ts);
//                }
            }

            if(ctx.lane_id == slot_num) {
                *evict_key = lane_key;
                bkeys[slot_num] = key;
                btimestamps[slot_num] = warpSize;
            }
            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>
    class LFUCache: public Cache<KeyType, ElemType>{
    public:
        explicit LFUCache(CacheConfig cfg):Cache<KeyType,ElemType>(cfg){
            // this->manager = new LRUCacheManager<KeyType,ElemType>();
        }
        ~LFUCache(){
            // delete this->manager;
        }
    };

    template<typename KeyType, typename ElemType>
    LFUCache<KeyType,ElemType>* NewLFUCache(CacheConfig cfg){
        return new LFUCache<KeyType,ElemType>(cfg);
    }

}