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

import torch
import torch.nn as nn
import pytorch_lightning as pl

from PIL import Image
from torch.utils.data import DataLoader

from torchvision import datasets, transforms
from torchmetrics.classification import Accuracy
# %% LeeNet

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.ReLU(),
            nn.AvgPool2d(kernel_size=2, stride=2),
            
            nn.Conv2d(6, 16, kernel_size=5),
            # nn.Tanh(),
            nn.ReLU(),
            nn.AvgPool2d(kernel_size=2, stride=2)          
        )
        
        
        self.mlp = nn.Sequential(
            nn.Linear(256, 120),
            # nn.LazyLinear(120),
            # nn.Tanh(),
            nn.ReLU(),
            nn.Linear(120, 84),
            # nn.Tanh(),
            nn.ReLU(),
            nn.Linear(84, out_classes)
        )
        
    def forward(self, image):
        x = self.feature(image)
        x = self.mlp(x.flatten(1))
        
        return x
    
class LeNetCustomLightning(pl.LightningModule):
    def __init__(self,
                 model: nn.Module,
                 lr_rate: float=0.1):
        super().__init__()
        
        self.save_hyperparameters(ignore=["model"])
        self.model = model
        
        self.criterion = nn.CrossEntropyLoss()
        self.train_metric = Accuracy(task="multiclass",
                                     num_classes=model._out_classes)
        self.val_metric = self.train_metric.clone()
        self.test_metric = self.train_metric.clone()

    def forward(self, image):
        return self.model(image)
    
    def _shared_step(self, sample, stage, metric):
        x, y = sample
        
        logits = self.forward(x.to(torch.float32))
        preds = torch.argmax(logits, 1)
        loss = self.criterion(logits, y.to(torch.int64))
        
        metric.update(preds, y)
        
        self.log(f"{stage}_loss", loss, prog_bar=True, 
                 on_epoch=True, on_step=False)
        
        return loss
    
    def _log_metric(self, stage, metric):
        computed = metric.compute()
        
        self.log(f"{stage}_acc", computed, prog_bar=True)
        
        metric.reset()
    
    def training_step(self, batch, batch_idx):
        return self._shared_step(batch, "train", self.train_metric)
    
    def on_train_epoch_end(self):
        self._log_metric("train", self.train_metric)
    
    def validation_step(self, batch, batch_idx):
        return self._shared_step(batch, "val", self.val_metric)
    
    def on_validation_epoch_end(self):
        self._log_metric("val", self.val_metric)
    
    def test_step(self, batch, batch_idx):
        return self._shared_step(batch, "test", self.test_metric)
    
    def on_test_epoch_end(self):
        self._log_metric("test", self.test_metric)
        
    def configure_optimizers(self):
        optimizer = torch.optim.SGD(self.parameters(), 
                                    lr=self.hparams.lr_rate)
        
        return optimizer

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

batch_size = 64

image_transforms = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

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

test_dataset = datasets.MNIST(
    root="./Datasets",
    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)

# %% Запуск обучения

num_epoch=10

lenet = LeNetCustom(in_channels=1,
                    out_classes=10)
lenet_lightning = LeNetCustomLightning(model=lenet,
                                       lr_rate=0.1)
lenet_lightning.compile()

trainer = pl.Trainer(max_epochs=num_epoch,
                     accelerator="gpu",
                     devices=1)

trainer.fit(model=lenet_lightning,
            train_dataloaders=train_loader,
            val_dataloaders=test_loader)



# %% Matlab Export

model = lenet_lightning.model
model.eval()
model = model.cpu()

dummy_input = torch.randn(1, 1, 28, 28, dtype=torch.float32)
with torch.no_grad():
    _ = model(dummy_input)
    
traced_model = torch.jit.trace(model, dummy_input)
traced_model.save("../models/LeNetCustom.pt")

# %% ONNX экспорт

model = lenet_lightning.model.cpu().eval()

dummy_input = torch.randn(1, 1, 28, 28, dtype=torch.float32)

torch.onnx.export(
    model,
    dummy_input,
    "../models/LeNetCustom.onnx",
    input_names=["input"],
    output_names=["logits"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "logits": {0: "batch_size"}
    },
    opset_version=17,
    # Проверка/оптимизация графа
    do_constant_folding=True
)

# %% Новый экспорт в ONNX

model = lenet_lightning.model.cpu().eval()
dummy_input = torch.randn(1, 1, 28, 28, dtype=torch.float32)

onnx_program = torch.onnx.export(
    model,
    (dummy_input,),
    dynamo=True,
    verify=True
)

onnx_program.save("../models/LeNetCustomNewExport.onnx")

# %% Инференс модели

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.jit.load("./models/LeNetCustom.pt", map_location=device)
model.eval()

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

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')}")