# Codeblock 8a
class DenseNet(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.first_conv = nn.Conv2d(in_channels=3, 
                                    out_channels=64, 
                                    kernel_size=7,    #(1)
                                    stride=2,         #(2)
                                    padding=3,        #(3)
                                    bias=False)
        self.first_pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)  #(4)
        channel_count = 64
        

        # Dense block #0
        self.dense_block_0 = DenseBlock(in_channels=channel_count,
                                        repeats=REPEATS[0])          #(5)
        channel_count = int(channel_count+REPEATS[0]*GROWTH)         #(6)
        self.transition_0 = Transition(in_channels=channel_count, 
                                       out_channels=int(channel_count*COMPRESSION))
        channel_count = int(channel_count*COMPRESSION)               #(7)
    

        # Dense block #1
        self.dense_block_1 = DenseBlock(in_channels=channel_count, 
                                        repeats=REPEATS[1])
        channel_count = int(channel_count+REPEATS[1]*GROWTH)
        self.transition_1 = Transition(in_channels=channel_count, 
                                       out_channels=int(channel_count*COMPRESSION))
        channel_count = int(channel_count*COMPRESSION)

        # # Dense block #2
        self.dense_block_2 = DenseBlock(in_channels=channel_count, 
                                        repeats=REPEATS[2])
        channel_count = int(channel_count+REPEATS[2]*GROWTH)
        
        self.transition_2 = Transition(in_channels=channel_count, 
                                       out_channels=int(channel_count*COMPRESSION))
        channel_count = int(channel_count*COMPRESSION)

        # Dense block #3
        self.dense_block_3 = DenseBlock(in_channels=channel_count, 
                                        repeats=REPEATS[3])
        channel_count = int(channel_count+REPEATS[3]*GROWTH)
        
        
        self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1,1))       #(8)
        self.fc = nn.Linear(in_features=channel_count, out_features=1000)  #(9)