import numpy as np
import torch
from torch.profiler import record_function
from typing import List, Optional
from contexttimer import Timer
from .copyer import LimitBuffIndexCopyer
from enum import Enum
import sys
from contextlib import contextmanager


class EvictionStrategy(Enum):
    LFU = 1
    # dataset aware eviction strategy
    DATASET = 2


def _wait_for_data(t, stream: Optional[torch.cuda.streams.Stream]) -> None:
    if stream is None:
        return
    torch.cuda.current_stream().wait_stream(stream)
    # As mentioned in https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html,
    # PyTorch uses the "caching allocator" for memory allocation for tensors. When a tensor is
    # freed, its memory is likely to be reused by newly constructed tensors.  By default,
    # this allocator traces whether a tensor is still in use by only the CUDA stream where it
    # was created.   When a tensor is used by additional CUDA streams, we need to call record_stream
    # to tell the allocator about all these streams.  Otherwise, the allocator might free the
    # underlying memory of the tensor once it is no longer used by the creator stream.  This is
    # a notable programming trick when we write programs using multi CUDA streams.
    cur_stream = torch.cuda.current_stream()
    assert isinstance(t, torch.Tensor)
    t.record_stream(cur_stream)


class CachedParamMgr(torch.nn.Module):
    """
    Manage Embedding Weights on CPU and CUDA memory uses a software cache.
    CPU maintains the entire original weight. 
    CUDA maintains a fraction of the weights used in the upcoming computation. The row number in CUDA is controlled by `cuda_row_num`.
    During training, GPU needs to transmit embedding rows between CPU and GPU.
    Args:
        weight (torch.Tensor): the weight of the Embedding layer.
        cuda_row_num (int, optional): the number of rows cached in CUDA memory. Defaults to 0.
        buffer_size (int, optional): the number of rows in a data transmitter buffer. Defaults to 50_000.
        pin_weight (bool, optional): use pin memory to store the cpu weight. If set `True`, the cpu memory usage will increase largely. Defaults to False.
        evict_strategy (EvictionStrategy, optional): the eviction strategy. There are two options.
        `EvictionStrategy.LFU`: use the least frequently used cache.
        `EvictionStrategy.DATASET`: use the stats collected from the target dataset. It usually leads to less cpu-gpu communication volume.
        Defaults to EvictionStrategy.DATASET.
    """

    def __init__(
        self,
        weight: torch.Tensor,
        cuda_row_num: int = 0,
        buffer_size: int = 0,
        pin_weight: bool = True,
        evict_strategy: EvictionStrategy = EvictionStrategy.DATASET,
        async_copy: bool = False,
    ) -> None:
        super(CachedParamMgr, self).__init__()
        self.buffer_size = buffer_size
        self.num_embeddings, self.embedding_dim = weight.shape
        self.cuda_row_num = cuda_row_num
        self._cuda_available_row_num = self.cuda_row_num
        self.pin_weight = pin_weight
        self.elem_size_in_byte = weight.element_size()

        # weight configure
        self._init_weight(weight)

        # Perf log
        self.num_hits_history = []
        self.num_miss_history = []
        self.num_write_back_history = []

        self._evict_strategy = evict_strategy

        self._async_copy = async_copy

        if self._async_copy:
            self._memcpy_stream = torch.cuda.Stream()

            print('use async copy')

        if self._evict_strategy == EvictionStrategy.LFU:
            # cache_row_idx -> frequency, freq of the cache rows.
            # classic lfu cache. evict the minimal freq value row in cuda cache.
            self.register_buffer("freq_cnter",
                                 torch.empty(self.cuda_row_num, device=torch.cuda.current_device(),
                                             dtype=torch.long).fill_(sys.maxsize),
                                 persistent=False)
        self._elapsed_dict = {}
        self._show_cache_miss = True
        self._reset_comm_stats()

    def _reset_comm_stats(self):
        for k in self._elapsed_dict.keys():
            self._elapsed_dict[k] = 0

        self._cpu_to_cuda_numel = 0
        self._cuda_to_cpu_numel = 0
        if self._show_cache_miss:
            self._cache_miss = 0
            self._total_cache = 0

    @contextmanager
    def timer(self, name):
        with Timer() as t:
            yield
            torch.cuda.synchronize()

        if name not in self._elapsed_dict.keys():
            self._elapsed_dict[name] = 0
        self._elapsed_dict[name] += t.elapsed

    def _find_evict_gpu_idxs(self, evict_num: int) -> torch.Tensor:
        """_find_evict_gpu_idxs 
        Find the gpu idxs to be evicted, according to their freq.
        Args:
            evict_num (int): how many rows has to be evicted
        Returns:
            torch.Tensor: a list tensor (1D), contains the gpu_row_idxs.
        """
        if self._evict_strategy == EvictionStrategy.LFU:
            # find the minimal evict_num freq entries in cached_idx_map
            _, evict_gpu_row_idxs = torch.topk(self.freq_cnter, evict_num, largest=False)
            return evict_gpu_row_idxs
        elif self._evict_strategy == EvictionStrategy.DATASET:
            # cached_idx_map itself implies the priority of eviction.
            # The value of self.cached_idx_map represents cpu_row_idx.
            # The larger it is, the less frequently it will appear in the dataset,
            # and the higher its eviction priority will be.
            _, evict_gpu_row_idxs = torch.topk(self.cached_idx_map, evict_num, largest=True)
            return evict_gpu_row_idxs
        else:
            raise TypeError

    def _init_weight(self, weight):
        if self.cuda_row_num > 0:
            # Enable cache with introducing auxiliary data structures
            self.cuda_cached_weight = torch.nn.Parameter(
                torch.zeros(self.cuda_row_num,
                            self.embedding_dim,
                            device=torch.cuda.current_device(),
                            dtype=weight.dtype))

            # pin memory cpu for higher CPU-GPU copy bandwidth
            self.weight = weight.pin_memory() if self.pin_weight else weight
            # map original id to new id with respect to frequency
            # id -> cpu_row_idx
            self.register_buffer(
                "idx_map",
                torch.arange(self.num_embeddings, dtype=torch.long, device=torch.cuda.current_device()),
                persistent=False,
            )

            # cached_idx_map: gpu_row_idx -> cpu_row_idx
            self.register_buffer("cached_idx_map",
                                 torch.empty(self.cuda_row_num, device=torch.cuda.current_device(),
                                             dtype=torch.long).fill_(-1),
                                 persistent=False)

            # cpu_row_id -> gpu_row_idx.
            # gpu_row_idx as -1 means cpu_row_id not in CUDA.
            self.register_buffer("inverted_cached_idx",
                                 torch.zeros(self.num_embeddings, device=torch.cuda.current_device(),
                                             dtype=torch.long).fill_(-1),
                                 persistent=False)

            self.evict_backlist = torch.tensor([], device=torch.cuda.current_device())

            # index copy buffer size should less than 10% of cuda weight.
            if self.buffer_size > 0:
                self.limit_buff_index_copyer = LimitBuffIndexCopyer(self.buffer_size)

        else:
            # Disable cache so that FreqCacheEmbedding is compatible with vanilla EmbeddingBag
            # self.weight = torch.nn.Parameter(weight)
            # self.cuda_cached_weight = self.weight
            raise NotImplementedError()

    def cpu_weight_data(self, row_idx: int) -> torch.Tensor:
        """
        access a row of CPU weight.
        Args:
            row_idx (int): the idx of rows
        Returns:
            torch.Tensor: a piece of memory in CPU weight corresponding to row id's payload. The tensor is 1-D.
        """

        return self.weight.data.view(-1).narrow(0,
                                                int(row_idx) * self.embedding_dim,
                                                self.embedding_dim).view(1, self.embedding_dim)

    @property
    def cuda_available_row_num(self):
        return self._cuda_available_row_num

    @torch.no_grad()
    def reorder(self, ids_freq_mapping: Optional[List[int]] = None, warmup_ratio=0.7):
        """reorder
        reorder the weight according to ids' frequency in dataset before training.
        Execute only once before training, also known as warmup phase.
        
        Note:
            If you would like to use the DATASET as the eviction strategy, you must call this function.
        Note:
            If you are use the LFU as the eviction strategy, you can skip this function. If you still use this function. It will initialize
            The frequency in LFU cache using the dataset statistics.
        Args:
            ids_freq_mapping (List[int]): a list, whose offset is id number, value is freq. if None then not reorder the cpu weight.
            warmup_ratio (float): the amount of chunks preloaded in cuda cache
        """
        # reorder phase: reorder the cpu weight according to their freq stats in the target dataset.
        # reorder only works for DATASET eviction strategy.

        if ids_freq_mapping is not None and not isinstance(ids_freq_mapping, torch.Tensor):
            ids_freq_mapping = torch.tensor(ids_freq_mapping)

        if self._evict_strategy == EvictionStrategy.DATASET:
            if ids_freq_mapping is not None:
                tmp_idx = torch.argsort(ids_freq_mapping, descending=True)
                sorted_idx = torch.argsort(tmp_idx)
                self.idx_map.data.copy_(sorted_idx)

        # warmup phase: copy #preload_row_num rows from cpu to gpu.
        preload_row_num = min(int(np.ceil(self.cuda_row_num * warmup_ratio)), self.num_embeddings)
        if preload_row_num > 0:
            with Timer() as timer:
                # extract rows from cpu weight
                if self._evict_strategy == EvictionStrategy.LFU and ids_freq_mapping is not None:
                    freq_value, preload_cpu_ids = torch.topk(ids_freq_mapping, preload_row_num, dim=0, largest=True)
                    preload_cuda_row_idxs = torch.arange(preload_row_num).cuda()
                else:
                    preload_cpu_ids = torch.arange(preload_row_num)
                    preload_cuda_row_idxs = preload_cpu_ids.cuda()
                if self.buffer_size > 0:
                    self.limit_buff_index_copyer.index_copy(0,
                                                            src_index=preload_cpu_ids,
                                                            tgt_index=preload_cuda_row_idxs,
                                                            src=self.weight.view(self.num_embeddings, -1),
                                                            tgt=self.cuda_cached_weight.view(self.cuda_row_num, -1))
                else:
                    preload_rows = self.weight.view(self.num_embeddings, -1).index_select(0, preload_cpu_ids).cuda()
                    self.cuda_cached_weight.view(self.cuda_row_num, -1).index_copy_(0, preload_cuda_row_idxs,
                                                                                    preload_rows)

                # update auxiliary info
                self.cached_idx_map[preload_cuda_row_idxs] = preload_cpu_ids.cuda()
                self.inverted_cached_idx[preload_cpu_ids] = preload_cuda_row_idxs
                self._cuda_available_row_num -= preload_row_num

                if self._evict_strategy == EvictionStrategy.LFU:
                    # if the ids_freq_mapping is not None, we initialize the embedding row's freq value in LFU as its freq in dataset.
                    if ids_freq_mapping is None:
                        self.freq_cnter.index_fill_(0, preload_cuda_row_idxs, 0)
                    else:
                        self.freq_cnter[preload_cuda_row_idxs] = freq_value.cuda()

            print(f'Cache warmup finished cost {timer.elapsed} sec.')

    def flush(self):
        """flush all CUDA rows to CPU.
        The function is usually called after training finished.
        """
        slots = torch.nonzero(self.cached_idx_map > -1).squeeze(1)
        row_ids = self.cached_idx_map[slots]
        rows = self.cuda_cached_weight.view(self.cuda_row_num, -1).index_select(0, slots).cpu()
        self.weight.view(self.num_embeddings, -1).index_copy_(0, row_ids.cpu(), rows)
        self.cached_idx_map.index_fill_(0, slots, -1)
        self.inverted_cached_idx.index_fill_(0, row_ids, -1)
        self._cuda_available_row_num += slots.numel()

        if self._show_cache_miss:
            self._cache_miss = 0
            self._total_cache = 0

        if self._evict_strategy == EvictionStrategy.LFU:
            self.freq_cnter.fill_(sys.maxsize)
        assert self._cuda_available_row_num == self.cuda_row_num
        assert torch.all(self.inverted_cached_idx == -1).item()
        assert torch.all(self.cached_idx_map == -1).item()

    def print_comm_stats(self):
        if self._cuda_to_cpu_numel > 0 and "3_evict_out" in self._elapsed_dict:
            elapsed = self._elapsed_dict["3_evict_out"]
            print(
                f"CUDA->CPU BWD {self._cuda_to_cpu_numel * self.elem_size_in_byte / 1e6 / elapsed} MB/s {self._cuda_to_cpu_numel / 1e6} M elem"
            )
            print(f'cuda_to_cpu_elapse {elapsed} sec')
        if self._cpu_to_cuda_numel > 0 and "5_evict_in" in self._elapsed_dict:
            elapsed = self._elapsed_dict["5_evict_in"]
            print(
                f"CPU->CUDA BWD {self._cpu_to_cuda_numel * self.elem_size_in_byte / 1e6 / elapsed} MB/s {self._cpu_to_cuda_numel / 1e6} M elem"
            )
            print(f'cpu_to_cuda_elapse {elapsed} sec')

        for k, v in self._elapsed_dict.items():
            print(f'{k}: {v}')

        print(f'cache miss ratio {self._cache_miss / self._total_cache}')

    @torch.no_grad()
    def _id_to_cached_cuda_id(self, ids: torch.Tensor) -> torch.Tensor:
        """
        convert ids to indices in self.cuda_cached_weight.
        Implemented with parallel operations on GPU.
        Args:
            ids (torch.Tensor): ids from the dataset
        Returns:
            torch.Tensor: contains indices in self.cuda_cached_weight
        """
        ids = self.idx_map.index_select(0, ids.view(-1))
        ret = self.inverted_cached_idx.index_select(0, ids)
        return ret

    @torch.no_grad()
    def prepare_ids(self, ids: torch.Tensor) -> torch.Tensor:
        """
        move the cpu embedding rows w.r.t. ids into CUDA memory
        Args:
            ids (torch.Tensor): the ids to be computed
        Returns:
            torch.Tensor: indices on the cuda_cached_weight.
        """
        torch.cuda.synchronize()
        with self.timer("cache_op") as gtimer:
            # identify cpu rows to cache
            with self.timer("1_identify_cpu_row_idxs") as timer:
                with record_function("(cache) get unique indices"):
                    if self._evict_strategy == EvictionStrategy.LFU:
                        cpu_row_idxs, repeat_times = torch.unique(ids, return_counts=True)
                    else:
                        cpu_row_idxs, repeat_times = torch.unique(self.idx_map.index_select(0, ids), return_counts=True)

                    assert len(cpu_row_idxs) <= self.cuda_row_num, \
                        f"You move {len(cpu_row_idxs)} embedding rows from CPU to CUDA. " \
                        f"It is larger than the capacity of the cache, which at most contains {self.cuda_row_num} rows, " \
                        f"Please increase cuda_row_num or decrease the training batch size."
                    self.evict_backlist = cpu_row_idxs
                    tmp = torch.isin(cpu_row_idxs, self.cached_idx_map, invert=True)
                    comm_cpu_row_idxs = cpu_row_idxs[tmp]

                    if self._show_cache_miss:
                        self._cache_miss += torch.sum(repeat_times[tmp])
                        self._total_cache += ids.numel()

            self.num_hits_history.append(len(cpu_row_idxs) - len(comm_cpu_row_idxs))
            self.num_miss_history.append(len(comm_cpu_row_idxs))
            self.num_write_back_history.append(0)

            # move sure the cuda rows will not be evicted!
            with record_function("(cache) prepare_rows_on_cuda"):
                with self.timer("prepare_rows_on_cuda") as timer:
                    self._prepare_rows_on_cuda(comm_cpu_row_idxs)

            self.evict_backlist = torch.tensor([], device=cpu_row_idxs.device, dtype=cpu_row_idxs.dtype)

            with self.timer("6_update_cache") as timer:
                with record_function("6_update_cache"):
                    gpu_row_idxs = self._id_to_cached_cuda_id(ids)

                # update for LFU.
                if self._evict_strategy == EvictionStrategy.LFU:
                    unique_gpu_row_idxs = self.inverted_cached_idx[cpu_row_idxs]
                    self.freq_cnter.scatter_add_(0, unique_gpu_row_idxs, repeat_times)

        return gpu_row_idxs

    def _row_in_cuda(self, row_id: int) -> bool:
        return self.inverted_cached_idx[row_id] != -1

    @torch.no_grad()
    def _prepare_rows_on_cuda(self, cpu_row_idxs: torch.Tensor) -> None:
        """prepare rows in cpu_row_idxs on CUDA memory
        Args:
            cpu_row_idxs (torch.Tensor): the rows to be placed on CUDA
        """
        evict_num = cpu_row_idxs.numel() - self.cuda_available_row_num

        cpu_row_idxs_copy = cpu_row_idxs.cpu()

        # move evict in rows to gpu
        if self._async_copy:
            if self.buffer_size == 0:
                evict_in_rows_gpu = self.weight.view(self.num_embeddings,
                                                     -1).index_select(0, cpu_row_idxs_copy).pin_memory()
                with torch.cuda.stream(self._memcpy_stream):
                    evict_in_rows_gpu = evict_in_rows_gpu.to(torch.cuda.current_device(), non_blocking=True)
            else:
                raise NotImplemented

        if evict_num > 0:
            with self.timer("2_identify_cuda_row_idxs") as timer:
                mask_cpu_row_idx = torch.isin(self.cached_idx_map, self.evict_backlist)
                invalid_idxs = torch.nonzero(mask_cpu_row_idx).squeeze(1)
                if self._evict_strategy == EvictionStrategy.DATASET:
                    # mask method.
                    # set cached_idx_map[invalid_idxs] to -2.
                    # so those idxs will be sorted to end, therefore not being chosen as victim
                    backup_idxs = self.cached_idx_map[mask_cpu_row_idx].clone()
                    self.cached_idx_map.index_fill_(0, invalid_idxs, -2)

                    with self.timer("2_1_find_evict_gpu_idxs") as timer:
                        evict_gpu_row_idxs = self._find_evict_gpu_idxs(evict_num)

                    # move evict out rows to cpu
                    if self._async_copy:
                        evict_out_rows_gpu = self.cuda_cached_weight.view(self.cuda_row_num,
                                                                          -1).index_select(0, evict_gpu_row_idxs)
                        evict_out_rows_cpu = torch.empty_like(evict_out_rows_gpu, device='cpu', pin_memory=True)
                        with torch.cuda.stream(None):
                            evict_out_rows_cpu.copy_(evict_out_rows_gpu, non_blocking=True)
                    self.cached_idx_map.index_copy_(0, invalid_idxs, backup_idxs)

                elif self._evict_strategy == EvictionStrategy.LFU:
                    with self.timer("2_1_backup_freqs") as timer:
                        backup_freqs = self.freq_cnter[invalid_idxs].clone()
                        self.freq_cnter.index_fill_(0, invalid_idxs, sys.maxsize)

                    with self.timer("2_2_find_evict_gpu_idxs") as timer:
                        evict_gpu_row_idxs = self._find_evict_gpu_idxs(evict_num)

                    if self._async_copy:
                        evict_out_rows_gpu = self.cuda_cached_weight.view(self.cuda_row_num,
                                                                          -1).index_select(0, evict_gpu_row_idxs)
                        evict_out_rows_cpu = torch.empty_like(evict_out_rows_gpu, device='cpu', pin_memory=True)
                        with torch.cuda.stream(None):
                            evict_out_rows_cpu.copy_(evict_out_rows_gpu, non_blocking=True)

                    with self.timer("2_3_revert_freqs") as timer:
                        self.freq_cnter.index_copy_(0, invalid_idxs, backup_freqs)

                evict_info = self.cached_idx_map[evict_gpu_row_idxs]

            with self.timer("3_evict_out") as timer:
                if self.buffer_size > 0:
                    self.limit_buff_index_copyer.index_copy(0,
                                                            src_index=evict_gpu_row_idxs,
                                                            tgt_index=evict_info.cpu(),
                                                            src=self.cuda_cached_weight.view(self.cuda_row_num, -1),
                                                            tgt=self.weight.view(self.num_embeddings, -1))
                else:
                    # allocate tmp memory on CPU and copy rows on CUDA to CPU.
                    # TODO async gpu -> cpu
                    if self._async_copy:
                        _wait_for_data(evict_out_rows_cpu, None)
                    else:
                        with self.timer("3_1_evict_out_index_select") as timer:
                            evict_out_rows_cpu = self.cuda_cached_weight.view(self.cuda_row_num,
                                                                              -1).index_select(0, evict_gpu_row_idxs)
                        with self.timer("3_2_evict_out_gpu_to_cpu_copy") as timer:
                            evict_out_rows_cpu = evict_out_rows_cpu.cpu()

                    with self.timer("3_2_evict_out_cpu_copy") as timer:
                        self.weight.view(self.num_embeddings, -1).index_copy_(0, evict_info.cpu(), evict_out_rows_cpu)

                self.cached_idx_map.index_fill_(0, evict_gpu_row_idxs, -1)
                self.inverted_cached_idx.index_fill_(0, evict_info, -1)
                # self.freq_cnter.index_fill(0, evict_gpu_row_idxs, sys.maxsize) # unnecessary
                self._cuda_available_row_num += evict_num

                weight_size = evict_gpu_row_idxs.numel() * self.embedding_dim
                self._cuda_to_cpu_numel += weight_size
            # print(f"evict embedding weight: {weight_size*self.elem_size_in_byte/1e6:.2f} MB")

        # slots of cuda weight to evict in
        with self.timer("4_identify_cuda_slot") as timer:
            slots = torch.nonzero(self.cached_idx_map == -1).squeeze(1)[:cpu_row_idxs.numel()]

        # TODO wait for optimize
        with self.timer("5_evict_in") as timer:
            # Here also allocate extra memory on CUDA. #cpu_row_idxs
            if self.buffer_size > 0:
                self.limit_buff_index_copyer.index_copy(0,
                                                        src_index=cpu_row_idxs_copy,
                                                        tgt_index=slots,
                                                        src=self.weight.view(self.num_embeddings, -1),
                                                        tgt=self.cuda_cached_weight.view(self.cuda_row_num, -1))
            else:
                if self._async_copy:
                    _wait_for_data(evict_in_rows_gpu, self._memcpy_stream)
                else:
                    with self.timer("5_1_evict_in_index_select") as timer:
                        # narrow index select to a subset of self.weight
                        # tmp = torch.narrow(self.weight.view(self.num_embeddings, -1), 0, min(cpu_row_idxs).cpu(), max(cpu_row_idxs) - min(cpu_row_idxs) + 1)
                        # evict_in_rows_gpu = tmp.index_select(0, cpu_row_idxs_copy - min(cpu_row_idxs).cpu())
                        evict_in_rows_gpu = self.weight.view(self.num_embeddings,
                                                             -1).index_select(0, cpu_row_idxs_copy).pin_memory()

                    with self.timer("5_2_evict_in_gpu_to_cpu_copy") as timer:
                        evict_in_rows_gpu = evict_in_rows_gpu.cuda()

                    with self.timer("5_3_evict_in_index_copy") as timer:
                        self.cuda_cached_weight.view(self.cuda_row_num, -1).index_copy_(0, slots, evict_in_rows_gpu)

        with self.timer("6_update_cache") as timer:
            self.cached_idx_map[slots] = cpu_row_idxs
            self.inverted_cached_idx.index_copy_(0, cpu_row_idxs, slots)
            if self._evict_strategy == EvictionStrategy.LFU:
                self.freq_cnter.index_fill_(0, slots, 0)
            self._cuda_available_row_num -= cpu_row_idxs.numel()

        weight_size = cpu_row_idxs.numel() * self.embedding_dim
        self._cpu_to_cuda_numel += weight_size
        # print(f"admit embedding weight: {weight_size*self.elem_size_in_byte/1e6:.2f} MB")

    def _find_free_cuda_row(self) -> int:
        if self._cuda_available_row_num == 0:
            return -1
        candidates = torch.nonzero(self.cached_idx_map == -1).squeeze(1)
        return candidates[0].item()

    def _evict(self) -> int:
        """
        deprecated
        evict one row from cuda to cpu.
        Returns: 
        (int) : the slot id be evicted.
        """
        mask = torch.logical_or(torch.isin(self.cached_idx_map, self.evict_backlist), self.cached_idx_map == -1)
        buf = self.cached_idx_map[mask].clone()
        idx = torch.nonzero(mask).squeeze(1)
        self.cached_idx_map.index_fill_(0, idx, -1)
        max_row, max_cpu_row_idx = torch.max(self.cached_idx_map, dim=0)
        max_gpu_row_idx = self.cached_idx_map[max_cpu_row_idx]

        if max_gpu_row_idx == -1:
            raise RuntimeError("Can not evict a row")

        max_gpu_row_idx = max_gpu_row_idx.item()
        max_offset = self.inverted_cached_idx[max_gpu_row_idx]
        # recover
        self.cached_idx_map.index_copy_(0, idx, buf)

        with Timer() as timer:
            cuda_tensor = torch.narrow(self.cuda_cached_weight.view(-1), 0, max_offset * self.embedding_dim,
                                       self.embedding_dim).view(1, self.embedding_dim)
            self.cpu_weight_data(max_gpu_row_idx).data.copy_(cuda_tensor)

        # update inverted_cached_idx, min_slot_id is evicted from cuda
        self.cached_idx_map[max_cpu_row_idx] = -1
        if self._evict_strategy == EvictionStrategy.LFU:
            self.freq_cnter[max_cpu_row_idx] = sys.maxsize
        self.inverted_cached_idx[max_gpu_row_idx] = -1

        self._cuda_available_row_num += 1

        self._cuda_to_cpu_numel += self.embedding_dim
        # self.num_write_back_history[-1] += 1
        return max_cpu_row_idx

    @torch.no_grad()
    def _admit(self, row_id: int):
        """
        deprecated
        move in row_id to CUDA
        Args:
            row_id (int): the id of row to be moved in
        """
        # find a free slot in partial cuda weight
        slot_id = self._find_free_cuda_row()

        if slot_id == -1:
            # evict one row
            slot_id = self._evict()
        slot_offset = slot_id
        # copy payload from cpu to cuda
        with Timer() as timer:
            cuda_tensor = torch.narrow(self.cuda_cached_weight.view(-1), 0, slot_offset * self.embedding_dim,
                                       self.embedding_dim).view(1, self.embedding_dim)
            cuda_tensor.data.copy_(self.cpu_weight_data(row_id))

        # update the inverted_cached_idx
        self.cached_idx_map[slot_id] = row_id
        if self._evict_strategy == EvictionStrategy.LFU:
            self.freq_cnter[slot_id] = 0
        self.inverted_cached_idx[row_id] = slot_offset

        self._cuda_available_row_num -= 1

        self._cpu_to_cuda_numel += self.embedding_dim
