import torch
from torch import nn

class MappingUnit(nn.Module):
    def __init__(self,dim, hidden_dim):
        super().__init__()
                   
        self.norm_token = nn.LayerNorm(dim)
        self.proj_1 =  nn.Linear(dim, hidden_dim, bias = False)
        self.proj_2 =  nn.Linear(dim, hidden_dim, bias = False)
        self.proj_3 = nn.Linear(hidden_dim, dim, bias = False)  
        self.gelu = nn.GELU()
             	   
    def forward(self, x):
    
    	x = self.norm_token(x)    	
    	u, v = x, x 
    	u = self.proj_1(u)
    	u = self.gelu(u)
    	v = self.proj_2(v)
    	g = u * v
    	x = self.proj_3(g)
    	    	
    	return x

class InteractionUnit(nn.Module):
    def __init__(self, dim):
        super().__init__()
            
        self.norm_token = nn.LayerNorm(dim)       
        self.gelu = nn.GELU()
             	   
    def forward(self, x):
    
    	x = self.norm_token(x)
    	dim0 = x.shape[0]
    	dim1 = x.shape[1]
    	dim2 = x.shape[2]
    	x = x.reshape([dim0, dim1 * dim2])
    	x = self.gelu(x)
    	x = x.reshape([dim0, dim1, dim2])
    	    	
    	return x

class InteractorBlock(nn.Module):
    def __init__(self, d_model, d_ffn):
        super().__init__()
                
        self.mapping = MappingUnit(d_model, d_ffn)
        self.interaction = InteractionUnit(d_model)
        
    def forward(self, x):
                  
        residual = x
        
        x = self.interaction(x)
    
        x = x + residual
        
        residual = x
        
        x = self.mapping(x)
                                                  
        out = x + residual
                
        return out

class Interactor(nn.Module):
    def __init__(self, d_model, d_ffn, num_layers):
        super().__init__()
        
        self.model = nn.Sequential(
            *[InteractorBlock(d_model, d_ffn) for _ in range(num_layers)]
        )

    def forward(self, x):
       
        return self.model(x)