import os
import torch
from torch import nn, optim
from torch.nn import functional as F
from torchvision import datasets, transforms
from torchvision.transforms.functional import to_pil_image

import PIL 
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from sklearn.metrics import accuracy_score



arquivo = 'lenet-modular2024-06-26.pt'

class Lenet5Modular(nn.Module):
  def __init__(self):
    super(Lenet5Modular,self).__init__()

    self.conv1 = nn.Sequential(
        nn.Conv2d(in_channels=3,out_channels=6,kernel_size=(5,5),stride=(1,1)),
        nn.Tanh(),
        nn.AvgPool2d(kernel_size=(2,2),stride=(2,2))
    )

    self.conv2 = nn.Sequential(
        nn.Conv2d(in_channels=6,out_channels=16,kernel_size=(5,5),stride=(1,1)),
        nn.Tanh(),
        nn.AvgPool2d(kernel_size=(2,2),stride=(2,2))
    )

    self.conv3 = nn.Sequential(
        nn.Conv2d(in_channels=16,out_channels=120,kernel_size=(5,5),stride=(1,1)),
        nn.Tanh()
    )

    self.flat = nn.Flatten()

    self.linear1 = nn.Sequential(
        nn.Linear(120,84),
        nn.Tanh()
    )

    self.linear2 = nn.Linear(84,10)

  def forward(self,x):
    x = self.conv1(x)
    x = self.conv2(x)
    x = self.conv3(x)
    x = self.flat(x)
    x = self.linear1(x)
    x = self.linear2(x)

    return x


modelo = Lenet5Modular()

# modelo.load_state_dict(torch.load(arquivo))

modelo.cpu()

a = torch.rand(1, 3, 32, 32)

b = modelo(a)

b.size()


# Função para criar o transformador
def create_transformer():
    return transforms.Compose([
        transforms.Resize((32, 32)),  # Redimensionar para 32x32
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(10),
        transforms.RandomAffine(0, shear=10, scale=(0.8,1.2)),
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
        transforms.ToTensor()
    ])