# Codeblock 3
class CAM(nn.Module):
    def __init__(self, num_channels, r=16):
        super().__init__()
        
        self.maxpool = nn.AdaptiveMaxPool2d(output_size=(1,1))  #(1)
        self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1,1))  #(2)
        
        self.mlp = nn.Sequential(
            nn.Linear(in_features=num_channels,
                      out_features=num_channels//r,   #(3)
                      bias=False),
            
            nn.ReLU(inplace=True),                    #(4)
            
            nn.Linear(in_features=num_channels//r,    #(5)
                      out_features=num_channels, 
                      bias=False)
        )
        
        self.sigmoid = nn.Sigmoid()                   #(6)
        
    def forward(self, x):         #(7)
        original = x
        print(f'original\t\t: {x.size()}\n')
        
        
        x_max = self.maxpool(x)   #(8)
        print(f'x after maxpool (x_max)\t: {x_max.size()}')
        
        x_avg = self.avgpool(x)   #(9)
        print(f'x after avgpool (x_avg)\t: {x_avg.size()}\n')
        
        
        x_max = torch.flatten(x_max, start_dim=1)    #(10)
        print(f'x_max after flatten\t: {x_max.size()}')
        
        x_avg = torch.flatten(x_avg, start_dim=1)    #(11)
        print(f'x_avg after flatten\t: {x_avg.size()}\n')
        
        
        x_max = self.mlp(x_max)   #(12)
        print(f'x_max after mlp\t\t: {x_max.size()}')
        
        x_avg = self.mlp(x_avg)   #(13)
        print(f'x_avg after mlp\t\t: {x_avg.size()}\n')
        
        
        x = x_max + x_avg         #(14)
        print(f'after sum\t\t: {x.size()}')
        
        x = self.sigmoid(x)       #(15)
        print(f'after sigmoid\t\t: {x.size()}')
        
        x = x[:, :, None, None]   #(16)
        print(f'after reshape\t\t: {x.size()}')
        
        x = x * original          #(17)           
        print(f'after multiply\t\t: {x.size()}')
        
        return x