import os
import torch
import time
import numpy as np
import torch.distributed as dist
from copy import deepcopy
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._post_validation()
        self.reset()
    
    @property
    def beacon_token(self):
        return self.config.vocab_size

    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, **kwargs):
        """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 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)]
        # the raw activations of recent tokens
        self.raw_activations = [(None, None) for _ in range(self.config.num_hidden_layers)]

        # NOTE: in case we want to resume the memory from a particular state
        for k, v in kwargs.items():
            # NOTE: deepcopy to untie the memory state and the model memory
            setattr(self, deepcopy(k), deepcopy(v))

    def export(self):
        """Export all necessary attributes of the memory module."""
        return {
            "start_idx": self.start_idx,
            "end_idx": self.end_idx,
            "all_beacon_sizes": self.all_beacon_sizes,
            "batch_loss": self.batch_loss,
            "valid_token_num": self.valid_token_num,
            "step_idx": self.step_idx,
            "compression_ratio": self.compression_ratio,
            "is_full_window": self.is_full_window,
            "raw_size_to_cache": self.raw_size_to_cache,
            "interleave_remainder": self.interleave_remainder,
            "interleave_compression_ratio": self.interleave_compression_ratio,
            "beacon_indices": self.beacon_indices,
            "all_input_ids": self.all_input_ids,
            "all_attention_mask": self.all_attention_mask,
            "all_labels": self.all_labels,
            "beacon_skip_first": self.beacon_skip_first,
            "beacon_skip_last": self.beacon_skip_last,
            # NOTE: deepcopy to untie the memory state and the model memory
            "sink_activations": deepcopy(self.sink_activations),
            "beacon_activations": deepcopy(self.beacon_activations),
            "raw_activations": deepcopy(self.raw_activations),
        }

    @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

    @property
    def beacon_window(self):
        if (
            self.beacon_skip_last is not None
            and self.start_idx < self.beacon_skip_last
            and self.start_idx + self.config.beacon_window > self.beacon_skip_last
        ):
            return self.beacon_skip_last - self.start_idx
        else:
            return self.config.beacon_window

    @property
    def beacon_stride(self):
        if (
            self.beacon_skip_last is not None
            and self.start_idx < self.beacon_skip_last
            and self.start_idx + self.config.beacon_window > self.beacon_skip_last
        ):
            return self.beacon_skip_last - self.start_idx
        else:
            return self.config.beacon_stride
 
    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.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.beacon_window) * compression_ratio + self.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.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
            # 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.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
        
        # NOTE: we do not compress the beacon_skip_first tokens at the beginning of the sequence
        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
        
        # NOTE: we do not compress tokens after beacon_skip_last tokens
        elif self.beacon_skip_last is not None and start_idx >= self.beacon_skip_last:
            end_idx = min(start_idx + self.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.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.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, input_ids.new_full((batch_size, beacon_size), self.beacon_token)], 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)
                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).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.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
        self.raw_size_to_cache = raw_size_to_cache
        # involked in self.output()
        self.all_beacon_sizes.append(beacon_size)
        # update start_idx and end_idx
        # NOTE: the update of start_idx will influence self.beacon_window and self.beacon_stride in case self.beacon_skip_last is not None
        # Therefore, we must make sure all calls to self.beacon_window and self.beacon_stride happen before the update of start_idx
        self.start_idx = next_start_idx
        self.end_idx = end_idx
        self.step_idx += 1

        # print(f"start_idx:          {start_idx}")
        # print(f"next_start_idx:     {next_start_idx}")
        # print(f"beacon_size:        {beacon_size}")
        # print(f"raw_size_to_cache:  {raw_size_to_cache}")
        # print(f"interleave_remainder:{self.interleave_remainder}")
        # print(f"input_ids:          {input_ids}")
        # print(f"input_len:          {input_len}")
        # print(f"beacon_indices:     {beacon_indices}")
        # print(f"position_ids:       {position_ids}")
        # print(f"attention_mask:\n{attention_mask == 0}")
        # x = input()
        # if x == "s":
        #     return

        return input_ids, attention_mask, position_ids, past_key_values, labels

    def update_memory(self, past_key_values):
        """
        Accumulate beacon activations and raw activations.
        """
        for layer_idx, (key, value, beacon_size, beacon_indices) in enumerate(past_key_values):
            # NOTE: the past_key_values are incrementally returned (only the new keys and values are returned)
            previous_raw_key, previous_raw_value = self.raw_activations[layer_idx]

            if self.beacon_skip_first is not None and self.sink_activations[layer_idx][0] is None:
                assert key.shape[self.k_seq_dim] == self.beacon_skip_first
                assert value.shape[self.k_seq_dim] == self.beacon_skip_first
                self.sink_activations[layer_idx] = [
                    key,
                    value,
                ]
                # NOTE: no need to update raw activations and beacon activations as all activations are kept as sink activations
                continue

            if self.beacon_activations[layer_idx][0] is None and self.config.beacon_sink_size > 0:
                # save the sink activations
                # NOTE: we do not slice the key/value activations, which may cause duplication when beacon_ratio=-1 for the first window, but it's okay
                self.sink_activations[layer_idx] = [
                    slice_tensor(key, end=self.config.beacon_sink_size, dim=self.k_seq_dim),
                    slice_tensor(value, end=self.config.beacon_sink_size, dim=self.v_seq_dim),
                ]

            if not self.is_full_window:
                # this means the current input does not fulfill a window
                # thus, the key and value are all raw activations, and we accumulate them until the window is fulfilled
                assert self.raw_size_to_cache == -1
                raw_key = cat_tensor([
                    previous_raw_key,
                    key
                ], dim=self.k_seq_dim)
                raw_value = cat_tensor([
                    previous_raw_value, 
                    value
                ], dim=self.v_seq_dim)
                self.raw_activations[layer_idx] = (raw_key, raw_value)

            else:
                # NOTE: use the correct previous_beacon_key and value!
                previous_beacon_key, previous_beacon_value = self.beacon_activations[layer_idx]
                
                beacon_key, beacon_value, raw_key, raw_value = self._extract_beacon_and_raw_memory(
                    key, 
                    value, 
                    previous_beacon_key, 
                    previous_beacon_value, 
                    previous_raw_key, 
                    previous_raw_value, 
                    beacon_indices,
                )

                self.beacon_activations[layer_idx] = (beacon_key, beacon_value)
                self.raw_activations[layer_idx] = (raw_key, raw_value)

    def update_loss(self, batch_loss, valid_token_num):
        """
        Accumulate loss for later perplexity computation and backward pass.
        """
        if self.batch_loss is None:
            # NOTE: multiply valid_token_num because batch_loss is divided by it in advance
            self.batch_loss = batch_loss * valid_token_num
            self.valid_token_num = valid_token_num
        else:
            # NOTE: avoid in-place operations, otherwise there will be gradient errors in training
            self.batch_loss = self.batch_loss + batch_loss * valid_token_num
            self.valid_token_num = self.valid_token_num + valid_token_num

    def output(self, model_outputs):
        """
        Override loss with accumulated loss. Update the next-token logits.
        """
        # override loss
        if self.batch_loss is not None:
            # here the batch_loss is the summation of all token losses in each element
            loss = self.batch_loss.sum() / self.valid_token_num.sum()

            # NOTE: prevent nan
            batch_loss = self.batch_loss / self.valid_token_num
            if (self.valid_token_num == 0).any():
                batch_loss = batch_loss.masked_fill(self.valid_token_num == 0, 0.)

            # NOTE: we must use dict to override values, otherwise trainer cannot find loss
            model_outputs["loss"] = loss
            model_outputs["batch_loss"] = batch_loss

        # override last_hidden_states (used in generation)
        beacon_size = self.all_beacon_sizes[-1]
        # remove logits corresponding to beacon tokens
        if beacon_size > 0:
            logits = model_outputs["logits"]
            beacon_indices = self.beacon_indices[-logits.shape[1]:]
            model_outputs["logits"] = logits[:, beacon_indices == 0]

        return model_outputs

    def _make_4d_attention_mask_and_position_ids(
        self, 
        attention_mask, 
        position_ids,
        mem_size, 
        beacon_size, 
        compression_ratio, 
    ):
        """
        Convert attention_mask into causal 4D attention_mask (batch_size, head_num, query_len, key_len).
        """
        tgt_size = attention_mask.size(-1) - mem_size
        dtype = self.dtype
        min_value = self.min_value
        device = self._device
        batch_size, src_size = attention_mask.size()

        # square for memory, and lower triangular for input_ids
        causal_mask = torch.full((tgt_size, tgt_size), min_value, device=device, dtype=dtype)
        mask_cond = torch.arange(causal_mask.size(-1), device=device)
        causal_mask.masked_fill_(mask_cond < (mask_cond + 1).view(causal_mask.size(-1), -1), 0)
        causal_mask = torch.cat([torch.zeros(tgt_size, mem_size, dtype=dtype, device=device), causal_mask], dim=-1)
        causal_mask = causal_mask[None, None, ...].expand(batch_size, 1, tgt_size, src_size)
        # 1 for non-padding tokens
        expand_mask = attention_mask[:, None, None, :].expand(batch_size, 1, tgt_size, src_size)
        invert_mask = 1.0 - expand_mask
        invert_mask.masked_fill_(invert_mask.bool(), min_value)

        attention_mask = causal_mask.masked_fill(invert_mask.bool(), min_value)

        if self.config.beacon_attn == "step-expansion":
            # each beacon can attend to one more sub-interval than its predecessor

            if self.config.beacon_pos == "append" and beacon_size > 0:
                window_size = self.beacon_window
                window_size_with_beacon = window_size + beacon_size
                beacon_start_idx = -beacon_size
                # batch_size, head_num, window_size
                reference_attention_mask = attention_mask[..., -beacon_size - 1, -window_size_with_beacon: -beacon_size]

                # compression_ratio, 2 * compression_ratio, ..., beacon_size * compression_ratio
                beacon_arange = torch.arange(1, beacon_size + 1, device=device) * compression_ratio
                # 0, 1, 2, ..., window_size - 1
                ordinal_arange = torch.arange(window_size, device=device)
                # beacon_size, window_size
                valid_pos = ordinal_arange.expand(beacon_size, window_size) < beacon_arange.unsqueeze(-1)
                # beacon_size, window_size
                ordinal_attention_mask = torch.where(valid_pos, 0, min_value)
                # NOTE: add reference attention_mask so that padding tokens are considered
                ordinal_attention_mask = ordinal_attention_mask[None, None, ...] + reference_attention_mask.unsqueeze(-2)

                if self.config.beacon_attend_prev:
                    beacon_attention_mask = attention_mask.new_full((beacon_size, beacon_size), min_value).triu(1)
                    # the beacon token is next to the last ordinal token it attends to
                    ordinal_position_ids = position_ids[:, -window_size_with_beacon: -beacon_size]
                    beacon_position_ids = ordinal_position_ids[:, compression_ratio - 1::compression_ratio] + torch.arange(1, beacon_size + 1, device=device)[None]
                    position_ids[:, beacon_start_idx:] = beacon_position_ids
                else:
                    beacon_attention_mask = attention_mask.new_full((beacon_size, beacon_size), min_value).fill_diagonal_(0)
                    # the beacon token is next to the last ordinal token it attends to
                    ordinal_position_ids = position_ids[:, -window_size_with_beacon: -beacon_size]
                    beacon_position_ids = ordinal_position_ids[:, compression_ratio - 1::compression_ratio] + 1
                    position_ids[:, beacon_start_idx:] = beacon_position_ids

                attention_mask[..., beacon_start_idx:, -window_size_with_beacon: -beacon_size] = ordinal_attention_mask
                attention_mask[..., beacon_start_idx:, beacon_start_idx:] = beacon_attention_mask

            # NOTE: the attention mask should be modified when there is beacon token within the window, not in the input_ids
            elif self.config.beacon_pos == "interleave" and (self.beacon_indices == 1).any():
                assert self.config.beacon_attend_prev == False, f"Make sure beacon_attend_prev is False if using 'interleave' beacon pos!"

                beacon_indices = self.beacon_indices

                cur_position_ids = position_ids[:, -len(beacon_indices):]
                base_position = cur_position_ids[:, 0] - 1
                # NOTE: alternate position so that the position of raw tokens are consistent
                position_template = cur_position_ids.new_ones(cur_position_ids.shape)
                position_template[:, compression_ratio + 1::compression_ratio + 1] = 0
                cur_position_ids = base_position + position_template.cumsum(-1)
                position_ids[:, -len(beacon_indices):] = cur_position_ids

                cur_input_length = len(beacon_indices)
                cur_attention_mask = attention_mask[..., -cur_input_length:, -cur_input_length:]
                # mask all beacon columns
                cur_attention_mask[..., beacon_indices] = min_value
                # beacon tokens can attend to themselves
                input_ids_attention_mask = cur_attention_mask[..., -tgt_size:, -tgt_size:]
                input_ids_attention_mask[..., range(tgt_size), range(tgt_size)] = 0

        elif self.config.beacon_attn == "segmentation":
            # each beacon can attend to its corresponding sub-interval

            if self.config.beacon_pos == "append" and beacon_size > 0:
                window_size = self.beacon_window
                window_size_with_beacon = window_size + beacon_size
                beacon_start_idx = -beacon_size
                # batch_size, head_num, window_size
                reference_attention_mask = attention_mask[..., -beacon_size - 1, -window_size_with_beacon: -beacon_size]

                # beacon_size, compression_ratio
                indices = torch.arange(compression_ratio * beacon_size, device=device).view(beacon_size, -1)
                # beacon_size, window_size
                ordinal_attention_mask = attention_mask.new_full((beacon_size, window_size), min_value)
                ordinal_attention_mask.scatter_(dim=-1, index=indices, value=0)

                # NOTE: add reference attention_mask so that padding tokens are considered
                ordinal_attention_mask = ordinal_attention_mask[None, None, ...] + reference_attention_mask.unsqueeze(-2)

                if self.config.beacon_attend_prev:
                    beacon_attention_mask = attention_mask.new_full((beacon_size, beacon_size), min_value).triu(1)
                    # the beacon token is next to the last ordinal token it attends to
                    beacon_position_ids = position_ids.new_full(beacon_size, fill_value=compression_ratio + mem_size)
                    beacon_position_ids = beacon_position_ids + torch.arange(beacon_size)
                    position_ids[:, beacon_start_idx:] = beacon_position_ids
                else:
                    beacon_attention_mask = attention_mask.new_full((beacon_size, beacon_size), min_value).fill_diagonal_(0)
                    # the beacon token is next to the last ordinal token it attends to
                    beacon_position_ids = position_ids.new_full(beacon_size, fill_value=compression_ratio + mem_size)
                    position_ids[:, beacon_start_idx:] = beacon_position_ids

                attention_mask[..., beacon_start_idx:, -window_size_with_beacon: -beacon_size] = ordinal_attention_mask
                attention_mask[..., beacon_start_idx:, beacon_start_idx:] = beacon_attention_mask
                # beacons of different ratios are blind to others
                attention_mask[..., beacon_start_idx:, -beacon_size: beacon_start_idx] = min_value

            elif self.config.beacon_pos == "interleave":
                raise NotImplementedError

        elif self.config.beacon_attn == "full-coverage":
            pass

        return attention_mask, position_ids

    def _extract_beacon_and_raw_memory(
        self, 
        key, 
        value, 
        previous_beacon_key, 
        previous_beacon_value, 
        previous_raw_key, 
        previous_raw_value, 
        beacon_indices,
    ):
        """Extract beacon and raw memory from the returned key and value when the window is full."""
        key = cat_tensor([
            previous_raw_key, 
            key
        ], dim=self.k_seq_dim)
        value = cat_tensor([
            previous_raw_value, 
            value
        ], dim=self.v_seq_dim)

        # NOTE: we use magic slice instead of boolean index here for efficiency
        beacon_key = slice_tensor(key, index=torch.logical_or(beacon_indices == 1, beacon_indices == -1), dim=self.k_seq_dim)
        beacon_value = slice_tensor(value, index=torch.logical_or(beacon_indices == 1, beacon_indices == -1), dim=self.v_seq_dim)

        if self.config.beacon_accum:
            beacon_key = cat_tensor([previous_beacon_key, beacon_key], dim=self.k_seq_dim)
            beacon_value = cat_tensor([previous_beacon_value, beacon_value], dim=self.v_seq_dim)

        if self.raw_size_to_cache > 0:
            raw_key = slice_tensor(key, index=beacon_indices == 0, dim=self.k_seq_dim)
            raw_key = slice_tensor(raw_key, start=-raw_size_to_cache, dim=self.k_seq_dim)

            raw_value = slice_tensor(value, index=beacon_indices == 0, dim=self.v_seq_dim)
            raw_value = slice_tensor(raw_value, start=-raw_size_to_cache, dim=self.v_seq_dim)        

        else:
            raw_key = None
            raw_value = None

        return beacon_key, beacon_value, raw_key, raw_value


def slice_tensor(x, start=None, end=None, step=None, index=None, dim=2):
    if x is None:
        return None
    if end == 0:
        return None
    if start == x.shape[dim]:
        return None
    if start is not None and start == end:
        return None
    if dim == 2:
        if index is not None:
            return x[:, :, index]
        elif start is None and end is not None:
            if step is None:
                return x[:, :, :end, ...]
            else:
                return x[:, :, :end:step, ...]
        elif start is not None and end is None:
            if step is None:
                return x[:, :, start:, ...]
            else:
                return x[:, :, start::step, ...]
        elif start is not None and end is not None:
            if step is None:
                return x[:, :, start:end, ...]
            else:
                return x[:, :, start:end:step, ...]
    elif dim == 1:
        if index is not None:
            return x[:, :, index]
        elif start is None and end is not None:
            if step is None:
                return x[:, :end, ...]
            else:
                return x[:, :end:step, ...]
        elif start is not None and end is None:
            if step is None:
                return x[:, start:, ...]
            else:
                return x[:, start::step, ...]
        elif start is not None and end is not None:
            if step is None:
                return x[:, start:end, ...]
            else:
                return x[:, start:end:step, ...]
    else:
        raise NotImplementedError

def cat_tensor(list_of_tensors, dim=-1):
    list_of_tensors = [t for t in list_of_tensors if t is not None]
    if len(list_of_tensors) > 1:
        result = torch.cat(list_of_tensors, dim=dim)
    elif len(list_of_tensors) == 1:
        result = list_of_tensors[0]
    else:
        result = None
    return result

def slice_activations(activations, start=None, end=None, k_seq_dim=2, v_seq_dim=2):
    new_activations = []
    for key, value in activations:
        new_key = slice_tensor(key, start=start, end=end, dim=k_seq_dim)
        new_value = slice_tensor(value, start=start, end=end, dim=v_seq_dim)
        new_activations.append([new_key, new_value])
    return new_activations

def cat_activations(list_of_activations, k_seq_dim=2, v_seq_dim=2):
    assert all(len(x) == len(list_of_activations[0]) for x in list_of_activations), f"Make sure all activations have the same number of layers! Found {[len(x) for x in list_of_activations]}."

    new_activations = []
    for layer_idx in range(len(list_of_activations[0])):
        keys = [x[layer_idx][0] for x in list_of_activations]
        values = [x[layer_idx][1] for x in list_of_activations]

        new_key = cat_tensor(keys, dim=k_seq_dim)
        new_value = cat_tensor(values, dim=v_seq_dim)
        new_activations.append([new_key, new_value])
    return new_activations

def interleave_activations(main_activations, augment_activations, main_spans, augment_spans, k_seq_dim=2, v_seq_dim=2, device=torch.device("cuda")):
    """ Interleave main_activations and augment_activations according to main_span and augment_span.

    Args:
        main_span: a list of tuples (start_idx, end_idx). when start_idx and end_idx is None, the augment_activations will be plugged in.
        augment_span: a list of tuples (start_idx, end_idx)
    """
    assert len(main_activations) == len(augment_activations) , f"Make sure main and augment activations have the same number of layers! Found {len(main_activations)} and {len(augment_activations)}!"
    assert sum(x[0] is None and x[1] is None for x in main_spans) == len(augment_spans), f"Make sure the number of slots for augmentation (start_idx=None and end_idx=None in main_spans) matches the number of augmentations. Found {sum(x for x in main_spans if x[0] is None and x[1] is None)} slots but {len(augment_spans)} augmentations!"

    new_activations = []
    for layer_idx in range(len(main_activations)):
        main_key, main_value = main_activations[layer_idx]
        augment_key, augment_value = augment_activations[layer_idx]

        sliced_keys = []
        sliced_values = []

        augment_idx = 0
        for start, end in main_spans:
            if start is None and end is None:
                # this means the augment key/value should be plugged in
                augment_start, augment_end = augment_spans[augment_idx]
                sliced_key = slice_tensor(
                    augment_key, 
                    start=augment_start, 
                    end=augment_end,
                    dim=k_seq_dim
                ).to(device)
                sliced_value = slice_tensor(
                    augment_value, 
                    start=augment_start, 
                    end=augment_end,
                    dim=v_seq_dim
                ).to(device)

            else:
                sliced_key = slice_tensor(
                    main_key, 
                    start=start, 
                    end=end,
                    dim=k_seq_dim
                )
                sliced_value = slice_tensor(
                    main_value, 
                    start=start, 
                    end=end,
                    dim=v_seq_dim
                )

            sliced_keys.append(sliced_key)
            sliced_values.append(sliced_value)

        new_key = cat_tensor(sliced_keys, dim=k_seq_dim)
        new_value = cat_tensor(sliced_values, dim=v_seq_dim)
        new_activations.append([new_key, new_value])

    return new_activations

def softmax(x:np.ndarray, axis=-1, temperature=1):
    if isinstance(x, list):
        x = np.array(x)
    x = x / temperature
    x = x - x.max(axis=axis, keepdims=True)
    y = np.exp(x)
    return y / y.sum(axis=axis, keepdims=True)

def l1_norm(x):
    sum_x = sum(x)
    x = [y/sum_x for y in x]
    return x