# Codeblock 7a
class ResNeXt(nn.Module):
    def __init__(self):
        super().__init__()
        

        # conv1 stage  #(1)
        self.resnext_conv1 = nn.Conv2d(in_channels=NUM_CHANNELS[0],
                                       out_channels=NUM_CHANNELS[1],
                                       kernel_size=7,    #(2) 
                                       stride=2,         #(3)
                                       padding=3, 
                                       bias=False)
        nn.init.kaiming_normal_(self.resnext_conv1.weight, 
                                nonlinearity='relu')
        self.resnext_bn1 = nn.BatchNorm2d(num_features=NUM_CHANNELS[1])
        self.relu = nn.ReLU()
        self.resnext_maxpool1 = nn.MaxPool2d(kernel_size=3,    #(4)
                                             stride=2, 
                                             padding=1)
        

        # conv2 stage  #(5)
        self.resnext_conv2 = nn.ModuleList([
            Block(in_channels=NUM_CHANNELS[1],
                  add_channel=True,       #(6)
                  channel_multiplier=4,
                  downsample=False)       #(7)
        ])
        for _ in range(NUM_BLOCKS[0]-1):  #(8)
            self.resnext_conv2.append(Block(in_channels=NUM_CHANNELS[2]))
            

        # conv3 stage  #(9)
        self.resnext_conv3 = nn.ModuleList([Block(in_channels=NUM_CHANNELS[2],  #(10)
                                                  add_channel=True, 
                                                  downsample=True)])
        for _ in range(NUM_BLOCKS[1]-1):    #(11)
            self.resnext_conv3.append(Block(in_channels=NUM_CHANNELS[3]))
            
            
        # conv4 stage  #(12)
        self.resnext_conv4 = nn.ModuleList([Block(in_channels=NUM_CHANNELS[3],  #(13)
                                                  add_channel=True, 
                                                  downsample=True)])
        
        for _ in range(NUM_BLOCKS[2]-1):    #(14)
            self.resnext_conv4.append(Block(in_channels=NUM_CHANNELS[4]))
            
            
        # conv5 stage  #(15)
        self.resnext_conv5 = nn.ModuleList([Block(in_channels=NUM_CHANNELS[4],  #(16)
                                                  add_channel=True, 
                                                  downsample=True)])
        
        for _ in range(NUM_BLOCKS[3]-1):    #(17)
            self.resnext_conv5.append(Block(in_channels=NUM_CHANNELS[5]))
 
       
        self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1,1))  #(18)

        self.fc = nn.Linear(in_features=NUM_CHANNELS[5],        #(19)
                            out_features=NUM_CLASSES)