import math
from typing import Optional, Tuple

import pdb
import torch
from torch import nn
import torch.utils.checkpoint

import torch.nn.functional as F

# from transformers.models.llama.configuration_llama import LlamaConfig
import tqdm
from transformers.models.llama.modeling_llama import (
    LlamaAttention,
    rotate_half,
    apply_rotary_pos_emb,
    LlamaRotaryEmbedding,
    apply_rotary_pos_emb,
    # LlamaForCausalLM,
)
from hip.models.modeling_llama_legacy import (
    LlamaLinearScalingRotaryEmbedding, 
    LlamaDynamicNTKScalingRotaryEmbedding,
)
from hip.models.modeling_llama import (
    # LlamaLinearScalingRotaryEmbedding, 
    # LlamaDynamicNTKScalingRotaryEmbedding,
    LlamaForCausalLM
)
import types

from transformers.utils import (
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    is_flash_attn_greater_or_equal_2_10,
    logging,
    replace_return_docstrings,
)

__all__ = ["H2OLlamaForCausalLM", "H2OLlamaAttention",
            'H2OLlamaAttention_streaming', 'H2OLlamaForCausalLM_streaming']

logger = logging.get_logger(__name__)

from transformers.configuration_utils import PretrainedConfig

LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP = {}

class LlamaConfig(PretrainedConfig):
    r"""
    This is the configuration class to store the configuration of a [`LlamaModel`]. It is used to instantiate an LLaMA
    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
    defaults will yield a similar configuration to that of the LLaMA-7B.

    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
    documentation from [`PretrainedConfig`] for more information.


    Args:
        vocab_size (`int`, *optional*, defaults to 32000):
            Vocabulary size of the LLaMA model. Defines the number of different tokens that can be represented by the
            `inputs_ids` passed when calling [`LlamaModel`]
        hidden_size (`int`, *optional*, defaults to 4096):
            Dimension of the hidden representations.
        intermediate_size (`int`, *optional*, defaults to 11008):
            Dimension of the MLP representations.
        num_hidden_layers (`int`, *optional*, defaults to 32):
            Number of hidden layers in the Transformer decoder.
        num_attention_heads (`int`, *optional*, defaults to 32):
            Number of attention heads for each attention layer in the Transformer decoder.
        num_key_value_heads (`int`, *optional*):
            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
            `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
            by meanpooling all the original heads within that group. For more details checkout [this
            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
            `num_attention_heads`.
        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
            The non-linear activation function (function or string) in the decoder.
        max_position_embeddings (`int`, *optional*, defaults to 2048):
            The maximum sequence length that this model might ever be used with. Llama 1 supports up to 2048 tokens,
            Llama 2 up to 4096, CodeLlama up to 16384.
        initializer_range (`float`, *optional*, defaults to 0.02):
            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
            The epsilon used by the rms normalization layers.
        use_cache (`bool`, *optional*, defaults to `True`):
            Whether or not the model should return the last key/values attentions (not used by all models). Only
            relevant if `config.is_decoder=True`.
        pad_token_id (`int`, *optional*):
            Padding token id.
        bos_token_id (`int`, *optional*, defaults to 1):
            Beginning of stream token id.
        eos_token_id (`int`, *optional*, defaults to 2):
            End of stream token id.
        pretraining_tp (`int`, *optional*, defaults to 1):
            Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
            document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is
            necessary to ensure exact reproducibility of the pretraining results. Please refer to [this
            issue](https://github.com/pytorch/pytorch/issues/76232).
        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
            Whether to tie weight embeddings
        rope_theta (`float`, *optional*, defaults to 10000.0):
            The base period of the RoPE embeddings.
        rope_scaling (`Dict`, *optional*):
            Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
            strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
            `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
            `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
            these scaling strategies behave:
            https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
            experimental feature, subject to breaking API changes in future versions.
        attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
            Whether to use a bias in the query, key, value and output projection layers during self-attention.
        attention_dropout (`float`, *optional*, defaults to 0.0):
            The dropout ratio for the attention probabilities.

    ```python
    >>> from transformers import LlamaModel, LlamaConfig

    >>> # Initializing a LLaMA llama-7b style configuration
    >>> configuration = LlamaConfig()

    >>> # Initializing a model from the llama-7b style configuration
    >>> model = LlamaModel(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```"""

    model_type = "llama"
    keys_to_ignore_at_inference = ["past_key_values"]

    def __init__(
        self,
        vocab_size=32000,
        hidden_size=4096,
        intermediate_size=11008,
        num_hidden_layers=32,
        num_attention_heads=32,
        num_key_value_heads=None,
        hidden_act="silu",
        max_position_embeddings=2048,
        initializer_range=0.02,
        rms_norm_eps=1e-6,
        use_cache=True,
        pad_token_id=None,
        bos_token_id=1,
        eos_token_id=2,
        pretraining_tp=1,
        tie_word_embeddings=False,
        rope_theta=10000.0,
        rope_scaling=None,
        attention_bias=False,
        attention_dropout=0.0,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.max_position_embeddings = max_position_embeddings
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads

        # for backward compatibility
        if num_key_value_heads is None:
            num_key_value_heads = num_attention_heads

        self.num_key_value_heads = num_key_value_heads
        self.hidden_act = hidden_act
        self.initializer_range = initializer_range
        self.rms_norm_eps = rms_norm_eps
        self.pretraining_tp = pretraining_tp
        self.use_cache = use_cache
        self.rope_theta = rope_theta
        self.rope_scaling = rope_scaling
        self._rope_scaling_validation()
        self.attention_bias = attention_bias
        self.attention_dropout = attention_dropout

        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )

    def _rope_scaling_validation(self):
        """
        Validate the `rope_scaling` configuration.
        """
        if self.rope_scaling is None:
            return

        if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
            raise ValueError(
                "`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
                f"got {self.rope_scaling}"
            )
        rope_scaling_type = self.rope_scaling.get("type", None)
        rope_scaling_factor = self.rope_scaling.get("factor", None)
        if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
            raise ValueError(
                f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
            )
        if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
            raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}")



def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
    """
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

def _make_causal_mask(
    bsz: int, tgt_len: int, past_key_values_length: int, dtype: torch.dtype, device: torch.device):
    """
    Make causal mask used for bi-directional self-attention.
    """
    mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
    mask_cond = torch.arange(mask.size(-1), device=device)
    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
    mask = mask.to(dtype)

    if past_key_values_length > 0:
        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)

def apply_rotary_pos_emb_single(x, cos, sin, position_ids, unsqueeze_dim=1, is_decoding=False):
    # cos, sin : [bsz, seq_len, dim], position_ids : [bsz, seq_len]
    # Gather values from a according to position_ids along the T dimension (dim=1)
    _, _, cos_dim = cos.shape
    bsz, pos_t = position_ids.shape
    assert sin.shape[-1] == cos_dim
    assert x.shape[-2] == pos_t

    if not is_decoding:
        cos = torch.gather(cos, 1, position_ids.unsqueeze(-1).expand(bsz, pos_t, cos_dim)).unsqueeze(1)  # [bs, 1, seq_len, dim]
        sin = torch.gather(sin, 1, position_ids.unsqueeze(-1).expand(bsz, pos_t, cos_dim)).unsqueeze(1)  # [bs, 1, seq_len, dim]
    else:
        cos = cos.unsqueeze(unsqueeze_dim)
        sin = sin.unsqueeze(unsqueeze_dim)
    
    x_embed = (x * cos) + (rotate_half(x) * sin) # [bs, head, pos_t, dim]
    return x_embed

import os
class H2OKVCache_LayerWise:
    def __init__(
        self,
        hh_size=4,
        recent_size=512,
        k_seq_dim=2,
        v_seq_dim=2,
    ):
        # print(f"H2OKVCache-LayerWise: {hh_size}, {recent_size}")
        self.hh_size = hh_size
        self.recent_size = recent_size
        self.cache_size = hh_size + recent_size
        self.k_seq_dim = k_seq_dim
        self.v_seq_dim = v_seq_dim
        self.hh_score = None

    def __call__(
        self, 
        past_key_values, 
        attn_score_cache, 
        num_key_value_groups=None, 
        reduction_for_gqa=None, 
        layer_idx=None, 
        cache_position=None,
        num_new_tokens=None,
    ): # , hh_score=None
        self._update_hh_score(
            attn_score_cache, 
            num_key_value_groups, 
            reduction_for_gqa,
            num_new_tokens=num_new_tokens,
        ) # , hh_score
        
        if past_key_values is None or len(past_key_values) <= layer_idx: # TODO purpose?
            return (False, None)
        
        # seq_len = past_key_values[layer_idx][0].size(self.k_seq_dim)
        bsz, num_heads, seq_len, head_dim = past_key_values[layer_idx][0].shape
        
        if seq_len <= self.cache_size: # TODO check
            return (False, past_key_values)
        
        # hh-selection

        select_hh_scores = self.hh_score[:, :, :seq_len - self.recent_size]
        _, keep_topk = torch.topk(select_hh_scores, self.hh_size, dim=-1)
        keep_topk = keep_topk.sort().values

        # keep_recent = torch.arange(seq_len - self.recent_size, seq_len).expand(keep_topk.shape[0], 1).to(keep_topk.device)
        keep_recent = torch.arange(seq_len - self.recent_size, seq_len, device=keep_topk.device).repeat(keep_topk.shape[0], keep_topk.shape[1], 1)
        keep_idx = torch.cat([keep_topk, keep_recent], dim=-1)

        k_hh_recent = torch.gather(past_key_values.key_cache[layer_idx][:, :, :seq_len, :], -2, (keep_idx.unsqueeze(-1)).expand(-1, -1, -1, head_dim))
        v_hh_recent = torch.gather(past_key_values.value_cache[layer_idx][:, :, :seq_len, :], -2, keep_idx.unsqueeze(-1).expand(-1, -1, -1, head_dim))
        
        
        self.hh_score = torch.gather(self.hh_score, -1, keep_idx.expand(bsz, num_heads, self.cache_size))

        return (True, k_hh_recent, v_hh_recent) # , self.hh_score

    def evict_for_space(self, past_key_values, num_coming):
        raise Exception()
        if past_key_values is None:
            return None
        seq_len = past_key_values[0][0].size(self.k_seq_dim)
        if seq_len + num_coming <= self.cache_size:
            return past_key_values

        # hh-selection
        bsz, num_heads, _, head_dim = past_key_values[0].shape

        select_hh_scores = self.hh_score[:, :seq_len - self.recent_size + num_coming]
        _, keep_topk = torch.topk(select_hh_scores, self.hh_size, dim=-1)
        keep_topk = keep_topk.sort().values

        # keep_recent = torch.arange(seq_len - self.recent_size, seq_len).expand(keep_topk.shape[0], 1).to(keep_topk.device)
        keep_recent = torch.arange(seq_len - self.recent_size + num_coming, seq_len, device=keep_topk.device).repeat(keep_topk.shape[0], 1)
        keep_idx = torch.cat([keep_topk, keep_recent], dim=-1)

        mask = torch.zeros(self.hh_score.shape, dtype=torch.bool).to(past_key_values[0].device)
        mask = mask.scatter(-1, keep_idx, 1)

        k_hh_recent = past_key_values[0].squeeze()[mask].view(bsz, num_heads, -1, head_dim)
        v_hh_recent = past_key_values[1].squeeze()[mask].view(bsz, num_heads, -1, head_dim)

        self.hh_score= self.hh_score[mask].view(num_heads, self.cache_size)

        return (k_hh_recent, v_hh_recent)

    def _update_hh_score(self, attn_score_cache, num_key_value_groups=None, reduction_for_gqa=None, num_new_tokens=None): # , hh_score
        # num_new_tokens = attn_score_cache.shape[2]
        assert num_new_tokens is not None
        assert attn_score_cache.ndim == 4
        attn_score_cache = attn_score_cache.sum(2)
        
        N, H, T = attn_score_cache.shape
        assert H % num_key_value_groups == 0
        
        if num_key_value_groups > 1:
            if reduction_for_gqa == 'first_head':
                attn_score_cache = attn_score_cache[:, ::num_key_value_groups, :]
            elif reduction_for_gqa == 'average':
                attn_score_cache = attn_score_cache.view(N, H//num_key_value_groups, num_key_value_groups, T)
                attn_score_cache = torch.mean(attn_score_cache, dim=2)
            elif reduction_for_gqa == 'max_total':
                b = torch.sum(attn_score_cache, dim=2)  # Shape of b is [N, H]
                k = H//num_key_value_groups  # Specify the value of k
                _, topk_indices = torch.topk(b, k, dim=1)  # Shape of topk_indices is [N, k]

                topk_indices_expanded = topk_indices.unsqueeze(2).expand(-1, -1, T)  # Expand indices for proper gathering
                attn_score_cache = torch.gather(attn_score_cache, 1, topk_indices_expanded)  # Shape of attn_score_cache is [N, k, T]

            elif reduction_for_gqa == 'max_group':
                b = torch.sum(attn_score_cache, dim=2)  # Shape of b is [N, H]
                k = H//num_key_value_groups  # Specify the value of k (number of groups)

                # Reshape b to [N, k, H//k]
                b_reshaped = b.view(N, k, num_key_value_groups)  # Shape of b_reshaped is [N, k, H//k]

                # Use torch.topk to get the indices of the top values within each group
                _, topk_indices = torch.topk(b_reshaped, 1, dim=2)  # Shape of topk_indices is [N, k, 1]

                topk_indices_expanded = topk_indices.squeeze(2).unsqueeze(2).expand(-1, -1, T)  # Shape [N, k, T]

                # Gather the topk heads from the original tensor a
                attn_score_cache = torch.gather(attn_score_cache.view(N, k, num_key_value_groups, T), 2, topk_indices_expanded.unsqueeze(2)).squeeze(2)

        if self.hh_score is None:
            self.hh_score = attn_score_cache
        else:
            attn_score_cache[:, :, :-num_new_tokens] += self.hh_score # BSZ, H, T
            self.hh_score = attn_score_cache
                
    def _clean_scores(self):
        self.hh_score = None


class H2OLlamaAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None): # layer_idx = 0?
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.hidden_size // self.num_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.max_position_embeddings = config.max_position_embeddings
        self.rope_theta = config.rope_theta

        if (self.head_dim * self.num_heads) != self.hidden_size:
            raise ValueError(
                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
                f" and `num_heads`: {self.num_heads})."
            )
        
        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=self.config.attention_bias)
        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.config.attention_bias)
        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.config.attention_bias)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=self.config.attention_bias)
        
        self._init_rope()

        self.kv_cache = H2OKVCache_LayerWise(
            hh_size=config.hh_size,
            recent_size=config.recent_size,
            k_seq_dim=2,
            v_seq_dim=2,
        )

    def _init_rope(self):
        if self.config.rope_scaling is None:
            self.rotary_emb = LlamaRotaryEmbedding(
                self.head_dim,
                max_position_embeddings=self.max_position_embeddings,
                base=self.rope_theta,
            )
        else:
            if self.config.rope_scaling is not None:
                self.rope_type = self.config.rope_scaling.get("rope_type", self.config.rope_scaling.get("type"))
            else:
                self.rope_type = "default"
            if self.rope_type == 'llama3':
                self.rope_type = 'linear'
            scaling_factor = self.config.rope_scaling["factor"]
            if self.rope_type == "linear":
                self.rotary_emb = LlamaLinearScalingRotaryEmbedding(
                    self.head_dim,
                    max_position_embeddings=self.max_position_embeddings,
                    scaling_factor=scaling_factor,
                    base=self.rope_theta,
                )
            elif self.rope_type == "dynamic":
                self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding(
                    self.head_dim,
                    max_position_embeddings=self.max_position_embeddings,
                    scaling_factor=scaling_factor,
                    base=self.rope_theta,
                )
            else:
                raise ValueError(f"Unknown RoPE scaling type {self.rope_type}")

    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

    def _clean_cache(self):
        self.kv_cache._clean_scores()
    
    def _h2o_attention_itself(
        self,
        query_states,
        key_states,
        value_states,
        
        bsz, num_heads, head_dim, q_len, kv_seq_len,
        attention_mask,
        
        past_key_value, 
        num_key_value_groups,
        reduction_for_gqa,
        layer_idx,
    ):            
        
        # upcast attention to fp32
        
        N, HEAD, TDST, HID = query_states.shape
        N, HEAD_KV, TSRC, HID = key_states.shape
        assert key_states.shape == value_states.shape
        
        attn_weight_accumulator = torch.zeros(
            (N, HEAD, 1, TSRC), 
            dtype=torch.float32, 
            device=query_states.device
        )
        
        chunk_count = math.ceil((TDST * TSRC) / (2040 * 2048))
        chunk_size = math.ceil(TDST / chunk_count)
        
        attn_output_final = torch.empty_like(query_states)
        
        for i_tdst_start in tqdm.tqdm(range(0, TDST, chunk_size), leave=False, delay=1):
            i_tdst_end = min(i_tdst_start + chunk_size, TDST)
            attn_weights = torch.matmul(
                query_states[:, :, i_tdst_start:i_tdst_end], 
                key_states[:, :, :i_tdst_end + TSRC - TDST].transpose(2, 3)
            ) / math.sqrt(head_dim)
            if TDST > 1:
                idx_tdst = torch.arange(i_tdst_start, i_tdst_end, device=query_states.device)
                idx_tsrc = torch.arange(0, i_tdst_end, device=query_states.device) + TSRC - TDST
                attn_weights = torch.where(
                    idx_tsrc[None, None, None, :] <= idx_tdst[None, None, :, None],
                    attn_weights,
                    -32000.0
                )
            attn_weights = nn.functional.softmax(
                attn_weights,
                dim=-1, 
                dtype=torch.float32
            ).to(query_states.dtype)
            attn_weight_accumulator[:, :, :, :i_tdst_end + TSRC - TDST].add_(attn_weights.sum(2, keepdim=True))
            attn_output = torch.matmul(attn_weights, value_states[:, :, :i_tdst_end + TSRC - TDST])
            
            attn_output_final.index_copy_(
                dim=2,
                index=torch.arange(i_tdst_start, i_tdst_end, device=attn_output.device),
                source=attn_output.to(attn_output_final.dtype)
            )
            # attn_output_final = attn_output

        attn_weights = attn_weight_accumulator
        attn_output = attn_output_final
        
        kv_hh = self.kv_cache(
            past_key_value, 
            attn_weights, 
            num_key_value_groups, 
            reduction_for_gqa, 
            layer_idx,
            num_new_tokens=TDST,
        ) # , hh_score TODO check
                
        if kv_hh[0] == True:
            _, k_hh_recent, v_hh_recent = kv_hh
            
            past_key_value.key_cache[layer_idx] = k_hh_recent
            past_key_value.value_cache[layer_idx] = v_hh_recent
        
        return attn_output, None
    
    def _h2o_attention(
        self,
        query_states,
        key_states,
        value_states,
        
        position_ids,
        past_key_value,
        output_attentions,
        use_cache,
        # hh_score,
        
        bsz,
        cache_position,
        reduction_for_gqa=None,
        kv_seq_len=None,
        decoding_loop_for_prefill=True,
        compute_final_attn_output=True,
        
        cos=None,
        sin=None,
        i=None
        # position_embeddings=None,
    ):
        assert use_cache == True
        q_len = query_states.shape[-2]
        
        if decoding_loop_for_prefill:
            attention_mask = None
            
            kv_seq_len = key_states.shape[-2]
        
            if past_key_value is not None and len(past_key_value) > self.layer_idx: # TODO check this for decoding
                kv_seq_len += past_key_value[self.layer_idx][0].shape[-2]
            
            """
            NOTE: H2O has 3 variants
            1. No touch on RoPE: official implementation
            2. shift query position of RoPE by min(k, pos idx): Implementation details in the official implementation that not API available. They describe this only with code comment
            3. StreamingLLM style RoPE: my own implementation
            """

            if self.config.h2o_streaming:
                if past_key_value is not None: # and len(past_key_value) != 0:
                    # sin and cos are specific to RoPE models; cache_position needed for the static cache
                    cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
                    key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) # TODO check past_key_value update
                    # reuse k, v, self_attention
                    # key_states = torch.cat([past_key_value[0], key_states], dim=2)
                    # value_states = torch.cat([past_key_value[1], value_states], dim=2)
                # past_key_value = (key_states, value_states) if use_cache else None
                
                position_ids = torch.arange(0, key_states.shape[-2], device=key_states.device)[None, :]
                
                # cos, sin = self.rotary_emb(value_states, seq_len=position_length)
                # NOTE: grab all position embeddings
                # cos, sin = self.rotary_emb(value_states, position_ids)
                
                ### Shift Pos: query pos is min(cache_size, idx)
                query_states = apply_rotary_pos_emb_single(
                    query_states, 
                    cos, 
                    sin, 
                    position_ids[:, -q_len:],
                    is_decoding=self.config.is_decoding
                )
                key_states = apply_rotary_pos_emb_single(
                    key_states, 
                    cos, 
                    sin, 
                    position_ids,
                    is_decoding=self.config.is_decoding
                )
            else:
                # cos, sin = self.rotary_emb(value_states, seq_len=position_length)
                # NOTE: grab all position embeddings
            
                ### Shift Pos: query pos is min(cache_size, idx)
                query_states = apply_rotary_pos_emb_single(
                    query_states, 
                    cos, 
                    sin, 
                    torch.clamp_max(position_ids, kv_seq_len) if self.config.h2o_shift_q_pos else position_ids,
                    is_decoding=self.config.is_decoding
                )
                key_states = apply_rotary_pos_emb_single(
                    key_states, 
                    cos, 
                    sin, 
                    position_ids,
                    is_decoding=self.config.is_decoding
                )
                if past_key_value is not None: # and len(past_key_value) != 0:
                    # sin and cos are specific to RoPE models; cache_position needed for the static cache
                    cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
                    key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) # TODO check past_key_value update
                    # reuse k, v, self_attention
                    # key_states = torch.cat([past_key_value[0], key_states], dim=2)
                    # value_states = torch.cat([past_key_value[1], value_states], dim=2)
                # past_key_value = (key_states, value_states) if use_cache else None

            key_states = repeat_kv(key_states, self.num_key_value_groups)
            value_states = repeat_kv(value_states, self.num_key_value_groups)

        attn_output, attn_weights= self._h2o_attention_itself(
            query_states,
            key_states,
            value_states,
            
            bsz, self.num_heads, self.head_dim, q_len, kv_seq_len,
            attention_mask,
            
            past_key_value, 
            self.num_key_value_groups,
            reduction_for_gqa,
            self.layer_idx,
        )
        
        if compute_final_attn_output:
            if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
                raise ValueError(
                    f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
                    f" {attn_output.size()}"
                )
                
            attn_output = attn_output.transpose(1, 2).contiguous()
            attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)

            if self.config.pretraining_tp > 1:
                attn_output = attn_output.split(
                    self.hidden_size // self.config.pretraining_tp, dim=2
                )
                o_proj_slices = self.o_proj.weight.split(
                    self.hidden_size // self.config.pretraining_tp, dim=1
                )
                attn_output = sum(
                    [
                        F.linear(attn_output[i], o_proj_slices[i])
                        for i in range(self.config.pretraining_tp)
                    ]
                )
            else:
                attn_output = self.o_proj(attn_output)

            if not output_attentions:
                attn_weights = None
        
        return attn_output, attn_weights, past_key_value # , hh_score
    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        cache_position: Optional[torch.Tensor] = None,
        position_embeddings: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        
        bsz, q_len, _ = hidden_states.size()

        if self.config.pretraining_tp > 1:
            key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp
            query_slices = self.q_proj.weight.split(
                (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0
            )
            key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
            value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)

            query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]
            query_states = torch.cat(query_states, dim=-1)

            key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]
            key_states = torch.cat(key_states, dim=-1)

            value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]
            value_states = torch.cat(value_states, dim=-1)

        else:
            query_states = self.q_proj(hidden_states)
            key_states = self.k_proj(hidden_states)
            value_states = self.v_proj(hidden_states)

        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
        
        assert self.config.hh_size == self.config.tree_k//2
        assert self.config.recent_size == self.config.tree_k//2
        assert self.config._attn_implementation == self.config.attn_implementation == 'eager'
        assert self.config.h2o_shift_q_pos is not None
        assert self.config.reduction_for_gqa is not None
        
        if self.config.attention_method == 'h2o_stream':
            assert self.config.h2o_streaming == True
        assert use_cache == True
        
        mask_k = self.config.tree_k
        
        if position_embeddings is None:
            logger.warning_once(
                "The attention layers in this model are transitioning from computing the RoPE embeddings internally "
                "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed "
                "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.45 `position_ids` will be "
                "removed and `position_embeddings` will be mandatory."
            )
            cos, sin = self.rotary_emb(value_states, position_ids)
        else:
            cos, sin = position_embeddings
        
        if not self.config.is_decoding:
            position_ids_k = position_ids[:, :mask_k]
            position_ids_loop = position_ids[:, mask_k:]
            
            query_states_k = query_states[:, :, :mask_k, :]
            query_states_loop = query_states[:, :, mask_k:, :]
            key_states_k = key_states[:, :, :mask_k, :]
            key_states_loop = key_states[:, :, mask_k:, :]
            value_states_k = value_states[:, :, :mask_k, :]
            value_states_loop = value_states[:, :, mask_k:, :]
            
            # assert past_key_value is None # TODO CHECK
            assert use_cache is True
            
            attn_output, attn_weights, past_key_value = self._h2o_attention( # , hh_score
                query_states_k,
                key_states_k,
                value_states_k,
                
                position_ids_k,
                past_key_value=past_key_value,
                output_attentions=output_attentions,
                use_cache=use_cache,
                # hh_score=hh_score,
                bsz=bsz,
                cache_position=cache_position,
                reduction_for_gqa=self.config.reduction_for_gqa,
                kv_seq_len=None,
                decoding_loop_for_prefill=True,
                compute_final_attn_output=True,
                
                cos=cos,
                sin=sin,
            )
            attn_output_loop = torch.zeros(
                (attn_output.shape[0], q_len - mask_k, attn_output.shape[-1]), 
                dtype=attn_output.dtype, device=attn_output.device
            )
            
            # loop one by one
            assert query_states_loop.shape[-2] == q_len - mask_k
            for i in range(q_len - mask_k):
                # print(f'>> loop {i}')
                attn_output_, attn_weights_, past_key_value = self._h2o_attention( # , hh_score
                    query_states_loop[:, :, i, :][:, :, None, :],
                    key_states_loop[:, :, i, :][:, :, None, :],
                    value_states_loop[:, :, i, :][:, :, None, :],
                    
                    position_ids_loop[:, i][:, None],
                    past_key_value=past_key_value,
                    output_attentions=output_attentions,
                    use_cache=use_cache,
                    # hh_score=hh_score,
                    bsz=bsz,
                    cache_position=cache_position,
                    reduction_for_gqa=self.config.reduction_for_gqa,
                    kv_seq_len=None,
                    decoding_loop_for_prefill=True,
                    compute_final_attn_output=True,
                    
                    cos=cos,
                    sin=sin,
                    i=i
                )
                
                attn_output_loop[:, i:i+1, :].copy_(attn_output_, non_blocking=True)
                
            attn_output = torch.cat((attn_output, attn_output_loop), dim=1)
        
            if output_attentions:
                raise Exception()
                attn_weights = torch.cat((attn_weights, attn_weigth_loop), dim=-2)
            
        else:
            assert use_cache is True
            # compute_final_attn_output = False
            
            attn_output, attn_weights, past_key_value = self._h2o_attention( # , hh_score
                query_states,
                key_states,
                value_states,
                
                position_ids,
                past_key_value=past_key_value,
                output_attentions=output_attentions,
                use_cache=use_cache,
                # hh_score=hh_score,
                bsz=bsz,
                cache_position=cache_position,
                reduction_for_gqa=self.config.reduction_for_gqa,
                kv_seq_len=None,
                decoding_loop_for_prefill=True,
                compute_final_attn_output=True, # compute_final_attn_output
                
                cos=cos,
                sin=sin
            )
        return attn_output, attn_weights, past_key_value
    
from hip.models.modeling_llama import LlamaCustomAttention

class H2OLlamaForCausalLM(LlamaForCausalLM):
    def __init__(self, config):
        super().__init__(config)
        num_layers = len(self.model.layers)
        for layer_idx in range(num_layers):
            if layer_idx in config.tree_dense_layers:
                self.model.layers[layer_idx].self_attn = LlamaCustomAttention(config, layer_idx=layer_idx)
            else:
                self.model.layers[layer_idx].self_attn = H2OLlamaAttention(config, layer_idx=layer_idx)

## H2O KV Cache dropping with Position rolling
class H2OLlamaAttention_streaming(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(self, config: LlamaConfig):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.hidden_size // self.num_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.max_position_embeddings = config.max_position_embeddings
        self.rope_theta = config.rope_theta

        if (self.head_dim * self.num_heads) != self.hidden_size:
            raise ValueError(
                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
                f" and `num_heads`: {self.num_heads})."
            )
        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
        self._init_rope()

        self.kv_cache = H2OKVCache_LayerWise(
            hh_size=config.hh_size,
            recent_size=config.recent_size,
            k_seq_dim=2,
            v_seq_dim=2,
        )

    def _init_rope(self):
        if self.config.rope_scaling is None:
            self.rotary_emb = LlamaRotaryEmbedding(
                self.head_dim,
                max_position_embeddings=self.max_position_embeddings,
                base=self.rope_theta,
            )
        else:
            scaling_type = self.config.rope_scaling["type"]
            scaling_factor = self.config.rope_scaling["factor"]
            if scaling_type == "linear":
                self.rotary_emb = LlamaLinearScalingRotaryEmbedding(
                    self.head_dim,
                    max_position_embeddings=self.max_position_embeddings,
                    scaling_factor=scaling_factor,
                    base=self.rope_theta,
                )
            elif scaling_type == "dynamic":
                self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding(
                    self.head_dim,
                    max_position_embeddings=self.max_position_embeddings,
                    scaling_factor=scaling_factor,
                    base=self.rope_theta,
                )
            else:
                raise ValueError(f"Unknown RoPE scaling type {scaling_type}")

    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

    def _clean_cache(self):
        self.kv_cache._clean_scores()

    def forward(
        self,
        hidden_states: torch.Tensor,
        query_states: torch.Tensor=None,
        key_states: torch.Tensor=None,
        value_states: torch.Tensor=None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        h2o_shift_q_pos: bool = False,
        mask_k: int = 512,
        reduction_for_gqa: str = 'average',
        cache_position: Optional[torch.Tensor] = None,
        position_embeddings: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:

        bsz, q_len, _ = hidden_states.size()

        if self.config.pretraining_tp > 1:
            key_value_slicing = (
                self.num_key_value_heads * self.head_dim
            ) // self.config.pretraining_tp
            query_slices = self.q_proj.weight.split(
                (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0
            )
            key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
            value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)

            query_states = [
                F.linear(hidden_states, query_slices[i])
                for i in range(self.config.pretraining_tp)
            ]
            query_states = torch.cat(query_states, dim=-1)

            key_states = [
                F.linear(hidden_states, key_slices[i])
                for i in range(self.config.pretraining_tp)
            ]
            key_states = torch.cat(key_states, dim=-1)

            value_states = [
                F.linear(hidden_states, value_slices[i])
                for i in range(self.config.pretraining_tp)
            ]
            value_states = torch.cat(value_states, dim=-1)

        else:
            query_states = self.q_proj(hidden_states)
            key_states = self.k_proj(hidden_states)
            value_states = self.v_proj(hidden_states)

        query_states = query_states.view(
            bsz, q_len, self.num_heads, self.head_dim
        ).transpose(1, 2)
        key_states = key_states.view(
            bsz, q_len, self.num_key_value_heads, self.head_dim
        ).transpose(1, 2)
        value_states = value_states.view(
            bsz, q_len, self.num_key_value_heads, self.head_dim
        ).transpose(1, 2)

        # remake causal mask
        attention_mask = _make_causal_mask(
            bsz=bsz,
            tgt_len=q_len,
            past_key_values_length=past_key_value[0].shape[-2] if (past_key_value is not None and len(past_key_value) != 0) else 0,
            dtype=query_states.dtype,
            device=query_states.device,
        )

        kv_seq_len = key_states.shape[-2]
        if past_key_value is not None and len(past_key_value) != 0:
            kv_seq_len += past_key_value[0].shape[-2]

        if not position_ids.nelement() > 1:
            position_ids[0][0] = kv_seq_len - 1

        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
        ### Shift Pos: query pos is min(cache_size, idx)
        # query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
        query_states = apply_rotary_pos_emb_single(query_states, cos, sin, position_ids)
        ###

        if past_key_value is not None and len(past_key_value) != 0:
            # reuse k, v, self_attention
            key_states = torch.cat([past_key_value[0], key_states], dim=2)
            value_states = torch.cat([past_key_value[1], value_states], dim=2)

        past_key_value = (key_states, value_states) if use_cache else None

        ### Shift Pos: key pos is the pos in cache (Rolling KV Cache and using relative pos emb)
        key_position_ids = torch.arange(kv_seq_len, device=position_ids.device).unsqueeze(0)
        key_states = apply_rotary_pos_emb_single(key_states, cos, sin, key_position_ids)
        ###

        # repeat k/v heads if n_kv_heads < n_heads
        key_states = repeat_kv(key_states, self.num_key_value_groups)
        value_states = repeat_kv(value_states, self.num_key_value_groups)

        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(
            self.head_dim
        )

        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
            raise ValueError(
                f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
                f" {attn_weights.size()}"
            )

        if attention_mask is not None:
            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
                )
            attn_weights = attn_weights + attention_mask

        # upcast attention to fp32
        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
            query_states.dtype
        )

        past_key_value = self.kv_cache(past_key_value, attn_weights.detach().clone(), self.num_key_value_groups, reduction_for_gqa, self.layer_idx)

        attn_output = torch.matmul(attn_weights, value_states)

        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
            raise ValueError(
                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
                f" {attn_output.size()}"
            )

        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)

        if self.config.pretraining_tp > 1:
            attn_output = attn_output.split(
                self.hidden_size // self.config.pretraining_tp, dim=2
            )
            o_proj_slices = self.o_proj.weight.split(
                self.hidden_size // self.config.pretraining_tp, dim=1
            )
            attn_output = sum(
                [
                    F.linear(attn_output[i], o_proj_slices[i])
                    for i in range(self.config.pretraining_tp)
                ]
            )
        else:
            attn_output = self.o_proj(attn_output)

        if not output_attentions:
            attn_weights = None

        return attn_output, attn_weights, past_key_value


class H2OLlamaForCausalLM_streaming(LlamaForCausalLM):
    def __init__(self, config):
        super().__init__(config)
        num_layers = len(self.model.layers)
        for layer_idx in range(num_layers):
            self.model.layers[layer_idx].self_attn = H2OLlamaAttention_streaming(config)
