from dataclasses import dataclass
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import torch.backends.cudnn as cudnn
import numpy as np
from torchvision.utils import make_grid
from torchvision.transforms import v2 
from torchvision import datasets, models
import matplotlib.pyplot as plt
import logging
import os
from tqdm import tqdm
from torch.utils.data import DataLoader
from PIL import Image
from tempfile import TemporaryDirectory
from torch.profiler import profile, record_function, ProfilerActivity
from pathlib import Path
from taolib.utils.logger import config_logging

temp_dir = Path(".").resolve() / ".temp"
temp_dir.mkdir(exist_ok=True)
config_logging(f'{temp_dir}/0-compile.log', "root", maxBytes=5000000, backupCount=7)
torch.cuda.empty_cache() # 清空 GPU 缓存

cudnn.benchmark = True
plt.ion()   # interactive mode

@dataclass
class CIFAR10:
    root_dir: str

    def __post_init__(self):
        self.root_dir = Path(self.root_dir)
        self.mean = [0.485, 0.456, 0.406]
        self.std = [0.229, 0.224, 0.225]
        self.train_transform = v2.Compose([ # 在高度和宽度上将图像放大到40像素的正方形
            v2.Resize(40),
            # 随机裁剪出一个高度和宽度均为40像素的正方形图像，
            # 生成一个面积为原始图像面积0.64到1倍的小正方形，
            # 然后将其缩放为高度和宽度均为32像素的正方形
            v2.RandomResizedCrop(32, scale=(0.64, 1.0), ratio=(1.0, 1.0)),
            v2.RandomHorizontalFlip(),
            v2.ToImage(),
        ])
        self.val_transform = v2.ToImage()
        self.train = datasets.CIFAR10(
            root=self.root_dir, 
            train=True, download=True, 
            transform=self.train_transform,
        )
        self.val = datasets.CIFAR10(
            root=self.root_dir, 
            train=False, download=True,
            transform=self.val_transform,
        )
        self.normalize = nn.Sequential(
            v2.ToDtype(torch.float32, scale=True),
            v2.Normalize(self.mean, self.std)
        )

dataset = CIFAR10(temp_dir/"data")
batch_size = 16
dataloaders = {
    "train": DataLoader(dataset.train, batch_size=batch_size, shuffle=True),
    "val": DataLoader(dataset.val, batch_size=batch_size, shuffle=False),
}
dataset_sizes = {
    "train": len(dataset.train),
    "val": len(dataset.val),
}
class_names = dataset.train.classes
# We want to be able to train our model on an `accelerator <https://pytorch.org/docs/stable/torch.html#accelerators>`__
# such as CUDA, MPS, MTIA, or XPU. If the current accelerator is available, we will use it. Otherwise, we use the CPU.
device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
print(f"Using {device} device")