#include "gtest/gtest.h"
#include "../src/lru_cache.cu"
//#include "cache.h"
#include <cuda_runtime.h>
#include <unordered_set>
#include <vector>
#include <numeric>
#include <random>
#include <algorithm>

namespace gpucache{
    void TestCache(Cache<uint32_t ,uint32_t>& cache){

        std::unordered_set<uint32_t> in_cache;
        const size_t n_iter = 32;
        const uint32_t n_keys = 1024;

        uint32_t* d_keys;
        uint32_t* keys;
        uint32_t* d_values;
        uint32_t* values;
        uint32_t* d_n_missing;
        uint32_t* n_missing;
        uint32_t* d_missing_keys;
        uint32_t* missing_keys;
        uint32_t* d_n_evict;
        uint32_t* n_evict;
        uint32_t* d_evict_keys;
        uint32_t* evict_keys;
        bool* find_mask;
        bool* d_find_mask;

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


        CHECK(cudaMalloc(&d_keys,keys_size));
        CHECK(cudaMalloc(&d_values, values_size));
        CHECK(cudaMalloc(&d_n_missing,sizeof(uint32_t)));
        CHECK(cudaMalloc(&d_find_mask, mask_size));
        CHECK(cudaMalloc(&d_missing_keys,keys_size));
        CHECK(cudaMalloc(&d_evict_keys,keys_size));
        CHECK(cudaMalloc(&d_n_evict,sizeof(uint32_t)));
        CHECK(cudaMallocHost(&keys,keys_size));
        CHECK(cudaMallocHost(&values,values_size));
        CHECK(cudaMallocHost(&n_missing,sizeof(uint32_t)));
        CHECK(cudaMallocHost(&find_mask,mask_size));
        CHECK(cudaMallocHost(&missing_keys,mask_size));
        CHECK(cudaMallocHost(&evict_keys,keys_size));
        CHECK(cudaMallocHost(&n_evict,sizeof(uint32_t)));



        std::vector<uint32_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::shuffle(random_keys.begin(), random_keys.end(),g);
            std::copy(random_keys.begin(),random_keys.end()+n_keys,keys);
            uint32_t expect_n_missing = 0;
            std::unordered_set<uint32_t> expect_missing_keys_set;
            std::unordered_set<uint32_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
            CHECK(cudaMemcpy(d_keys,keys,keys_size,cudaMemcpyHostToDevice));
            cudaStream_t* stream;
            CHECK(cudaStreamCreate(stream));
            CHECK(cudaDeviceSynchronize());
            cache.Get(stream,n_keys,d_keys,d_values,d_find_mask);
            CHECK(cudaDeviceSynchronize());

            dim3 block(256);
            dim3 grid((n_keys + block.x - 1)/block.x);
            CollectMissingKeys<uint32_t,uint32_t><<<grid,block>>>(n_keys,d_keys,find_mask,d_n_missing,d_missing_keys);

            CHECK(cudaMemcpy(n_missing,d_n_missing, sizeof(uint32_t),cudaMemcpyDefault));
            CHECK(cudaMemcpy(missing_keys,d_missing_keys,keys_size,cudaMemcpyDefault));
            ASSERT_EQ(expect_n_missing, *n_missing);
            std::unordered_set<uint32_t> missing_keys_set(missing_keys, missing_keys + *n_missing);
            ASSERT_EQ(missing_keys_set,expect_missing_keys_set);

            // test put
            for (size_t i = 0; i < n_keys; i++){
                values[i] = keys[i] + 123;
            }
            CHECK(cudaMemcpy(d_values,values,values_size,cudaMemcpyDefault));
            cache.Put(stream,n_keys,d_keys,d_values,d_n_evict,d_evict_keys);
            CHECK(cudaMemcpy(n_evict,d_n_evict,sizeof(uint32_t),cudaMemcpyDefault));
            CHECK(cudaMemcpy(evict_keys,d_evict_keys,*n_evict * sizeof(uint32_t), cudaMemcpyDefault));
            std::unordered_set<uint32_t> evict_keys_set(evict_keys, evict_keys + *n_evict);
            CHECK(cudaDeviceSynchronize());
            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);
            }
            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]);}

            CHECK(cudaFree(&d_keys));
            CHECK(cudaFree(&d_values));
            CHECK(cudaFree(&d_n_missing));
            CHECK(cudaFree(&d_find_mask));
            CHECK(cudaFree(&d_missing_keys));
            CHECK(cudaFree(&d_evict_keys));
            CHECK(cudaFree(&d_n_evict));
            CHECK(cudaFreeHost(&keys));
            CHECK(cudaFreeHost(&values));
            CHECK(cudaFreeHost(&n_missing));
            CHECK(cudaFreeHost(&find_mask));
            CHECK(cudaFreeHost(&missing_keys));
            CHECK(cudaFreeHost(&evict_keys));
            CHECK(cudaFreeHost(&n_evict));


        }

    }

    TEST(GPUCACHE,LRUCACHE){
        CacheConfig cfg{};
        cfg.strategy = CacheConfig::CacheEvictStrategy::LRU;
        cfg.valueSize = 1;
        cfg.capacity = 65536;
        cfg.keySize = 1;

        LRUCache<uint32_t,uint32_t> cache(cfg);
        TestCache(cache);
    }
}

