#include "gtest/gtest.h"
#include <cuda_runtime.h>
#include <unordered_set>
#include <vector>
#include <numeric>
#include <random>
#include <algorithm>
#include "cache.cuh"
#include "lru_cache.cuh"
#include "lfu_cache.cuh"
#include "fifo_cache.cuh"
#include "cache_wrapper.cu"
#include "utils.cu"
#include "hash/hash_function.cu"



#define CUDA_CHECK(call)                                                       \
{                                                                              \
    const cudaError_t error = call;                                            \
    if (error != cudaSuccess)                                                  \
    {                                                                          \
        fprintf(stderr, "Error: %s:%d, ", __FILE__, __LINE__);                 \
        fprintf(stderr, "code: %d, reason: %s\n", error,                       \
                cudaGetErrorString(error));                                    \
    }                                                                          \
}


namespace gpucache{
//    template<typename KeyType>
//    __global__ void CollectMissingKeys(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];
//            }
//        }
//    }

    void TestCache(Cache<int32_t ,uint32_t>& cache, uint32_t num_elem_per_value){
        constexpr int warpsize = 32;
        std::unordered_set<uint32_t> in_cache;
        const size_t n_iter = 32;
        const uint32_t n_keys = 2048;


        int32_t* d_keys;
        int32_t* keys;
        uint32_t* d_values;
        uint32_t* values;
        uint32_t* d_n_missing;
        uint32_t* n_missing;
        int32_t* d_missing_keys;
        int32_t* missing_keys;
        uint32_t* d_n_evict;
        uint32_t* n_evict;
        int32_t* d_evict_keys;
        int32_t* evict_keys;
        uint32_t* missing_indexs;
        uint32_t* d_missing_indexs;

        const size_t keys_size = n_keys * sizeof(int32_t);

        const size_t values_size = n_keys * num_elem_per_value * sizeof(uint32_t);

        const size_t mask_size = n_keys * sizeof(bool);
        const size_t index_size = n_keys * sizeof(uint32_t);


        CUDA_CHECK(cudaMalloc(&d_keys,keys_size));
        CUDA_CHECK(cudaMalloc(&d_values, values_size));
        CUDA_CHECK(cudaMalloc(&d_n_missing,sizeof(uint32_t)));
        CUDA_CHECK(cudaMalloc(&d_missing_indexs, index_size));
        CUDA_CHECK(cudaMalloc(&d_missing_keys,keys_size));
        CUDA_CHECK(cudaMalloc(&d_evict_keys,keys_size));
        CUDA_CHECK(cudaMalloc(&d_n_evict,sizeof(uint32_t)));

        CUDA_CHECK(cudaMallocHost(&keys,keys_size));
        CUDA_CHECK(cudaMallocHost(&values,values_size));
        CUDA_CHECK(cudaMallocHost(&n_missing,sizeof(uint32_t)));
        CUDA_CHECK(cudaMallocHost(&missing_indexs,index_size));

        CUDA_CHECK(cudaMallocHost(&missing_keys,keys_size));

        CUDA_CHECK(cudaMallocHost(&evict_keys,keys_size));
        CUDA_CHECK(cudaMallocHost(&n_evict,sizeof(uint32_t)));




        std::vector<int32_t> random_keys(n_keys * n_iter);
        std::iota(random_keys.begin(),random_keys.end(),1);
        std::random_device rd;
        std::mt19937 g(rd());

        for (size_t iter = 0; iter < n_iter; iter++){

            std::cout << "iter " << iter << std::endl;
            uint32_t expect_n_missing = 0;
            CUDA_CHECK(cudaMemset(d_n_missing,0,sizeof(uint32_t)));
            CUDA_CHECK(cudaMemset(d_n_evict,0,sizeof(uint32_t)));
            // CUDA_CHECK(cudaMemset(d_missing_keys,0,keys_size));
            // CUDA_CHECK(cudaMemset(d_evict_keys, 0, keys_size));


            std::shuffle(random_keys.begin(), random_keys.end(),g);
            std::copy(random_keys.begin(),random_keys.begin()+n_keys,keys);
            // for (auto i = keys; i < keys + n_keys; i++){
            //     std::cout << *i << " ";
            // }
            // std::cout << std::endl;

            std::unordered_set<int32_t> expect_missing_keys_set;
            std::unordered_set<int32_t> keys_set; // store current iter keys
            for (size_t i = 0; i < n_keys; ++i) {
                keys_set.emplace(keys[i]);
                if (in_cache.count(keys[i]) == 0) {
                    expect_missing_keys_set.emplace(keys[i]);
                    expect_n_missing += 1;
                }
            }

            // test get

            std::cout << "test get" << std::endl;
            CUDA_CHECK(cudaMemcpy(d_keys,keys,keys_size,cudaMemcpyHostToDevice));
            cudaStream_t stream;
            CUDA_CHECK(cudaStreamCreate(&stream));
            dim3 cblock(256);
            dim3 cgrid(((cache.Capacity() + warpsize - 1) / warpsize  + cblock.x - 1) / cblock.x);
            // std::cout << "check locks..............." << std::endl;
            // checkLocks<<<cgrid, cblock>>>((cache.Capacity() + warpsize - 1) / warpsize, cache.Mutex());
            // CUDA_CHECK(cudaDeviceSynchronize());
            cache.Get(stream,n_keys,d_keys,d_values,d_missing_indexs,d_missing_keys,d_n_missing);
            CUDA_CHECK(cudaDeviceSynchronize());
            CUDA_CHECK(cudaMemcpy(missing_indexs,d_missing_indexs,keys_size,cudaMemcpyDeviceToHost));
            CUDA_CHECK(cudaMemcpy(missing_keys,d_missing_keys,keys_size,cudaMemcpyDeviceToHost));
            CUDA_CHECK(cudaMemcpy(n_missing,d_n_missing,sizeof(uint32_t),cudaMemcpyDeviceToHost));

            dim3 block(256);
            dim3 grid((n_keys + block.x - 1)/block.x);
            // CollectMissingKeys<uint32_t><<<grid,block>>>(n_keys,d_keys,d_find_mask,d_n_missing,d_missing_keys);
            // CollectMissingKeys<int32_t><<<grid,block>>>(n_keys,d_keys,d_find_mask,d_n_missing,d_missing_keys);
//            CUDA_CHECK(cudaDeviceSynchronize());
//            CUDA_CHECK(cudaMemcpy(n_missing,d_n_missing, sizeof(uint32_t),cudaMemcpyDefault));
//            CUDA_CHECK(cudaMemcpy(missing_keys,d_missing_keys,keys_size,cudaMemcpyDefault));
            ASSERT_EQ(expect_n_missing, *n_missing) << "expect_n_missing is " << expect_n_missing << " n_missing is " << *n_missing;
            std::unordered_set<int32_t> missing_keys_set(missing_keys, missing_keys + *n_missing);
            ASSERT_EQ(missing_keys_set,expect_missing_keys_set);
            // check get value
            CUDA_CHECK(cudaMemcpy(values,d_values,values_size,cudaMemcpyDefault));
            for (size_t i = 0; i < n_keys ; i += 1) {
                if(missing_keys_set.count(keys[i]) == 0){
                    ASSERT_EQ(keys[i] + 123,values[i * num_elem_per_value]) << "key[" << i << "] = " << keys[i] <<  " doesn't get correct value should be " << keys[i] + 123 << " get " << values[i * num_elem_per_value];
                }

            }


            //test put
            std::cout << "test put" << std::endl;
            memset(values, 0,values_size);
            for (size_t i = 0; i < n_keys; i++){ // for key i, the first number of embeddding[key] should be key i +123
                values[i * num_elem_per_value] = keys[i] + 123;
            }
            CUDA_CHECK(cudaMemcpy(d_values,values,values_size,cudaMemcpyDefault));
            cache.Put(stream,n_keys,d_keys,d_values,d_n_evict,d_evict_keys);
            CUDA_CHECK(cudaDeviceSynchronize());
            CUDA_CHECK(cudaMemcpy(n_evict,d_n_evict,sizeof(uint32_t),cudaMemcpyDefault));
            CUDA_CHECK(cudaMemcpy(evict_keys,d_evict_keys,*n_evict * sizeof(int32_t), cudaMemcpyDefault));

            std::unordered_set<int32_t> evict_keys_set(evict_keys, evict_keys + *n_evict);
            CUDA_CHECK(cudaDeviceSynchronize());
            std::cout << "check evict keys in cache or in current insert batch, n_evict = " << *n_evict << std::endl;
            for(size_t i = 0; i < *n_evict; i++){
                ASSERT_TRUE(in_cache.count(evict_keys[i]) > 0 || keys_set.count(evict_keys[i]) > 0) << "evict_keys[" << i << "] = " << evict_keys[i];
            }
            for (size_t i = 0; i < n_keys; ++i){in_cache.emplace(keys[i]);}
            for (size_t i = 0; i < *n_evict; ++i){in_cache.erase(evict_keys[i]);}

            // std::cout << "check locks..............." << std::endl;
            // checkLocks<<<cgrid, cblock>>>((cache.Capacity() + warpsize - 1) / warpsize, cache.Mutex());
            // CUDA_CHECK(cudaDeviceSynchronize());
            printf("---------------------------------------------------\n");

        }
        CUDA_CHECK(cudaFree(d_keys));
        CUDA_CHECK(cudaFree(d_values));
        CUDA_CHECK(cudaFree(d_n_missing));
        CUDA_CHECK(cudaFree(d_missing_keys));
        CUDA_CHECK(cudaFree(d_evict_keys));
        CUDA_CHECK(cudaFree(d_n_evict));
        CUDA_CHECK(cudaFree(d_missing_indexs));
        CUDA_CHECK(cudaFreeHost(keys));
        CUDA_CHECK(cudaFreeHost(values));
        CUDA_CHECK(cudaFreeHost(n_missing));
        CUDA_CHECK(cudaFreeHost(missing_indexs));
        CUDA_CHECK(cudaFreeHost(missing_keys));
        CUDA_CHECK(cudaFreeHost(evict_keys));
        CUDA_CHECK(cudaFreeHost(n_evict));
    }

    TEST(GPUCACHE,LRUCACHE){
        CacheConfig cfg{};
        cfg.strategy = CacheConfig::CacheEvictStrategy::LRU;

        cfg.value_size = 172*4;
        cfg.capacity = 4096 * 2;
        cfg.key_size = 4;
        cfg.max_query_num = 2048;
        cfg.device_id = 0;
        cfg.dim = 172;

        LRUCache<int32_t,uint32_t> *cache = NewLRUCache<int32_t,uint32_t>(cfg);
        TestCache(*cache,cfg.dim);

    }

    TEST(GPUCACHE,FIFOCACHE){
        CacheConfig cfg{};
        cfg.strategy = CacheConfig::CacheEvictStrategy::FIFO;

        cfg.value_size = 172*4;
        cfg.capacity = 4096 * 2;
        cfg.key_size = 4;
        cfg.max_query_num = 2048;
        cfg.device_id = 0;
        cfg.dim = 172;

        FIFOCache<int32_t,uint32_t> *cache = NewFIFOCache<int32_t,uint32_t>(cfg);
        TestCache(*cache,cfg.dim);

    }

    TEST(GPUCACHE,LFUCACHE){
        CacheConfig cfg{};
        cfg.strategy = CacheConfig::CacheEvictStrategy::LFU;

        cfg.value_size = 172*4;
        cfg.capacity = 4096 * 2;
        cfg.key_size = 4;
        cfg.max_query_num = 2048;
        cfg.device_id = 0;
        cfg.dim = 172;

        LFUCache<int32_t,uint32_t> *cache = NewLFUCache<int32_t,uint32_t>(cfg);
        TestCache(*cache,cfg.dim);

    }

    TEST(GPUCACHE, FIFOCACHEWRAPPER){
        CacheConfig cfg{CacheConfig::CacheEvictStrategy::FIFO,8192,4,32,2048,0,8};
        auto t = torch::empty({1},torch::dtype(torch::kInt32).device(torch::kCUDA,0));

        auto cache = NewCache(t,cfg);
        auto keys = torch::arange(0,5,torch::dtype(torch::kInt32)).to(torch::kCUDA, 0);
        auto [values, missing_indexs, missing_keys, n_missing] = cache->Get(5,keys);

        torch::Tensor put_values = torch::reshape(torch::arange(0, 5 * 8,torch::dtype(torch::kInt32)),{5,8}).to(torch::kCUDA, 0);
        std::cout << "put_values: \n" << put_values << std::endl;
        cache->Put(5,keys,put_values);

        auto result = cache->Get(5,keys);
        values = std::get<0>(result);
        missing_indexs = std::get<1>(result);
        missing_keys = std::get<2>(result);
        n_missing = std::get<3>(result);
        std::cout << " values: \n" << values << "missing_indexs: " << missing_indexs << std::endl;
        std::cout << "missing_keys\n"<<missing_keys << "n_missing: " << n_missing <<std::endl;

        cache->Clear();

        CUDA_CHECK(cudaDeviceSynchronize());
    }

    TEST(GPUCACHE, LRUCACHEWRAPPER){
        CacheConfig cfg{CacheConfig::CacheEvictStrategy::LRU,8192,4,32,2048,0,8};
        auto t = torch::empty({1},torch::dtype(torch::kInt32).device(torch::kCUDA,0));

        auto cache = NewCache(t,cfg);
        auto keys = torch::arange(0,5,torch::dtype(torch::kInt32)).to(torch::kCUDA, 0);
        auto [values, missing_indexs, missing_keys, n_missing] = cache->Get(5,keys);

        torch::Tensor put_values = torch::reshape(torch::arange(0, 5 * 8,torch::dtype(torch::kInt32)),{5,8}).to(torch::kCUDA, 0);
        std::cout << "put_values: \n" << put_values << std::endl;
        cache->Put(5,keys,put_values);

        auto result = cache->Get(5,keys);
        values = std::get<0>(result);
        missing_indexs = std::get<1>(result);
        missing_keys = std::get<2>(result);
        n_missing = std::get<3>(result);
        std::cout << " values: \n" << values << "missing_indexs: " << missing_indexs << std::endl;
        std::cout << "missing_keys\n"<<missing_keys << "n_missing: " << n_missing <<std::endl;

        cache->Clear();

        CUDA_CHECK(cudaDeviceSynchronize());
    }

    TEST(GPUCACHE, LFUCACHEWRAPPER){
        CacheConfig cfg{CacheConfig::CacheEvictStrategy::LFU,8192,4,32,2048,0,8};
        auto t = torch::empty({1},torch::dtype(torch::kInt32).device(torch::kCUDA,0));

        auto cache = NewCache(t,cfg);
        auto keys = torch::arange(0,5,torch::dtype(torch::kInt32)).to(torch::kCUDA, 0);
        auto [values, missing_indexs, missing_keys, n_missing] = cache->Get(5,keys);

        torch::Tensor put_values = torch::reshape(torch::arange(0, 5 * 8,torch::dtype(torch::kInt32)),{5,8}).to(torch::kCUDA, 0);
        std::cout << "put_values: \n" << put_values << std::endl;
        cache->Put(5,keys,put_values);

        auto result = cache->Get(5,keys);
        values = std::get<0>(result);
        missing_indexs = std::get<1>(result);
        missing_keys = std::get<2>(result);
        n_missing = std::get<3>(result);
        std::cout << " values: \n" << values << "missing_indexs: " << missing_indexs << std::endl;
        std::cout << "missing_keys\n"<<missing_keys << "n_missing: " << n_missing <<std::endl;

        cache->Clear();

        CUDA_CHECK(cudaDeviceSynchronize());
    }
}