# Codeblock 6
class PANet(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.downsample_n2 = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=2, padding=1)
        self.downsample_n3 = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=2, padding=1)
        self.downsample_n4 = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=2, padding=1)
        
        self.conv_n2down_p3 = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1)
        self.conv_n3down_p4 = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1)
        self.conv_n4down_p5 = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1)
        
        self.relu = nn.ReLU()

    def forward(self, p2, p3, p4, p5):
        
        n2 = p2                                         #(1)
        print(f'n2\t\t: {n2.size()}\n')
        
        #######################################
        
        n2down = self.relu(self.downsample_n2(n2))      #(2)
        print(f'n2 downsampled\t: {n2down.size()}')

        n2down_p3 = n2down + p3                         #(3)
        print(f'after sum\t: {n2down_p3.size()}')
        
        n3 = self.relu(self.conv_n2down_p3(n2down_p3))  #(4)
        print(f'n3\t\t: {n3.size()}\n')
        
        #######################################
        
        n3down = self.relu(self.downsample_n3(n3))
        print(f'n3 downsampled\t: {n3down.size()}')
        
        n3down_p4 = n3down + p4
        print(f'after sum\t: {n3down_p4.size()}')
        
        n4 = self.relu(self.conv_n3down_p4(n3down_p4))
        print(f'n4\t\t: {n4.size()}\n')
        
        #######################################
        
        n4down = self.relu(self.downsample_n4(n4))
        print(f'n4 downsampled\t: {n4down.size()}')
        
        n4down_p5 = n4down + p5
        print(f'after sum\t: {n4down_p5.size()}')
        
        n5 = self.relu(self.conv_n4down_p5(n4down_p5))
        print(f'n5\t\t: {n5.size()}')
        
        return n2, n3, n4, n5