import torch

from transformers import PreTrainedModel
from typing import OrderedDict
from .configuration_lenet import LeNetConfig


class LeNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = torch.nn.Sequential(
            OrderedDict(
                [
                    ("conv1", torch.nn.Conv2d(1, 6, 5, padding=2)),
                    ("pooling1", torch.nn.AvgPool2d(2, stride=2)),
                    ("sigmoid1", torch.nn.Sigmoid()),
                    ("conv2", torch.nn.Conv2d(6, 16, 5)),
                    ("pooling2", torch.nn.AvgPool2d(2, stride=2)),
                    ("sigmoid1", torch.nn.Sigmoid()),
                    ("flatten", torch.nn.Flatten()),
                    ("dense1", torch.nn.Linear(400, 120)),
                    ("sigmoid3", torch.nn.Sigmoid()),
                    ("dense2", torch.nn.Linear(120, 84)),
                    ("sigmoid4", torch.nn.Sigmoid()),
                    ("dense3", torch.nn.Linear(84, 10)),
                ]
            )
        )
        
    def forward(self, pixel_values) -> torch.Tensor:
        return self.model(pixel_values)
        

class LeNetModel(PreTrainedModel):
    config_class = LeNetConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = LeNet()

    def forward(self, pixel_values):
        return self.model.forward_features(pixel_values)
        

class LeNetModelForImageClassification(PreTrainedModel):
    config_class = LeNetConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = LeNet()

    def forward(self, pixel_values, labels=None):
        logits = self.model(pixel_values)
        if labels is not None:
            loss = torch.nn.functional.cross_entropy(logits, labels)
            return {"loss": loss, "logits": logits}
        return {"logits": logits}