import einops
import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class SinusoidalEmbedding(torch.nn.Module):
    def __init__(self, num_positions: int, d_model: int):
        super().__init__()
        assert d_model % 2 == 0, "Cannot use sin/cos positional encoding with odd dim (got dim={:d})".format(d_model)
        embedding = torch.zeros(num_positions, d_model)
        position = torch.arange(0, num_positions, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        embedding[:, 0::2] = torch.sin(position * div_term)
        embedding[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('embedding', embedding)

    def forward(self, indices: torch.Tensor):
        return self.embedding[indices]
    
class DropPath(nn.Module):
    def __init__(self, drop_prob: float = 0.):
        super().__init__()
        self.drop_prob = drop_prob

    def forward(self, x: torch.Tensor):
        if not self.training or self.drop_prob == 0.:
            return x
        keep_prob = 1 - self.drop_prob
        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
        random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
        return x * random_tensor.div(keep_prob)

class SwiGLU(nn.Module):
    def forward(self, x: torch.Tensor):
        x, gate = x.chunk(2, dim=-1)
        return x * F.silu(gate)
       
class MLP(nn.Module):
    def __init__(self, dim_in, dim_out: int = None, dim_hidden: int = None, bias: bool = False):
        super().__init__()
        dim_out = dim_out or dim_in
        dim_hidden = dim_hidden or dim_in
        self.hidden_proj = nn.Linear(dim_in, dim_hidden * 2, bias=bias)
        self.activation = SwiGLU()
        self.out_proj = nn.Linear(dim_hidden, dim_out, bias=bias)

    def forward(self, x: torch.Tensor):
        x = self.hidden_proj(x)
        x = self.activation(x)
        x = self.out_proj(x)
        return x

class RMSNorm(nn.Module):
    def __init__(self, dim: int, learnable: bool = True):
        super().__init__()
        self.dim = dim
        self.alpha = nn.Parameter(torch.ones(dim)) if learnable else torch.ones(dim)
        self.gamma = dim ** 0.5

    def forward(self, x: torch.Tensor):
        return F.normalize(x, p=2, dim=-1) * self.alpha * self.gamma

class Attention(nn.Module):
    def __init__(self, dim_q: int, dim_kv: int = None, dim_out: int = None, dim_head: int = 64, num_heads: int = None, **kwargs):
        super().__init__()
        dim_kv = dim_kv or dim_q
        dim_out = dim_out or dim_q
        self.nhead = num_heads or dim_q // dim_head
        dim_inner = dim_head * self.nhead
        self.to_q = nn.Linear(dim_q, dim_inner, bias = False)
        self.to_k = nn.Linear(dim_kv, dim_inner, bias = False)
        self.to_v = nn.Linear(dim_kv, dim_inner, bias = False)
        self.out_proj = nn.Linear(dim_inner, dim_out, bias = False)

    def forward(self, q, k, v, is_causal = False, **kwargs):
        # q: (B, *, N, D), k: (B, *, M, D), v: (B, *, M, D)
        B, ndim = q.size(0), q.ndim # needed to handle group queries

        # prepare queries, keys, values
        q, k, v = self.to_q(q), self.to_k(k), self.to_v(v) # (B, *, N, D) -> (B, *, N, H * d)
        q, k, v = map(lambda x: einops.rearrange(x, '... n (h d) -> (...) h n d', h = self.nhead), (q, k, v)) # (B, *, N, H * d) -> (B*, H, N, d)

        # scaled dot-product attention
        attn = F.scaled_dot_product_attention(q, k, v, is_causal = is_causal, **kwargs)
        out = einops.rearrange(attn, 'b h n d ->  b n (h d)') if ndim == 3 else einops.rearrange(attn, '(b g) h n d -> b g n (h d)', b = B)

        # project back to original dimension
        out = self.out_proj(out)
        return out
    
class TransformerBlock(nn.Module):
    def __init__(self, dim: int, dim_head: int = 64, num_heads=None, drop_path: float = 0., is_causal=False, **kwargs):
        super().__init__()
        self.is_causal = is_causal
        self.norm_attn = RMSNorm(dim)
        self.norm_ff = RMSNorm(dim)
        self.drop_path = DropPath(drop_path)
        self.attn = Attention(dim, dim_head = dim_head, num_heads = num_heads)
        self.ff = MLP(dim)

    def forward(self, x):
        # x: (B, *, N, D)
        # attention
        skip = x
        x = self.norm_attn(x)
        x = self.attn(x, x, x, is_causal = self.is_causal)
        x = self.drop_path(x) + skip
        
        # feed-forward
        skip = x
        x = self.norm_ff(x)
        x = self.ff(x)
        return self.drop_path(x) + skip

class PerceiverBlock(nn.Module):
    def __init__(self, dim: int, dim_tgt: int = None, dim_src: int = None, dim_head: int = 64, num_heads: int = None, **kwargs):
        super().__init__()
        dim_tgt = dim_tgt or dim
        dim_src = dim_src or dim
        self.norm_tgt = RMSNorm(dim_tgt)
        self.norm_src = RMSNorm(dim_src)
        self.norm_ff = RMSNorm(dim)
        self.attn = Attention(dim_q = dim_tgt, dim_kv = dim_src, dim_head = dim_head, dim_out = dim, num_heads = num_heads) 
        self.ff = MLP(dim)

    def forward(self, src, tgt):
        # src: (B, *, N, D), tgt: (B, *, M, D)
        # cross-attention
        tgt = self.norm_tgt(tgt)
        src = self.norm_src(src)
        tgt = self.attn(tgt, src, src) # (B, *, M, D) x (B, *, N, D) -> dim: int(B, *, M, D)

        # feed-forward
        skip = tgt
        tgt = self.norm_ff(tgt)
        tgt = self.ff(tgt)
        return tgt + skip


