import os
import torch
import time
import numpy as np
import torch.distributed as dist
from transformers.utils import logging
from transformers import AutoTokenizer
from itertools import cycle
from typing import List

logger = logging.get_logger(__name__)


class Memory(torch.nn.Module):
    def __init__(
        self, 
        model_config, 
        k_seq_dim:int=2, 
        v_seq_dim:int=2, 
    ):
        """Setup necessary attributes."""
        super().__init__()

        self.config = model_config

        # initialize necessary parameters
        self.k_seq_dim = k_seq_dim
        self.v_seq_dim = v_seq_dim
        self.rng = np.random.default_rng(42)

        self.beacon_token = torch.tensor([self.config.vocab_size])

        self._post_validation()
        self.reset()

    def _post_validation(self, verbose=True):
        assert self.config.beacon_window >= self.config.beacon_stride, f"Make sure the beacon_window {self.config.beacon_window} >= beacon_stride {self.config.beacon_stride}!"
        for ratio in self.config.beacon_ratio:
            assert ratio >= 0, f"Make sure all beacon ratios are greater than or equal to 0, found {self.config.beacon_ratio}!"
        assert self.config.beacon_attn in ["segmentation", "step-expansion", "full-coverage"], f"beacon_attn {self.config.beacon_attn} not implemented!"
        assert self.config.beacon_ratio_mix in ["instance-random", "step-random", "sequence"] or "adapt-" in self.config.beacon_ratio_mix, f"beacon_ratio_mix {self.config.beacon_ratio_mix} not implemented!"
        # assert self.config.beacon_pos in ["append", "interleave"], f"beacon_pos {self.config.beacon_pos} not implemented!"
        if self.config.beacon_pos == "interleave":
            assert self.config.beacon_window == self.config.beacon_stride, f"Make sure the beacon_window equals to beacon_stride when using interleaving mode."
        if self.config.beacon_parallel_window > 1:
            assert self.config._attn_implementation != "flash_attention_2", f"Currently parallel window does not support flash_attention_2!"

        self._cpu = torch.device("cpu")

        if verbose:
            info = f"applying activation beacon on {self.config.beacon_param} (the beacon embedding is initialized from {'bos' if self.config.beacon_embed_init == 'bos' else 'eos'} embedding, the beacon tokens are positioned with '{self.config.beacon_pos}' method), with window size {self.config.beacon_window}, stride {self.config.beacon_stride}, {self.config.beacon_attn} attention{' (attending to previous beacons)' if self.config.beacon_attend_prev else ' (no attending to previous beacons)'}, sink size {self.config.beacon_sink_size}, compression ratio {self.config.beacon_ratio} (mixed by {self.config.beacon_ratio_mix})..."
            logger.info(info)

    def set(self, verbose=True, **kwargs):
        """
        Set attributes out of the constructor.
        """
        for k, v in kwargs.items():
            setattr(self.config, k, v)
        self._post_validation(verbose=verbose)

    def reset(self):
        """Initialize attributes for a new sequence."""
        # the cursor pointing to the start of the current window
        self._start_idx = 0
        # the cursor pointing to the end of the current window
        self._end_idx = 0
        # the beacon sizes of all strides
        self._all_beacon_sizes = []
        # the loss per batch
        self._batch_loss = None
        # the valid token number per batch
        self._valid_token_num = None
        # the step index for processing the input_ids
        self._step_idx = 0
        # used in set_compression_ratio
        self._compression_ratio = None
        # the previous inputs is a full window or not, defaults to True
        self._is_full_window = True
        # the number of raw activations to preserve in update_memory (only useful when beacon_stride < beacon_window)
        self._raw_size_to_cache = 0

        # the number of tokens in previous stride that should be compressed by the upcoming beacon
        self._interleave_remainder = 0
        # compression ratio for the unfinished window
        self._interleave_compression_ratio = None
        self._beacon_indices = None

        self.all_input_ids = None
        self.all_attention_mask = None
        self.all_labels = None

        # NOTE: will be reset in prepare()
        self.beacon_skip_first = None
        self.beacon_skip_last = None

        # the raw activations of recent tokens
        self.raw_activations = [(None, None) for _ in range(self.config.num_hidden_layers)]
        # the attention sink activations
        self.sink_activations = [(None, None) for _ in range(self.config.num_hidden_layers)]
        # the beacon activations
        self.beacon_activations = [(None, None) for _ in range(self.config.num_hidden_layers)]

    @property
    def all_sequence_length(self):
        if self.all_input_ids is None:
            return 0
        else:
            return self.all_input_ids.shape[1]

    @property
    def batch_size(self):
        if self.all_input_ids is None:
            return 0
        else:
            return self.all_input_ids.shape[0]

    @property
    def finish(self):
        is_finish = self._end_idx == self.all_sequence_length
        return is_finish

    @property
    def dtype(self):
        return self.config.torch_dtype

    @property
    def min_value(self):
        return torch.finfo(self.dtype).min

    @property
    def max_position_embeddings(self):
        max_position_embeddings = self.config.max_position_embeddings
        if getattr(self.config, "rope_scaling", None) is not None:
            scaling_factor = self.config.rope_scaling["factor"]
            max_position_embeddings = max_position_embeddings * scaling_factor
        return max_position_embeddings
 
    def get_memory_size(self):
        """
        Sink memory size, beacon memory size and raw memory size.
        """
        sink_memory_size = 0
        beacon_memory_size = 0
        raw_memory_size = 0
        if self.sink_activations[0][0] is not None:
            sink_memory_size += self.sink_activations[0][0].shape[self.k_seq_dim]
        if self.beacon_activations[0][0] is not None:
            beacon_memory_size += self.beacon_activations[0][0].shape[self.k_seq_dim]
        if self.raw_activations[0][0] is not None:
            raw_memory_size += self.raw_activations[0][0].shape[self.k_seq_dim]
        return sink_memory_size, beacon_memory_size, raw_memory_size

    def prepare(self, input_ids, attention_mask, labels, skip_first=None, skip_last=None):
        """
        Prepare inputs for the model. These inputs belong to the same sequence.
        """
        # assert input_ids.shape[0] == 1, "Make sure the batch size is 1!"
        # assert attention_mask is None or (attention_mask == 1).all(), "Make sure there is no padding!"

        self._device = input_ids.device

        # accumulate input_ids
        if self.all_input_ids is None:
            self.all_input_ids = input_ids.cpu()
        else:
            self.all_input_ids = torch.cat([self.all_input_ids, input_ids.cpu()], dim=1)

        # accumulate attention_mask
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids, device=torch.device("cpu"))
        if self.all_attention_mask is None:
            self.all_attention_mask = attention_mask.cpu()
        else:
            self.all_attention_mask = torch.cat([self.all_attention_mask, attention_mask.cpu()], dim=1)

        # accumulate labels if exisits
        if labels is not None:
            # rotate labels in advance so that the loss of the last token is not ignored in every window
            labels = torch.cat([labels[:, 1:].cpu(), torch.tensor([-100]).expand(labels.shape[0], 1)], dim=1)
            if self.all_labels is None:
                self.all_labels = labels.cpu()
            else:
                self.all_labels = torch.cat([self.all_labels, labels], dim=1)
            assert self.all_input_ids.shape[1] == self.all_labels.shape[1], f"Found inconsistent all_input_ids {self.all_input_ids.shape} and all_labels {self.all_labels.shape}!"
        
        # how many tokens to skip at the beginning of the sequence? (They will be packed in a single chunk and processed by the model, after which their activations will be cached in sink_activations.)
        if skip_first is not None:
            assert self.config.beacon_parallel_window == 1, f"Make sure the parallel window is set to 1 when using beacon_skip!"
            assert self.config.beacon_window == self.config.beacon_stride, f"Make sure the beacon_window equals to beacon_stride when using beacon_skip."
            assert self.config.beacon_sink_size == 0, f"Make sure the beacon_sink_size is set to 0 when using beacon_skip!"
        # stop compression after how many tokens
        if skip_last is not None:
            skip_first = skip_first if skip_first is not None else 0
            assert (skip_last - skip_first) % self.config.beacon_window == 0, f"skip_last ({skip_last}) - skip_first ({skip_first}) = {skip_last - skip_first} is not divisible by window size {self.config.beacon_window}"
            assert self.config.beacon_sink_size == 0, "Make sure the beacon_sink_size is zero when using skip_last!"
        self.beacon_skip_first = skip_first
        self.beacon_skip_last = skip_last

    def set_compression_ratio(self, start_idx, end_idx):
        """Choose a condensing ratio from self.config.beacon_ratio"""
        def filter_ratio(ratios, stride):
            valid_ratios = []
            for ratio in ratios:
                # stride must be bigger than condensing ratio because we there must be at least one beacon
                if stride < ratio:
                    continue
                # the stride must be evenly divisible by condensing ratio
                if ratio > 0 and (stride % ratio) != 0:
                    continue
                # when training, ratio=0 is valid if previous windows contain beacon or later windows contain beacon
                if ratio == 0 and self.training:
                    previous_has_zero = -1 in self._all_beacon_sizes
                    following_has_nonzero = (start_idx + stride + self.config.beacon_window) <= self.all_sequence_length
                    if previous_has_zero or (not following_has_nonzero):
                        continue
                valid_ratios.append(ratio)
            assert len(valid_ratios), f"Cannot find valid condensing ratio (among {ratios}) for stride {stride}!"
            return valid_ratios

        def get_max_length(ratios):
            max_lengths = []
            for compression_ratio in ratios:
                if compression_ratio > 0:
                    # NOTE: here we must use the scaled position embeddings
                    max_lengths.append((self.max_position_embeddings - self.config.beacon_window) * compression_ratio + self.config.beacon_window)
                else:
                    max_lengths.append(self.max_position_embeddings)
            return max_lengths

        if len(self.config.beacon_ratio) == 1:
            return self.config.beacon_ratio[0]

        ratio_mix = self.config.beacon_ratio_mix

        beacon_ratio = filter_ratio(self.config.beacon_ratio, self.config.beacon_stride)

        if ratio_mix == "instance-random":
            if self._compression_ratio is None:
                beacon_ratio = self.rng.choice(beacon_ratio).tolist()
                self._compression_ratio = beacon_ratio
            else:
                beacon_ratio = self._compression_ratio

        elif ratio_mix == "step-random":
            beacon_ratio = self.rng.choice(beacon_ratio).tolist()
        
        elif ratio_mix == "sequence":
            if self._compression_ratio is None:
                self._compression_ratio = cycle(beacon_ratio)
            beacon_ratio = next(self._compression_ratio)

        elif "adapt" in ratio_mix:
            if self._compression_ratio is None:
                future_length = int(ratio_mix.split("-")[1])
                sequence_length = self.all_input_ids.shape[1] + future_length
                max_lengths = get_max_length(beacon_ratio)
                # ascendingly sort the max lengths
                valid_max_lengths_and_indices = [x for x in enumerate(max_lengths) if x[1] >= sequence_length]
                if len(valid_max_lengths_and_indices):
                    minimum_length_index = min(valid_max_lengths_and_indices, key=lambda x: x[1])[0]
                    # use the minimal possible length for this sequence (the smallest fold ratio)
                    beacon_ratio = beacon_ratio[minimum_length_index]
                else:
                    beacon_ratio = max(beacon_ratio)
                    # logger.warning(f"Failed to find valid fold window and size for sequence length {sequence_length}, as the maximum theoretical length is {max(max_lengths)}. Fall back to use the maximum one: {beacon_ratio}.")
                self._compression_ratio = beacon_ratio
            else:
                beacon_ratio = self._compression_ratio

        return beacon_ratio

    def step(self):
        # parallel does not support stride < window
        # parallel does not support non-compression
        # the input_ids is not long enough for parallel
        if \
        (self.config.beacon_parallel_window > 1) and \
        (self.config.beacon_stride == self.config.beacon_window) and \
        (0 not in self.config.beacon_ratio) and \
        (self.all_input_ids[:, self._end_idx:].shape[1] >= self.config.beacon_parallel_window * self.config.beacon_window):
            input_ids_list = []
            attention_mask_list = []
            position_ids_list = []
            labels_list = []

            beacon_size_list = []
            beacon_indices_list = []

            for i in range(self.config.beacon_parallel_window):
                if i == 0:
                    _input_ids, _attention_mask, _position_ids, _past_key_values, _labels = self._step()
                else:
                    _input_ids, _attention_mask, _position_ids, _past_key_values, _labels = self._step(ignore_memory=True)

                input_ids_list.append(_input_ids)
                attention_mask_list.append(_attention_mask)
                position_ids_list.append(_position_ids)
                labels_list.append(_labels)
                beacon_size_list.append(_past_key_values[0][2])
                beacon_indices_list.append(_past_key_values[0][3])

                if i == 0:
                    past_key_values = _past_key_values
                    if past_key_values[0][0] is None:
                        mem_size = 0
                    else:
                        mem_size = past_key_values[0][0].shape[self.k_seq_dim]

                else:
                    # no memory
                    assert _past_key_values[0][0] is None
            
            batch_size = self.all_input_ids.shape[0]
            # NOTE: we do not need to repliace beacon tokens for the last window
            seq_len = sum(x.shape[1] for x in input_ids_list) + sum(beacon_size_list) - beacon_size_list[-1]

            input_ids = _input_ids.new_zeros((batch_size, seq_len)) + self.beacon_token.to(_input_ids.device)
            # all 0
            attention_mask = _attention_mask.new_zeros((batch_size, 1, seq_len, mem_size + seq_len)) + self.min_value
            position_ids = torch.arange(mem_size + seq_len, device=self._device).expand(batch_size, mem_size + seq_len)
            # 2 indicates the beacon token is used for replication
            beacon_indices = beacon_indices_list[0].new_zeros(seq_len) + 2
            if _labels is not None:
                # -100 because no loss on beacon tokens
                labels = _labels.new_zeros((batch_size, seq_len)) - 100
            else:
                labels = None

            start_idx = 0
            position_offset = mem_size
            for i in range(self.config.beacon_parallel_window):
                beacon_size = beacon_size_list[i]

                # populate input_ids
                _input_ids = input_ids_list[i]
                cur_seq_len = _input_ids.shape[1]
                input_ids[:, start_idx: start_idx + cur_seq_len] = _input_ids
                
                # populate attention_mask and position_ids
                _attention_mask = attention_mask_list[i]
                _position_ids = position_ids_list[i]
                # the attention mask in the first window contains the mask for memory, which is redundant here
                if i == 0:
                    _attention_mask = _attention_mask[:, :, :, mem_size:]
                    _position_ids = _position_ids[:, mem_size:] - mem_size

                attention_mask[:, :, start_idx: start_idx + cur_seq_len, mem_size + start_idx: mem_size + start_idx + cur_seq_len] = _attention_mask
                position_ids[:, mem_size + start_idx: mem_size + start_idx + cur_seq_len] = _position_ids + position_offset

                # populate beacon_indices
                _beacon_indices = beacon_indices_list[i]
                beacon_indices[start_idx: start_idx + cur_seq_len] = _beacon_indices

                # populate labels
                if labels is not None:
                    # populate labels
                    _labels = labels_list[i]
                    labels[:, start_idx: start_idx + cur_seq_len] = _labels

                # NOTE: when there is sink activations, we need to bias the position_ids for the first window
                if i == 0 and self.config.beacon_sink_size > 0 and self.sink_activations[0][0] is None:
                    position_offset += 1

                # modify the attention and position for replicated beacon tokens
                if i != self.config.beacon_parallel_window - 1:
                    replicate_beacon_row_start = start_idx + cur_seq_len
                    replicate_beacon_col_start = mem_size + start_idx + cur_seq_len
                    # NOTE: any attention mask is okay for replicated beacon tokens, but for convenience we use the causal mask
                    attention_mask[:, :, replicate_beacon_row_start: replicate_beacon_row_start + beacon_size, replicate_beacon_col_start: replicate_beacon_col_start + beacon_size] = _attention_mask.new_full((beacon_size, beacon_size), self.min_value).triu(1)
                    # NOTE: all future tokens can attend to the replicated beacon tokens
                    attention_mask[:, :, replicate_beacon_row_start + beacon_size:, replicate_beacon_col_start: replicate_beacon_col_start + beacon_size] = 0
                    # NOTE: the position of replicated beacon tokens start from 0
                    position_ids[:, mem_size + start_idx + cur_seq_len: mem_size + start_idx + cur_seq_len + beacon_size] = torch.arange(position_offset, position_offset + beacon_size, device=_input_ids.device)[None:]

                start_idx += cur_seq_len + beacon_size
                position_offset += beacon_size

            # the memory is visible to all subsequent tokens
            attention_mask[:, :, :, :max(mem_size, self.config.beacon_sink_size)] = 0

            # NOTE: modify beacon_indices
            for i, (key, value, _, _) in enumerate(past_key_values):
                past_key_values[i] = (key, value, sum(beacon_size_list), beacon_indices)

            # NOTE: update _beacon_indices so that the next-token logits can be properly sliced out in self.output()
            self._beacon_indices = beacon_indices
            
            return input_ids, attention_mask, position_ids, past_key_values, labels

        else:
            return self._step()

    def _step(self, ignore_memory=False):
        """
        Yield inputs for the current sliding window, including the input_ids, attention_mask, position_ids, and past_key_values.
        """
        #============================================#
        # Check whether the inputs fulfills a window.
        #============================================#

        # the starting position of the current window w.r.t. the start of the current input sequence
        start_idx = self._start_idx
        # the end position of the current window w.r.t. the start of the current input sequence
        end_idx = start_idx + self.config.beacon_window
        # indicates if the current window is completely filled by raw activations and new tokens
        # we only append beacon tokens for full windows
        if end_idx > self.all_sequence_length:
            # the input is shorter than the initial window size
            end_idx = self.all_sequence_length
            is_full_window = False
        else:
            is_full_window = True

        # NOTE: in training, the entire sequence is input to the model at once
        # In the last window, we do not need to append beacons because they will not be used at all
        if self.training and end_idx == self.all_sequence_length:
            next_start_idx = start_idx
            is_full_window = False
            raw_size_to_cache = -1
            beacon_size = 0
            compression_ratio = -1
        
        elif self._step_idx == 0 and self.beacon_skip_first is not None:
            end_idx = start_idx + self.beacon_skip_first
            assert end_idx < self.all_sequence_length
            next_start_idx = end_idx
            is_full_window = True
            raw_size_to_cache = -1
            beacon_size = 0
            compression_ratio = -1
        
        elif self.beacon_skip_last is not None and start_idx >= self.beacon_skip_last:
            end_idx = min(start_idx + self.config.beacon_window, self.all_sequence_length)
            next_start_idx = end_idx
            is_full_window = False
            raw_size_to_cache = -1
            beacon_size = 0
            compression_ratio = -1

        else:
            #============================================#
            # Set compression ratio
            #============================================#
            if self.config.beacon_pos == "append":
                if is_full_window:
                    # determine compression ratio for the current window
                    beacon_stride = self.config.beacon_stride
                    compression_ratio = self.set_compression_ratio(start_idx=start_idx, end_idx=end_idx)

                    if compression_ratio > 0:
                        # the stride must be evenly divisible by compression_ratio
                        beacon_size = beacon_stride // compression_ratio
                    else:
                        # the raw activations are used as beacon activations
                        beacon_size = -1

                    # forward start_idx and end_idx
                    next_start_idx = start_idx + beacon_stride
                    # how many raw activations to save
                    raw_size_to_cache = end_idx - next_start_idx
                else:
                    # no stride because the sequence has finished
                    next_start_idx = start_idx
                    # cache all raw activations
                    raw_size_to_cache = -1
                    beacon_size = 0
                    compression_ratio = 0

            elif self.config.beacon_pos == "interleave":
                # the number of raw tokens in the input_ids
                input_size = end_idx - self._end_idx
                # set compression ratio once the previous window has finished, otherwise, reuse the interleave_compression_ratio if the input belongs to an unfinished window
                if self._is_full_window:
                    compression_ratio = self.set_compression_ratio(start_idx=start_idx, end_idx=end_idx)
                    self._interleave_compression_ratio = compression_ratio
                else:
                    compression_ratio = self._interleave_compression_ratio

                # the beacon size is non-zero even if the window is not full
                if compression_ratio > 0:
                    # this number of beacon tokens will be inserted among the raw tokens
                    beacon_size = (input_size + self._interleave_remainder) // compression_ratio
                else:
                    # the raw activations are used as beacon activations
                    beacon_size = -1

                if is_full_window:
                    # move forward one window
                    next_start_idx = start_idx + self.config.beacon_stride
                    # no save raw activations
                    raw_size_to_cache = 0
                else:
                    # no stride because the sequence has not finished
                    next_start_idx = start_idx
                    # cache all recent raw activations to be used in the next window
                    raw_size_to_cache = -1

        #============================================#
        # Slice out input_ids (raw tokens in the current window)
        #============================================#
        input_ids = self.all_input_ids[:, self._end_idx: end_idx].to(self._device)
        attention_mask = self.all_attention_mask[:, self._end_idx: end_idx].to(self._device)
        if self.all_labels is not None:
            labels = self.all_labels[:, self._end_idx: end_idx].to(self._device)
        else:
            labels = None
        batch_size = input_ids.shape[0]

        #============================================#
        # Insert beacon tokens if necessary.
        #============================================#
        # t1 = time.time()

        if self.config.beacon_pos == "append":
            # append beacons if necessary
            if is_full_window and beacon_size > 0:
                input_ids = torch.cat([input_ids, self.beacon_token.expand(batch_size, beacon_size).to(input_ids.device, dtype=input_ids.dtype)], dim=1)
                # NOTE: prepend 1 to attention_mask because we have past_key_values
                attention_mask = torch.cat([attention_mask, attention_mask.new_ones(batch_size, beacon_size)], dim=1)
                if labels is not None:
                    labels = torch.cat([labels, labels.new_zeros(batch_size, beacon_size) - 100], dim=1)

        elif self.config.beacon_pos == "interleave":
            input_len = input_ids.shape[1]
            if beacon_size > 0:
                # insert beacon tokens in between raw tokens
                input_ids_with_beacons = input_ids.new_full((input_ids.shape[0], input_len + beacon_size), self.beacon_token.item())
                raw_token_indices = torch.arange(input_ids_with_beacons.shape[1], device=input_ids.device)
                interleave_start_idx = compression_ratio - self._interleave_remainder
                raw_token_indices = raw_token_indices[raw_token_indices % (compression_ratio + 1) != interleave_start_idx].unsqueeze(0).expand_as(input_ids)
                input_ids_with_beacons = input_ids_with_beacons.scatter(dim=1, index=raw_token_indices, src=input_ids)
                input_ids = input_ids_with_beacons
                # attention mask
                attention_mask_with_beacons = attention_mask.new_full((attention_mask.shape[0], attention_mask.shape[1] + beacon_size), 1)
                attention_mask_with_beacons = attention_mask_with_beacons.scatter(dim=1, index=raw_token_indices, src=attention_mask)
                attention_mask = attention_mask_with_beacons
                # labels
                if labels is not None:
                    labels_with_beacons = labels.new_full((labels.shape[0], labels.shape[1] + beacon_size), -100)
                    labels_with_beacons = labels_with_beacons.scatter(dim=1, index=raw_token_indices, src=labels)
                    labels = labels_with_beacons

            if compression_ratio > 0:
                # update the reminder
                self._interleave_remainder = (input_len + self._interleave_remainder) % compression_ratio

        # NOTE: skip computing loss in the very first window because the beacon tokens will be used in the next window
        if self.training and self._step_idx == 0 and not (self.config.beacon_pos == 'interleave' and self.config.beacon_attn == 'full-coverage'):
            labels[:] = -100

        # t2 = time.time()

        #============================================#
        # Prepare beacon_indices for interleave beacon_pos, a boolean mask where True indicates the beacon tokens.
        # The mask is applied on the inputs of the entire window, including the cached activations and the input_ids.
        #============================================#
        beacon_indices = (input_ids[0] == self.beacon_token.item()).long()
        if self._is_full_window:
            self._beacon_indices = torch.tensor([], dtype=torch.long, device=input_ids.device)
        # the beacon_indices always tracks the beacon tokens in both the cached activations and the input_ids
        beacon_indices = torch.cat([self._beacon_indices, beacon_indices])
        # record the beacon_indices for the next window
        self._beacon_indices = beacon_indices
        if is_full_window and beacon_size == -1:
            # NOTE: the first beacon_stride raw tokens serve as beacon tokens
            # we use -1 to indicate these raw tokens, so that the attention mask and position ids will not be modified
            beacon_indices[:self.config.beacon_stride] = -1

        # t3 = time.time()

        #============================================#
        # Prepare past_key_values.
            # beacon_size: how many beacon tokens are there in the input_ids
            # beacon_indices: the boolean mask for the entire window where True indicates the beacon tokens (for append, the beacon_indices corresponds to input_ids, while for 'interleave', the beacon_indices corresponds to the entire window including both the input_ids and the cached activations)
        #============================================#
        past_key_values = []
        for layer_idx in range(self.config.num_hidden_layers):
            if ignore_memory:
                key, value = None, None
            else:
                sink_key, sink_value = self.sink_activations[layer_idx]
                beacon_key, beacon_value = self.beacon_activations[layer_idx]
                raw_key, raw_value = self.raw_activations[layer_idx]

                key = cat_tensor([
                    sink_key, beacon_key, raw_key,
                ], dim=self.k_seq_dim)
                value = cat_tensor([
                    sink_value, beacon_value, raw_value,
                ], dim=self.v_seq_dim)

            layer_past_key_values = (key, value, beacon_size, beacon_indices)
            past_key_values.append(layer_past_key_values)

        # t4 = time.time()
        
        #============================================#
        # Prepare attention_mask and position_ids.
        #============================================#
        first_key = past_key_values[0][0]
        mem_size = first_key.shape[self.k_seq_dim] if first_key is not None else 0
        if mem_size > 0:
            attention_mask = torch.cat([attention_mask.new_ones(batch_size, mem_size), attention_mask], dim=1)

        input_length = input_ids.shape[1]
        position_ids = torch.arange(attention_mask.shape[-1], dtype=torch.long, device=self._device).repeat(batch_size, 1)

        if self.config._attn_implementation == "flash_attention_2":
            assert self.config.beacon_attn == "full-coverage", f"Make sure to set beacon_attn='full-coverage' when using flash attention! Found {self.config.beacon_attn}."
            if 0 in attention_mask:
                pass
            else:
                attention_mask = None
        elif self.config._attn_implementation == "sdpa" and self.config.beacon_pos == "append" and beacon_size <= 0 and (input_length == 1 or mem_size == 0):
            attention_mask = None
        else:
            attention_mask, position_ids = self._make_4d_attention_mask_and_position_ids(
                attention_mask, 
                position_ids, 
                mem_size, 
                beacon_size, 
                compression_ratio, 
            )

        # t5 = time.time()

        # print(f"prepare inputs {t2-t1}, prepare indices {t3-t2}, prepare memory {t4-t3}, prepare attention mask {t5-t4}")

        #============================================#
        # Update necessary attributes.
        #============================================#
        # keep track of whether the current inputs is a full_window
        self._is_full_window = is_full_window
        # keep track of the raw_size_to_cache
