#include "cache_wrapper.cuh"
#include "fifo_cache.cuh"
#include "lru_cache.cuh"
#include "lfu_cache.cuh"

// CacheWrapper API
namespace gpucache {

        template<typename ElemType>
        void NewInt32Cache(void **cache, CacheConfig cfg) {
            void* c;
            switch (cfg.strategy) {
                case CacheConfig::CacheEvictStrategy::LRU:
                    c = new LRUCache<int32_t, ElemType>(cfg);
                    break;
                case CacheConfig::CacheEvictStrategy::LFU:
                    c = new LFUCache<int32_t, ElemType>(cfg);
                    break;
                case CacheConfig::CacheEvictStrategy::FIFO:
                    c = new FIFOCache<int32_t, ElemType>(cfg);
                    break;
            }

            //std::cout << "NewInt32Cache Cache keysize " << c->keySize << std::endl;
            *cache = reinterpret_cast<void *>(c);
        }

        template<typename ElemType>
        void NewInt64Cache(void **cache, CacheConfig cfg) {
            void* c;
            switch (cfg.strategy) {
                case CacheConfig::CacheEvictStrategy::LRU:
                    c = new LRUCache<int64_t, ElemType>(cfg);
                    break;
                case CacheConfig::CacheEvictStrategy::LFU:
                    c = new LFUCache<int64_t, ElemType>(cfg);
                    break;
                case CacheConfig::CacheEvictStrategy::FIFO:
                    c = new FIFOCache<int64_t, ElemType>(cfg);
                    break;
            }
            *cache = reinterpret_cast<void *>(c);
        }

        CacheWrapper::CacheWrapper(torch::Tensor t, CacheConfig cfg) {
            //std::cout << "CacheWrapper constructor" << std::endl;
            vtype = t.scalar_type();
            cache_cfg = cfg;
            key_is_int32 = cfg.capacity < INT32_MAX;
            if (key_is_int32) {
                ktype = torch::kInt32;
                AT_DISPATCH_ALL_TYPES(t.scalar_type(), "int32_cache", [&] {
                    NewInt32Cache<scalar_t>(&cache, cfg);
                });
            } else {
                ktype = torch::kInt64;
                AT_DISPATCH_ALL_TYPES(t.scalar_type(), "int64_cache", [&] {
                    NewInt64Cache<scalar_t>(&cache, cfg);
                });
            }
//        AT_DISPATCH_ALL_TYPES(dtype,"cache_test",[&]{
//            auto c = reinterpret_cast<Cache<int32_t,scalar_t>*>(cache);
//            std::cout << "CacheWrapper constructor Cache dim: " << c->Dim() << " keySize: " << c->KeySize() << "\n";
//        });

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



        CacheWrapper::~CacheWrapper() {
            if (key_is_int32) {
                AT_DISPATCH_ALL_TYPES(vtype, "int32_destructor", [&] {
                    auto c = reinterpret_cast<Cache<int32_t, scalar_t> *>(cache);
                    delete c;
                });
            } else {
                AT_DISPATCH_ALL_TYPES(vtype, "int64_destructor", [&] {
                    auto c = reinterpret_cast<Cache<int64_t, scalar_t> *>(cache);
                    delete c;
                });
            }
            //std::cout << "~LRUCacheWrapper()" << std::endl;
        }


        std::tuple<torch::Tensor,torch::Tensor,torch::Tensor, torch::Tensor> CacheWrapper::Get(uint32_t num_query, const torch::Tensor queries) {
            using namespace torch::indexing;
            assert(queries.device().index() == cache_cfg.device_id &&
                   "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.max_query_num && "num_query must less than max_query_num");
            assert(queries.dtype().toScalarType() == ktype && "key type doesn't match");
            //std::cout << "Get queries: " << queries << std::endl;
            //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.device_id).stream();
            auto result = torch::empty({num_query, cache_cfg.dim},
                                       torch::dtype(vtype).device(torch::kCUDA, cache_cfg.device_id));
            auto missing_indexs = torch::empty({num_query},torch::dtype(c10::ScalarType::Int).device(torch::kCUDA,cache_cfg.device_id));
            auto missing_keys = torch::empty({num_query},torch::dtype(ktype).device(torch::kCUDA,cache_cfg.device_id));
            auto n_missing = torch::zeros({1},torch::dtype(ktype).device(torch::kCUDA,cache_cfg.device_id));

//            auto find_mask = torch::zeros({num_query},
//                                          torch::dtype(torch::kBool).device(torch::kCUDA, cache_cfg.device_id));

            if (key_is_int32) {
                AT_DISPATCH_ALL_TYPES(vtype, "int32_get", [&] {
                    Cache<int32_t, scalar_t> *c = reinterpret_cast<Cache<int32_t, scalar_t> *>(cache);
//                cache->Test();
                    //std::cout <<  "keysize: " << cache->KeySize() << "\n";
                    //std::cout << "dim: " << cache->Dim() << "\n";
                    c->Get(stream, num_query, reinterpret_cast<int32_t *>(keys),
                               reinterpret_cast<scalar_t *>(result.data_ptr()),
                               reinterpret_cast<uint32_t *>(missing_indexs.data_ptr()),reinterpret_cast<int32_t*>(missing_keys.data_ptr()),
                               reinterpret_cast<uint32_t*>(n_missing.data_ptr()));
                });
            } else {
                //std::cout << "int64_lru_get" << "\n";
                AT_DISPATCH_ALL_TYPES(vtype, "int64_lru_get", [&] {
                    Cache<int64_t, scalar_t> *c = reinterpret_cast<Cache<int64_t, scalar_t> *>(cache);
                    //std::cout << "cast success!" << std::endl;
                    c->Get(stream, num_query, reinterpret_cast<int64_t *>(keys),
                               reinterpret_cast<scalar_t *>(result.data_ptr()),
                               reinterpret_cast<uint32_t *>(missing_indexs.data_ptr()), reinterpret_cast<int64_t *>(missing_keys.data_ptr()),
                               reinterpret_cast<uint32_t *>(n_missing.data_ptr()));
                });
            }
            // std::cout << "find_mask:" << find_mask << std::endl;
            return std::make_tuple(result, missing_indexs.index({Slice(0,n_missing.index({0}).item<int64_t >(),None)}), missing_keys.index({Slice(0,n_missing.index({0}).item<int64_t >(),None)}),n_missing);
        }

        void CacheWrapper::Put(uint32_t num_query, const torch::Tensor keys, const torch::Tensor values) {
            assert(keys.device().index() == cache_cfg.device_id && cache_cfg.device_id == 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.max_query_num && "num_query must less than max_query_num");
            assert(values.dtype() == vtype);
            auto stream = at::cuda::getCurrentCUDAStream(cache_cfg.device_id).stream();
            if (key_is_int32) {
                //std::cout << "int32_put" << "\n";
                AT_DISPATCH_ALL_TYPES(vtype, "int32_put", [&] {
                    auto c = reinterpret_cast<Cache<int32_t, scalar_t> *>(cache);
                    //std::cout << "put_dim" << cache->Dim() << std::endl;
                    c->Put(stream, num_query, reinterpret_cast<int32_t *>(keys.data_ptr()),
                               reinterpret_cast<scalar_t *>(values.data_ptr()));
                });
            } else {
                //std::cout << "int64_put" << "\n";
                AT_DISPATCH_ALL_TYPES(vtype, "int64_put", [&] {
                    auto c = reinterpret_cast<Cache<int64_t, scalar_t> *>(cache);
                    c->Put(stream, num_query, reinterpret_cast<int64_t *>(keys.data_ptr()),
                               reinterpret_cast<scalar_t *>(values.data_ptr()));
                });
            }
        }

        CacheConfig::CacheEvictStrategy CacheWrapper::Strategy() {
            return cache_cfg.strategy;
        }

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

        uint32_t CacheWrapper::KeySize() { return cache_cfg.key_size; }

        uint32_t CacheWrapper::ValueSize() { return cache_cfg.value_size; }

        uint32_t CacheWrapper::MaxQueryNum() { return cache_cfg.max_query_num; }

        uint64_t CacheWrapper::DeviceId() { return cache_cfg.device_id; }

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

        void CacheWrapper::Clear() {
            if (key_is_int32) {
                AT_DISPATCH_ALL_TYPES(vtype, "int32_clear", [&] {
                    auto c = reinterpret_cast<Cache<int32_t, scalar_t> *>(cache);
                    c->Clear();
                });
            } else {
                AT_DISPATCH_ALL_TYPES(vtype, "int64_clear", [&] {
                    auto c = reinterpret_cast<Cache<int64_t, scalar_t> *>(cache);
                    c->Clear();
                });
            }
        }

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

} // namespace gpucache