from transformers.models.llama.modeling_llama import *
from beagle.mixin import *


class LlamaBeagleAttention(LlamaAttention, BeagleAttentionMixin):

    def 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,
        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        **kwargs,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:

        bsz, q_len, _ = hidden_states.size()
        query_states, key_states, value_states = self.qkv_transform(
            hidden_states, past_key_value, use_cache, position_embeddings, **kwargs)

        ################################################
        ### everything kept original starting from here
        ################################################

        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 attention_mask is not None:  # no matter the length, we just slice it
            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
            attn_weights = attn_weights + causal_mask

        # upcast attention to fp32
        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
        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, -1)

        attn_output = self.o_proj(attn_output)

        if not output_attentions:
            attn_weights = None

        return attn_output, attn_weights, past_key_value


class LlamaBeagleDecoderLayer(LlamaDecoderLayer):
    def __init__(self, config, layer_id):
        super().__init__(config, layer_id)

        if not config.beagle_use_fc_eagle:
            delattr(self, 'self_attn')
            recycle_vram()
            self.self_attn = LlamaBeagleAttention(
                config=config, layer_idx=layer_id
            )


class LlamaForSpeculativeCausalLM(LlamaForCausalLM, BeagleMixin):
    _no_split_modules = ["LlamaDecoderLayer", "LlamaBeagleDecoderLayer"]

    def __init__(self, config):
        super().__init__(config)

        BeagleMixin.__init__(self, config)
        self.speculative_decoder = LlamaBeagleDecoderLayer(config, layer_id=0)

        self.post_init()

    def forward(self, *args, **kwargs) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPast]:
        return self.beagle_forward(*args, **kwargs)