"""Common transformer and attention blocks."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../src/models/transformers/attention.ipynb.

# %% auto #0
__all__ = ['FeedForwardBlock', 'BasisSelfAttnBlock', 'BasisCrossAttnBlock', 'SpatialTransformerSelfAttn', 'SpatialTransformer']

# %% ../../../src/models/transformers/attention.ipynb #06272f6f-b4e3-4504-a90a-feebbf6ad821
from ...imports import *

# %% ../../../src/models/transformers/attention.ipynb #cc6ea7ff-988c-4cab-96ba-fa32e3b6da2c
class FeedForwardBlock(nn.Module):
    """
    A small dense feed-forward network as used in `transformers`. Assumes channel last.
    Inspired by https://arxiv.org/pdf/2401.11605.
    From https://arxiv.org/pdf/2002.05202 a modification to SiGLU
    """

    def __init__(self, in_dim: int, hidden_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.hidden_dim = hidden_dim
        self.proj_in  = nn.Linear(in_dim, 2*hidden_dim) # factor two for GLU part split
        self.proj_out = nn.Linear(hidden_dim, in_dim) 
        self.act   = nn.SiLU()
        self.drop  = nn.Dropout(dropout)
        
    def siglu(self, x: torch.Tensor) -> torch.Tensor:
        x = self.proj_in(x) 
        return x[..., :self.hidden_dim] * self.act(x[..., self.hidden_dim:])

    #@torch.compile
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.siglu(x)
        x = self.drop(x)
        x = self.proj_out(x)
        return x

# %% ../../../src/models/transformers/attention.ipynb #a88bad1d-1ce4-4dab-92c0-6e3af703ae5a
class BasisSelfAttnBlock(nn.Module):
    """A self attention block, i.e. a `transformer` encoder."""
    def __init__(self, ch, num_heads, dropout=0.0, batch_first=False):
        super().__init__()
        self.self_att  = nn.MultiheadAttention(ch, num_heads=num_heads, batch_first=batch_first) #[t, b, c]
        self.ff    = FeedForwardBlock(ch, 2*ch)   
        self.norm1 = nn.LayerNorm(ch)
        self.norm2 = nn.LayerNorm(ch)
        self.drop  = nn.Dropout(dropout)
               
    def forward(self, x, attn_mask=None, key_padding_mask=None, need_weights=False):
        #x     ... [  t, batch, ch]       
        
        self_out    = self.norm1(x)  
        self_out, _ = self.self_att(self_out, key=self_out, value=self_out, attn_mask=attn_mask, key_padding_mask=key_padding_mask, need_weights=need_weights)
        self_out    = self.drop(self_out) + x      
        
        feed_out = self.norm2(self_out)              
        feed_out = self.ff(feed_out)
        feed_out = self.drop(feed_out) + self_out            
                   
        return feed_out

# %% ../../../src/models/transformers/attention.ipynb #e7291e33-7ca6-48a2-8556-f5a2cf3e476c
class BasisCrossAttnBlock(nn.Module):
    """A cross attention block, i.e. a `transformer` decoder."""
    def __init__(self, ch, num_heads, dropout=0.0, batch_first=False):
        super().__init__()
        self.self_att  = nn.MultiheadAttention(ch, num_heads=num_heads, batch_first=batch_first) #[t, b, c]
        self.cross_att = nn.MultiheadAttention(ch, num_heads=num_heads, batch_first=batch_first) 
        self.ff    = FeedForwardBlock(ch, 2*ch)   
        self.norm1 = nn.LayerNorm(ch)
        self.norm2 = nn.LayerNorm(ch)
        self.norm3 = nn.LayerNorm(ch)
        self.drop  = nn.Dropout(dropout)
        
    def forward(self, x, c_emb, attn_mask=None, key_padding_mask=None, need_weights=False):
        #x     ... [  t, batch, ch]       
        #c_emb ... [seq, batch, ch]
        
        self_out    = self.norm1(x)  
        self_out, _ = self.self_att(self_out, key=self_out, value=self_out, attn_mask=attn_mask, key_padding_mask=key_padding_mask, need_weights=need_weights)
        self_out    = self.drop(self_out) + x      
        
        cross_out    = self.norm2(self_out)   
        cross_out, _ = self.cross_att(cross_out, key=c_emb, value=c_emb, need_weights=need_weights)
        cross_out    = self.drop(cross_out) + self_out         
        
        feed_out = self.norm3(cross_out)              
        feed_out = self.ff(feed_out)
        feed_out = self.drop(feed_out) + cross_out            
                   
        return feed_out

# %% ../../../src/models/transformers/attention.ipynb #42ff2a79-66b6-4e1f-9a86-527894be54ad
class SpatialTransformerSelfAttn(nn.Module):
    """A spatial residual `transformer`, only uses self-attention."""
    def __init__(self, ch, num_heads, depth, dropout=0.0, num_groups=32):
        super().__init__()       
        self.norm               = torch.nn.GroupNorm(num_groups=num_groups, num_channels=ch, eps=1e-6, affine=True)
        self.transformer_blocks = nn.ModuleList([BasisSelfAttnBlock(ch=ch, num_heads=num_heads, dropout=dropout) for d in range(depth)])
        
    def forward(self, x, attn_mask=None, key_padding_mask=None):
        #x      ... [batch, ch, space, time]  
        #c_emb  ... [batch, seq, ch]
        b, ch, space, time = x.shape
            
        x_in = x
        
        #-------------------------
        x = self.norm(x) 
        
        x = torch.reshape(x, (b, ch, space*time))
        x = torch.permute(x, (2, 0, 1))#.contiguous()           # to [t, batch, ch]    
        
        #-------------------------           
        for transformer_block in self.transformer_blocks:
            x = transformer_block(x, attn_mask, key_padding_mask)
                
        #-------------------------
            
        x = torch.permute(x, (1, 2, 0))           # back to [batch, ch, t] 
        x = torch.reshape(x, (b, ch, space, time))#.contiguous()
                
        return x + x_in

# %% ../../../src/models/transformers/attention.ipynb #bcbc5dca-1f77-4276-85e2-b99afeee1128
class SpatialTransformer(nn.Module):
    """A spatial residual `transformer`, uses self- and cross-attention on conditional input."""
    
    def __init__(self, ch, cond_emb_size, num_heads, depth, dropout=0.0, num_groups=32):
        super().__init__()       
        self.cat_proj           = nn.Linear(cond_emb_size, ch)  
        self.norm               = torch.nn.GroupNorm(num_groups=num_groups, num_channels=ch, eps=1e-6, affine=True)
        self.transformer_blocks = nn.ModuleList([BasisCrossAttnBlock(ch=ch, num_heads=num_heads, dropout=dropout) for d in range(depth)])
        
    def forward(self, x, c_emb, attn_mask=None, key_padding_mask=None):
        #x      ... [batch, ch, space, time]  
        #c_emb  ... [batch, seq, ch]
        b, ch, space, time = x.shape
            
        x_in = x
        
        #-------------------------
        x = self.norm(x) 
        
        x = torch.reshape(x, (b, ch, space*time))
        x = torch.permute(x, (2, 0, 1))#.contiguous()           # to [t, batch, ch]    
       
        c_emb = self.cat_proj(c_emb)        
        c_emb = torch.permute(c_emb, (1, 0, 2))#.contiguous()  # to [seq, batch, ch]
        
        #-------------------------           
        for transformer_block in self.transformer_blocks:
            x = transformer_block(x, c_emb, attn_mask, key_padding_mask)
                
        #-------------------------
            
        x = torch.permute(x, (1, 2, 0))              # back to [batch, ch, t] 
        x = torch.reshape(x, (b, ch, space, time))#.contiguous()
                
        return x + x_in
