import torch
import math

from .utils import apply_rope, apply_rope_x

class MHA(torch.nn.Module):

    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.qkv = torch.nn.Parameter(0.01*torch.randn((3*d_model, d_model)))
        self.wo = torch.nn.Parameter(0.01*torch.randn((d_model, d_model)))

    def forward(self, x, kv_cache=None, past_length=0):
        # queries, keys, and values
        B, S, D = x.shape
        QKV = x @ self.qkv.T # B, S, 3D

        Q, K, V = torch.chunk(QKV, 3, -1)

        # split into multiple heads
        dh = D//self.n_heads

        # We do this all at once because the reshape options are super expensive
        q_heads = Q.view(B, S, self.n_heads, dh).transpose(1,2)
        k_heads = K.view(B, S, self.n_heads, dh).transpose(1,2)
        v_heads = V.view(B, S, self.n_heads, dh).transpose(1,2)

        if kv_cache is not None:
            k_cache, v_cache = kv_cache
            k_heads = torch.cat([k_cache, k_heads], dim=2)
            v_heads = torch.cat([v_cache, v_heads], dim=2)

        updated_kv_cache = (k_heads, v_heads)

        S_full = k_heads.size(2)

        # make attention mask
        mask = torch.ones((S,S_full), device=x.device)
        mask = torch.tril(mask, diagonal=past_length)
        mask = mask[None, None, :, :]

        sq_mask = mask == 1
        
        # attention
        x = torch.nn.functional.scaled_dot_product_attention(
            q_heads, k_heads, v_heads,
            attn_mask=sq_mask
        )

        x = x.transpose(1, 2).reshape(B, S, D)

        # apply projection
        x = x @ self.wo.T

        return x, updated_kv_cache

class Rope_MHA(torch.nn.Module):

    def __init__(self, d_model, n_heads, max_len=1024, rope_theta=10000.0):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.dh = d_model // n_heads
        self.qkv = torch.nn.Parameter(0.01*torch.randn((3*d_model, d_model)))
        self.wo = torch.nn.Parameter(0.01*torch.randn((d_model, d_model)))

        # RoPE
        self.max_seq_len = max_len
        self.rope_theta = rope_theta

        # https://github.com/lucidrains/rotary-embedding-torch/tree/main
        # visualize emb later to make sure it looks ok
        freqs = 1.0 / (rope_theta ** (torch.arange(0, self.dh, 2).float() / self.dh))
        emb = torch.outer(torch.arange(self.max_seq_len).float(), freqs)
        cos_cached = emb.cos()[None, None, :, :]
        sin_cached = emb.sin()[None, None, :, :]

        # https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_buffer
        # This is like a parameter but its a constant so we can use register_buffer
        self.register_buffer("cos_cached", cos_cached)
        self.register_buffer("sin_cached", sin_cached)

    def forward(self, x, kv_cache=None, past_length=0):
        # queries, keys, and values
        B, S, D = x.shape
        QKV = x @ self.qkv.T # B, S, 3D

        Q, K, V = torch.chunk(QKV, 3, -1)

        # split into multiple heads
        q_heads = Q.view(B, S, self.n_heads, self.dh).transpose(1,2)
        k_heads = K.view(B, S, self.n_heads, self.dh).transpose(1,2)
        v_heads = V.view(B, S, self.n_heads, self.dh).transpose(1,2)

        ## Apply RoPE
        cos = self.cos_cached[:, :, past_length:past_length+S, :self.dh//2].repeat(1, 1, 1, 2)
        sin = self.sin_cached[:, :, past_length:past_length+S, :self.dh//2].repeat(1, 1, 1, 2)
        q_heads, k_heads = apply_rope(q_heads, k_heads, cos, sin)

        if kv_cache is not None:
            k_cache, v_cache = kv_cache
            k_heads = torch.cat([k_cache, k_heads], dim=2)
            v_heads = torch.cat([v_cache, v_heads], dim=2)

        S_full = k_heads.size(2)

        # make attention mask
        mask = torch.ones((S,S_full), device=x.device)
        mask = torch.tril(mask, diagonal=past_length)
        mask = mask[None, None, :, :]

        sq_mask = mask == 1
        
        # attention
        x = torch.nn.functional.scaled_dot_product_attention(
            q_heads, k_heads, v_heads,
            attn_mask=sq_mask
        )

        x = x.transpose(1, 2).reshape(B, S, D)

        # apply projection
        x = x @ self.wo.T

        return x, (k_heads, v_heads)
    

class Decoupled_Rope_MHA(torch.nn.Module):

    def __init__(self, d_model, n_heads, max_len=1024, rope_theta=10000.0):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.dh = d_model // n_heads

        self.qk_nope_dim = self.dh // 2
        self.qk_rope_dim = self.dh // 2
        
        self.qkv = torch.nn.Parameter(0.01*torch.randn((2*d_model, d_model)))
        self.wk = torch.nn.Parameter(0.01*torch.randn((self.n_heads * self.qk_nope_dim) + self.qk_rope_dim,
                                                      d_model))
        self.wo = torch.nn.Parameter(0.01*torch.randn((d_model, d_model)))

        # RoPE
        self.max_seq_len = max_len
        self.rope_theta = rope_theta

        # https://github.com/lucidrains/rotary-embedding-torch/tree/main
        # visualize emb later to make sure it looks ok
        freqs = 1.0 / (rope_theta ** (torch.arange(0, self.dh, 2).float() / self.dh))
        emb = torch.outer(torch.arange(self.max_seq_len).float(), freqs)
        cos_cached = emb.cos()[None, None, :, :]
        sin_cached = emb.sin()[None, None, :, :]

        # https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_buffer
        # This is like a parameter but its a constant so we can use register_buffer
        self.register_buffer("cos_cached", cos_cached)
        self.register_buffer("sin_cached", sin_cached)

    def forward(self, x, kv_cache=None, past_length=0):
        # queries, keys, and values
        B, S, D = x.shape
        QV = x @ self.qkv.T # B, S, 2D
        K = x @ self.wk.T # B, S, n_heads*nope_dim + rope_dim

        Q, V = torch.chunk(QV, 2, -1)

        # split into multiple heads
        Q = Q.view(B, S, self.n_heads, self.dh).transpose(1,2)
        v_heads = V.view(B, S, self.n_heads, self.dh).transpose(1,2)

        # separate for RoPE
        # note: this one we do after head split to split along the heads
        Q, Q_for_rope = torch.split(Q, [self.qk_nope_dim, self.qk_rope_dim],
                                    dim=-1)
        # note: this one we do before head split so we can separate the rope part
        K, K_for_rope = torch.split(K, [self.n_heads * self.qk_nope_dim, self.qk_rope_dim],
                                    dim=-1)
        K = K.view(B, S, self.n_heads, self.qk_nope_dim).transpose(1,2)
        K_for_rope = K_for_rope.view(B, -1, 1, self.qk_rope_dim).transpose(1,2)

        ## Apply RoPE
        cos_q = self.cos_cached[:, :, past_length:past_length+S,
                                :self.qk_rope_dim//2].repeat(1, 1, 1, 2)
        sin_q = self.sin_cached[:, :, past_length:past_length+S,
                                :self.qk_rope_dim//2].repeat(1, 1, 1, 2)
        Q_for_rope = apply_rope_x(Q_for_rope, cos_q, sin_q)

        if kv_cache is None:
            S_full = K.size(2)
            
            cos_k = self.cos_cached[:, :, :S_full,
                                    :self.qk_rope_dim//2].repeat(1, 1, 1, 2)
            sin_k = self.sin_cached[:, :, :S_full,
                                    :self.qk_rope_dim//2].repeat(1, 1, 1, 2)
            K_for_rope = apply_rope_x(K_for_rope, cos_k, sin_k)

            # repeat for each head
            K_for_rope = K_for_rope.repeat(1, self.n_heads, 1, 1)

        else:
            k_cache, v_cache = kv_cache
            v_heads = torch.cat([v_cache, v_heads], dim=2)

            # decoupled RoPE for k cache
            old_K, old_K_for_rope = torch.split(k_cache, [self.qk_nope_dim, self.qk_rope_dim],
                                                dim=-1)

            K = torch.cat([old_K, K], dim=2)
            S_full = K.size(2)

            cos_k = self.cos_cached[:, :, past_length:past_length+S,
                                    :self.qk_rope_dim//2].repeat(1, 1, 1, 2)
            sin_k = self.sin_cached[:, :, past_length:past_length+S,
                                    :self.qk_rope_dim//2].repeat(1, 1, 1, 2)
            K_for_rope = apply_rope_x(K_for_rope, cos_k, sin_k)

            # repeat for each head
            K_for_rope = K_for_rope.repeat(1, self.n_heads, 1, 1)
            K_for_rope = torch.cat([old_K_for_rope, K_for_rope], dim=2)

        # assemble into heads
        q_heads = torch.cat([Q, Q_for_rope], dim=-1)
        k_heads = torch.cat([K, K_for_rope], dim=-1)

        # make attention mask
        mask = torch.ones((S,S_full), device=x.device)
        mask = torch.tril(mask, diagonal=past_length)
        mask = mask[None, None, :, :]

        sq_mask = mask == 1
        
        # attention
        x = torch.nn.functional.scaled_dot_product_attention(
            q_heads, k_heads, v_heads,
            attn_mask=sq_mask
        )

        x = x.transpose(1, 2).reshape(B, S, D)

        # apply projection
        x = x @ self.wo.T

        return x, (k_heads, v_heads)
    
