import torch.nn as nn # neural network module
import torch.nn.functional as F



# On va utiliser le dataset MNIST pour tester notre modèle, 
# il contient des images de chiffres manuscrits(0-9) de 28x28 pixels

# MNIST  est déjà disponible dans la librairie torchvision.datasets,
# on peut l'importer et l'utiliser facilement.
# cette fois ci on aura pas besoin de creer une classe Dataset personnalisée,
# Le Dataset est déjà dans PyTorch, on peut l'utiliser directement.

# MNIST composé de 60000 images pour l'entrainement et 10000 images pour le test,
# si on avait des images brutes, on pouvait partir de 
# la classe Dataset pour créer notre propre classe Dataset personnalisée,


# L'objective c'est de faire la classification de ces images, 
# c'est à dire prédire le chiffre représenté sur l'image

class Lenet5(nn.Module):
    def __init__(self, n_classes=10):
        super().__init__()
        self.conv1  = nn.Conv2d(1,   6,   kernel_size=5)
        self.pool1  = nn.AvgPool2d(2, 2)
        self.conv2  = nn.Conv2d(6,   16,  kernel_size=5)
        self.pool2  = nn.AvgPool2d(2, 2)
        self.conv3  = nn.Conv2d(16,  120, kernel_size=5)
        self.linear1 = nn.Linear(120, 84)
        self.linear2 = nn.Linear(84,  n_classes)

    def forward(self, image):
        x = F.tanh(self.conv1(image))
        x = F.tanh(self.pool1(x))
        x = F.tanh(self.conv2(x))
        x = F.tanh(self.pool2(x))
        x = F.tanh(self.conv3(x))
        x = x.view(x.size(0), -1)
        x = F.tanh(self.linear1(x))
        x = self.linear2(x)          # ← logits bruts, CrossEntropyLoss gère le softmax
        return x