import torch
import torch.nn as nn
import torch.F as F
import math

class RMSNorm(nn.Module):
    def __init__(self, hidden_dim, eps = 1e-6):
        self.weight = nn.Parameter(hidden_dim)
        self.eps = eps

    def forward(self, hidden_state):
        norm = hidden_state.pow(2).mean(-1, keepdim = True)
        hidden_state = hidden_state * torch.rsqrt(norm + self.eps) * self.weight
        return hidden_state


class MultiHeadAttention(nn.Module):
    def __init__(self, hidden_dim, num_heads):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.head_dim = hidden_dim // num_heads
        self.q_linear = nn.linear(hidden_dim, hidden_dim)
        self.k_linear = nn.linear(hidden_dim, hidden_dim)
        self.v_linear = nn.linear(hidden_dim, hidden_dim)
        self.o_linear = nn.linear(hidden_dim, hidden_dim)
        self.norm = RMSNorm(hidden_dim)

    def forward(self, hidden_state, mask, past_kv = None, use_cache = True):
        bs = hidden_state.shape[0]
        residual = hidden_state
        hidden_state = self.norm(hidden_state)
        q = self.q_linear(hidden_state)  # (bs, seqlen, hidden_dim)
        k = self.k_linear(hidden_state)  # (bs, seqlen, hidden_dim)
        v = self.v_linear(hidden_state)  # (bs, seqlen, hidden_dim)

        q = q.view(bs, -1, self.num_heads, self.head_dim).transpose(1, 2)  # (bs, num_heads, seqlen, head_dim)
        k = k.view(bs, -1, self.num_heads, self.head_dim).transpose(1, 2)  # (bs, num_heads, seqlen, head_dim)
        v = v.view(bs, -1, self.num_heads, self.head_dim).transpose(1, 2)  # (bs, num_heads, seqlen, head_dim)

        if past_kv is not None:
            past_k, past_v = past_kv
            k = torch.cat([past_k, k], dim = 2)
            v = torch.cat([past_v, v], dim = 2)

        new_past_kv = (k, v) if use_cache else None

        attention_score = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.head_dim)

        attention_score += mask * -1e9

        attention_score = F.softmax(attention_score)
        output = torch.matmul(attention_score, v)

        output = output.transpose(1, 2).view(bs, -1, self.hidden_dim)
        output = self.o_linear(output) + residual

        return output, new_past_kv if use_cache else output