# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import math
from typing import Optional, Tuple

import torch
from einops import rearrange
from transformers.utils import is_torch_npu_available
from transformers.cache_utils import Cache
from transformers.integrations.sdpa_attention import repeat_kv

if is_torch_npu_available():
    import torch_npu
    from .rope import apply_rotary_pos_emb


def sdpa_attention_forward(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: Optional[torch.Tensor],
    dropout: float = 0.0,
    scaling: Optional[float] = None,
    is_causal: Optional[bool] = None,
    **kwargs,
) -> Tuple[torch.Tensor, None]:
    if hasattr(module, "num_key_value_groups"):
        key = repeat_kv(key, module.num_key_value_groups)
        value = repeat_kv(value, module.num_key_value_groups)

    causal_mask = attention_mask
    if attention_mask is not None and causal_mask.ndim == 4:
        causal_mask = causal_mask[:, :, :, : key.shape[-2]]

    # SDPA with memory-efficient backend is bugged with non-contiguous inputs and custom attn_mask for some torch versions
    query = query.contiguous()
    key = key.contiguous()
    value = value.contiguous()

    # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
    # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
    # Note that it is important to check first for the shape, otherwise compile will fail with `argument 'is_causal' must be bool, not SymBool`
    if is_causal is None:
        is_causal = query.shape[2] > 1 and causal_mask is None

    # Shapes (e.g. query.shape[2]) are tensors during jit tracing, resulting in `is_causal` being a tensor.
    # We convert it to a bool for the SDPA kernel that only accepts bools.
    if torch.jit.is_tracing() and isinstance(is_causal, torch.Tensor):
        is_causal = is_causal.item()


    if is_causal:
        atten_mask_npu = torch.triu(torch.ones([2048, 2048]), diagonal=1).bool().to(query.device) if causal_mask else None
        sparse_mode = 2 if causal_mask else 0
        head_num = query.shape[1]
        attn_output = torch_npu.npu_fusion_attention(
            query, key, value, head_num, input_layout="BNSD",
            pse=None,
            atten_mask=atten_mask_npu,
            sparse_mode=sparse_mode,
            scale=1.0 / math.sqrt(query.shape[-1]),
            pre_tockens=2147483647,
            next_tockens=2147483647,
            keep_prob=1,
        )[0]
    else:
        if causal_mask.dtype == torch.bool:
            atten_mask_npu = torch.logical_not(causal_mask.bool()).to(query.device)
        else:
            atten_mask_npu = causal_mask.bool().to(query.device)
        head_num = query.shape[1]
        attn_output = torch_npu.npu_fusion_attention(
            query, key, value, head_num, input_layout="BNSD",
            pse=None,
            atten_mask=atten_mask_npu,
            scale=1.0 / math.sqrt(query.shape[-1]),
            pre_tockens=2147483647,
            next_tockens=2147483647,
            keep_prob=1
        )[0]

    attn_output = attn_output.transpose(1, 2).contiguous()

    return attn_output, None


def internlm2_sdpa_forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
    past_key_value: Optional[Cache] = None,
    output_attentions: bool = False,
    use_cache: bool = False,
    cache_position: Optional[torch.LongTensor] = None,
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
    if output_attentions:
        return super().forward(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
            cache_position=cache_position,
        )

    bsz, q_len, _ = hidden_states.size()

    qkv_states = self.wqkv(hidden_states)

    qkv_states = rearrange(
        qkv_states,
        "b q (h gs d) -> b q h gs d",
        gs=2 + self.num_key_value_groups,
        d=self.head_dim,
    )

    query_states = qkv_states[..., : self.num_key_value_groups, :]
    query_states = rearrange(query_states, "b q h gs d -> b q (h gs) d")
    key_states = qkv_states[..., -2, :]
    value_states = qkv_states[..., -1, :]

    query_states = query_states.transpose(1, 2)
    key_states = key_states.transpose(1, 2)
    value_states = value_states.transpose(1, 2)

    cos, sin = self.rotary_emb(value_states, position_ids)
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

    if past_key_value is not None:
        # 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)

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

    causal_mask = attention_mask
    if attention_mask is not None:
        causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]

    # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with
    # custom attn_mask, Reference: https://github.com/pytorch/pytorch/issues/112577.
    if query_states.device.type == "npu" and causal_mask is not None:
        query_states = query_states.contiguous()
        key_states = key_states.contiguous()
        value_states = value_states.contiguous()

    # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of
    # an inline conditional assignment in SDPA to support both torch.compile's dynamic shapes and full graph
    # options. An inline conditional prevents dynamic shapes from compiling.
    is_causal = bool(causal_mask is None and q_len > 1)

    if is_causal:
        atten_mask_npu = torch.triu(torch.ones([2048, 2048]), diagonal=1).bool().to(query_states.device) if causal_mask else None
        sparse_mode = 2 if causal_mask else 0
        head_num = query_states.shape[1]
        attn_output = torch_npu.npu_fusion_attention(
            query_states, key_states, value_states, head_num, input_layout="BNSD",
            pse=None,
            atten_mask=atten_mask_npu,
            sparse_mode=sparse_mode,
            scale=1.0 / math.sqrt(query_states.shape[-1]),
            pre_tockens=2147483647,
            next_tockens=2147483647,
            keep_prob=1,
        )[0]
    else:
        if causal_mask.dtype == torch.bool:
            atten_mask_npu = torch.logical_not(causal_mask.bool()).to(query_states.device)
        else:
            atten_mask_npu = causal_mask.bool().to(query_states.device)
        head_num = query_states.shape[1]
        attn_output = torch_npu.npu_fusion_attention(
            query_states, key_states, value_states, head_num, input_layout="BNSD",
            pse=None,
            atten_mask=atten_mask_npu,
            scale=1.0 / math.sqrt(query_states.shape[-1]),
            pre_tockens=2147483647,
            next_tockens=2147483647,
            keep_prob=1
        )[0]


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

    attn_output = self.wo(attn_output)

    return attn_output, None, past_key_value
