# Codeblock 10
class MobileNetV2(nn.Module):
    def __init__(self):
        super().__init__()
        
        # Input shape: 3x224x224
        self.first_conv = Conv(first=True)
        
        # Input shape: 32x112x112
        self.inv_residual0 = InvResidualS1(in_channels=32, 
                                           out_channels=16, 
                                           t=1)
        
        # Input shape: 16x112x112
        self.inv_residual1 = nn.ModuleList([InvResidualS2(in_channels=16, 
                                                          out_channels=24, 
                                                          t=6)])
        
        self.inv_residual1.append(InvResidualS1(in_channels=24, 
                                                out_channels=24, 
                                                t=6))
        
        # Input shape: 24x56x56
        self.inv_residual2 = nn.ModuleList([InvResidualS2(in_channels=24, 
                                                          out_channels=32, 
                                                          t=6)])
        
        for _ in range(2):
            self.inv_residual2.append(InvResidualS1(in_channels=32, 
                                                    out_channels=32, 
                                                    t=6))
        
        # Input shape: 32x28x28
        self.inv_residual3 = nn.ModuleList([InvResidualS2(in_channels=32, 
                                                          out_channels=64, 
                                                          t=6)])
        
        for _ in range(3):
            self.inv_residual3.append(InvResidualS1(in_channels=64, 
                                                    out_channels=64, 
                                                    t=6))
            
        # Input shape: 64x14x14
        self.inv_residual4 = nn.ModuleList([InvResidualS1(in_channels=64, 
                                                          out_channels=96, 
                                                          t=6)])
        
        for _ in range(2):
            self.inv_residual4.append(InvResidualS1(in_channels=96, 
                                                    out_channels=96, 
                                                    t=6))
        
        
        # Input shape: 96x14x14
        self.inv_residual5 = nn.ModuleList([InvResidualS2(in_channels=96, 
                                                          out_channels=160, 
                                                          t=6)])
        
        for _ in range(2):
            self.inv_residual5.append(InvResidualS1(in_channels=160, 
                                                    out_channels=160, 
                                                    t=6))
        
        # Input shape: 160x7x7
        self.inv_residual6 = InvResidualS1(in_channels=160, 
                                           out_channels=320, 
                                           t=6)
        
        # Input shape: 320x7x7
        self.last_conv = Conv(first=False)
        
        self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1,1))        #(1)
        self.dropout = nn.Dropout(p=0.2)                              #(2)
        self.fc = nn.Linear(in_features=int(1280*WIDTH_MULTIPLIER),   #(3)
                            out_features=1000)

    def forward(self, x):
        x = self.first_conv(x)
        print(f"after first_conv\t: {x.shape}")
        
        x = self.inv_residual0(x)
        print(f"after inv_residual0\t: {x.shape}")
            
        for i, layer in enumerate(self.inv_residual1):
            x = layer(x)
            print(f"after inv_residual1 #{i}\t: {x.shape}")
            
        for i, layer in enumerate(self.inv_residual2):
            x = layer(x)
            print(f"after inv_residual2 #{i}\t: {x.shape}")
            
        for i, layer in enumerate(self.inv_residual3):
            x = layer(x)
            print(f"after inv_residual3 #{i}\t: {x.shape}")
            
        for i, layer in enumerate(self.inv_residual4):
            x = layer(x)
            print(f"after inv_residual4 #{i}\t: {x.shape}")
            
        for i, layer in enumerate(self.inv_residual5):
            x = layer(x)
            print(f"after inv_residual5 #{i}\t: {x.shape}")
        
        x = self.inv_residual6(x)
        print(f"after inv_residual6\t: {x.shape}")
        
        x = self.last_conv(x)
        print(f"after last_conv\t\t: {x.shape}")
        
        x = self.avgpool(x)
        print(f"after avgpool\t\t: {x.shape}")
        
        x = torch.flatten(x, start_dim=1)
        print(f"after flatten\t\t: {x.shape}")
        
        x = self.dropout(x)
        print(f"after dropout\t\t: {x.shape}")
        
        x = self.fc(x)
        print(f"after fc\t\t: {x.shape}")
                
        return x