# %% Импорт библиотек

import tqdm
import torch
import torch.nn as nn

from torch.utils.data import DataLoader
from torchmetrics import MetricCollection
from torchvision import datasets, transforms
from torchmetrics.classification import Accuracy, Precision, Recall, F1Score

from PIL import Image
from beartype import beartype

# %% Гиперпараметры
num_epoch = 8
batch_size = 64
learning_rate = 1e-3
device = "cuda" if torch.cuda.is_available() else "cpu"

# %% Трансформация данных

# Resize(size) — изменяет размер изображения.
# Grayscale() — преобразование изображения в оттенки серого.
# ToTensor() — преобразует изображение PIL или NumPy-массив в PyTorch-тензор.
# Normalize(mean, std) — нормализация тензора изображения.
# GaussianBlur(kernel_size) — размытие по Гауссу.


image_transforms = transforms.Compose([
    transforms.ToTensor()
])

# %% Загрузка набора данных

train_dataset = datasets.MNIST(
    root="./MNIST",
    train=True,
    download=False, # Если скачен, то убрать
    transform=image_transforms
)

test_dataset = datasets.MNIST(
    root="./MNIST",
    train=False,
    download=False, # Если скачен, то убрать
    transform=image_transforms    
)

train_loader = DataLoader(dataset=train_dataset,
                          batch_size=batch_size,
                          shuffle=True)

test_loader = DataLoader(dataset=test_dataset,
                         batch_size=batch_size,
                         shuffle=False)
# %% Структура модели

class FFN(nn.Module):
    
    @beartype
    def __init__(self,
                 input_size: int,
                 output_size: int):
        super().__init__()
        
        self.input_layer = nn.Linear(input_size, 32*32)
        self.hidden_layer_1 = nn.Linear(32*32, 32*32)
        self.hidden_layer_2 = nn.Linear(32*32, 32*32)
        self.output_layer = nn.Linear(32*32, output_size)
    
    @beartype
    def forward(self, image: torch.Tensor) -> torch.Tensor:
        x = image.flatten(start_dim=1)  # batch_size x 28*28
        x = torch.relu(self.input_layer(x))
        x = torch.relu(self.hidden_layer_1(x))
        x = torch.relu(self.hidden_layer_2(x))
        x = self.output_layer(x)
        
        return x # batch_size x 10

model = FFN(input_size=28*28,
            output_size=10)
model.to(device)

# %% Метрики

train_metrics = MetricCollection({
    "accuracy": Accuracy(task="multiclass", num_classes=10),
    "precision": Precision(task="multiclass", num_classes=10, average="macro"),    
    "recall": Recall(task="multiclass", num_classes=10, average="macro"),
    "f1": F1Score(task="multiclass", num_classes=10, average="macro")    
}).to(device)

test_metrics = MetricCollection({
    "accuracy": Accuracy(task="multiclass", num_classes=10),
    "precision": Precision(task="multiclass", num_classes=10, average="macro"),    
    "recall": Recall(task="multiclass", num_classes=10, average="macro"),
    "f1": F1Score(task="multiclass", num_classes=10, average="macro")    
}).to(device)

# %% Оптимизаторы

criterion = nn.CrossEntropyLoss()
optim = torch.optim.Adam(model.parameters(), lr=learning_rate)
# %% Тренировочный цикл

for epoch in tqdm.tqdm(range(num_epoch)):
    model.train()
    running_loss = 0
    
    for data, target in train_loader:
        data, target = data.to(device), target.to(device)
        
        optim.zero_grad()
        outputs = model(data)
        loss = criterion(outputs, target)
        loss.backward()
        running_loss += loss.item()
        optim.step()
        
        preds = outputs.argmax(dim=1)
        train_metrics.update(preds, target)
    
    results = train_metrics.compute()        
    print(f"Epoch: {epoch+1}")
    print("=== TRAIN ===")
    print(f"  Loss: {(running_loss/len(train_loader)):.4f}")
    print(f"  Accuracy:  {results['accuracy']:.4f}")
    print(f"  Precision: {results['precision']:.4f}")
    print(f"  Recall:    {results['recall']:.4f}")
    print(f"  F1:        {results['f1']:.4f}")
    print()
    
    train_metrics.reset()
    
    model.eval()
    with torch.no_grad():
        test_loss = 0
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            
            outputs = model(data)
            loss = criterion(outputs, target)
            test_loss += loss.item()
            preds = outputs.argmax(dim=1)
            test_metrics.update(preds, target)
        
        results = test_metrics.compute()        
        
        print("=== TEST ===")
        print(f"  Loss: {(test_loss/len(train_loader)):.4f}")
        print(f"  Accuracy:  {results['accuracy']:.4f}")
        print(f"  Precision: {results['precision']:.4f}")
        print(f"  Recall:    {results['recall']:.4f}")
        print(f"  F1:        {results['f1']:.4f}")
        print()
        
    test_metrics.reset()
        
    
torch.save(model.state_dict(), "./models/mnist_model.pth")
# %% Инференс модели

image_transforms = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize((28,28)),
    transforms.ToTensor()    
])

model.eval()
model.load_state_dict(torch.load("./models/mnist_model.pth"))

def predict(image_path):
    image = Image.open(image_path)
    image = image_transforms(image)
    image = image.unsqueeze(0)
    image = image.to(device)
    
    with torch.no_grad():
        output = model(image)
        
        pred = output.argmax(dim=1).item()
        
    return pred

print(f"Predicted label: {predict('./test_images/mnist/3.png')}")
print(f"Predicted label: {predict('./test_images/mnist/5.png')}")
print(f"Predicted label: {predict('./test_images/mnist/6.png')}")