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

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

# %% ../../src/models/transformers.ipynb 2
from ..imports import *
import genQC.models.layers as layers

# %% ../../src/models/transformers.ipynb 4
class BasisSelfAttnBlock(nn.Module):
    """A self attention block, i.e. a `transformer` encoder."""
    def __init__(self, ch, num_heads, dropout=0):
        super().__init__()
        self.self_att  = nn.MultiheadAttention(ch, num_heads=num_heads, batch_first=False) #[t, b, c]
        self.ff    = layers.FeedForward(ch, 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]       
        #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      
        
        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.ipynb 5
class BasisCrossAttnBlock(nn.Module):
    """A cross attention block, i.e. a `transformer` decoder."""
    def __init__(self, ch, cond_emb_size, num_heads, dropout=0.0):
        super().__init__()
        self.self_att  = nn.MultiheadAttention(ch, num_heads=num_heads, batch_first=False) #[t, b, c]
        self.cross_att = nn.MultiheadAttention(ch, num_heads=num_heads, batch_first=False) 
        self.ff    = layers.FeedForward(ch, 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.ipynb 7
class SpatialTransformerSelfAttn(nn.Module):
    """A spatial residual `transformer`, only uses self-attention."""
    def __init__(self, ch, num_heads, depth, dropout=0.0):
        super().__init__()       
        self.norm               = torch.nn.GroupNorm(num_groups=32, num_channels=ch, eps=1e-6, affine=True)
        self.transformer_blocks = nn.ModuleList([BasisSelfAttnBlock(ch, num_heads, 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]    
        
        #-------------------------   
        # x = self.proj_in(x) #NEW   only used so that ch is a multiple of heads
        
        for transformer_block in self.transformer_blocks:
            x = transformer_block(x, attn_mask, key_padding_mask)
                
        # feed_out = self.proj_out(feed_out) #NEW
        #-------------------------
            
        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.ipynb 8
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):
        super().__init__()       
        self.cat_proj           = nn.Linear(cond_emb_size, ch)  
        self.norm               = torch.nn.GroupNorm(num_groups=32, num_channels=ch, eps=1e-6, affine=True)
        self.transformer_blocks = nn.ModuleList([BasisCrossAttnBlock(ch, cond_emb_size, num_heads, 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]
        
        #-------------------------   
        # x = self.proj_in(x) #NEW   only used so that ch is a multiple of heads
        
        for transformer_block in self.transformer_blocks:
            x = transformer_block(x, c_emb, attn_mask, key_padding_mask)
                
        # feed_out = self.proj_out(feed_out) #NEW
        #-------------------------
            
        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
