import torch
from utils import * 

dtype = torch.cuda.FloatTensor


class OML(nn.Module):                          
    def __init__(self,in_channel,out_channel):  
        super(OML, self).__init__()
        self.conv2 = nn.Sequential(nn.Conv2d(in_channel, out_channel,kernel_size=1,stride=1,padding=0,bias=False),          
                                nn.LeakyReLU())

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

class ST(nn.Module):                          
    def __init__(self,in_channel,out_channel):  
        super(ST, self).__init__()
        self.conv1 = nn.Sequential(nn.Conv2d(in_channel, out_channel,kernel_size=3,stride=1,padding=1,bias=False, padding_mode='circular'),                              
                                nn.LeakyReLU())
    def forward(self,x):
        return self.conv1(x)



# (conv1 conv3 conv1 conv3)*3


class CNT3D(nn.Module):                                              
    def __init__(self,n1,n4_1, n2,n4_2 ,n3,n4_3):                       
        super(CNT3D, self).__init__() 
                                        
        self.f1 = nn.Sequential(OML(n1,n4_1),ST(n4_1,n4_1))
        self.g1 = nn.Sequential(OML(n4_1,n1),ST(n1,n1))
        
        self.f2 = nn.Sequential(OML(n2,n4_2),ST(n4_2,n4_2))
        self.g2 = nn.Sequential(OML(n4_2,n2),ST(n2,n2))

        self.f3 = nn.Sequential(OML(n3,n4_3),ST(n4_3,n4_3))
        self.g3 = nn.Sequential(OML(n4_3,n3),ST(n3,n3))
            
    def forward(self, x):

        x1 = x                     # 第一个维度为通道数

        x2 = x.squeeze(0)
        x2 = x2.permute(1, 0, 2)   # 第二个维度为通道数
        x2 = x2.unsqueeze(0)

        x3 = x.squeeze(0)
        x3 = x3.permute(2, 1, 0)   # 第三个维度为通道数
        x3 = x3.unsqueeze(0)

        f1 = self.f1(x1)
        g1 = self.g1(f1)

        f2 = self.f2(x2)
        g2 = self.g2(f2)

        f3 = self.f3(x3)
        g3 = self.g3(f3)

        return f1,g1, f2,g2, f3,g3
    