import torch.nn as nn


class LeNetCustom(nn.Module):
    
    def __init__(self, 
                 in_channels: int=1,
                 out_classes: int=10):
        super().__init__()
        
        self._out_classes=out_classes
        
        self.feature = nn.Sequential(
            nn.Conv2d(in_channels, 6, kernel_size=5),
            nn.Tanh(),
            nn.AvgPool2d(kernel_size=2, stride=2),
            
            nn.Conv2d(6, 16, kernel_size=5),
            nn.Tanh(),
            nn.AvgPool2d(kernel_size=2, stride=2),

            nn.Conv2d(16, 120, kernel_size=5),
            nn.Tanh(),
        )
        
        
        self.mlp = nn.Sequential(
            nn.LazyLinear(84),
            nn.Tanh(),
            nn.Linear(84, out_classes)
        )
        
    def forward(self, image):
        x = self.feature(image)
        x = self.mlp(x.flatten(1))
        
        return x