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

import kagglehub
import splitfolders
import pytorch_lightning as pl

from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision.transforms import v2 as transforms

from pathlib import Path
from typing import Optional, Callable, Any

# %% Скачивание набора данных

path = kagglehub.dataset_download("jehanbhathena/weather-dataset",
                                  output_dir="./Datasets",
                                  force_download=False)

# %% Разделение набора данных

splitfolders.ratio("./Datasets/weather-dataset",
                   output="./Datasets/weather-dataset-split",
                   ratio=(.8, .1, .1),
                   group_prefix=None,
                   move=False)

# %% Пример модуля для загрузки данных

TransformType = Optional[Callable[..., Any]]

class MyImageDataset(pl.LightningDataModule):
    def __init__(self,
                 batch_size: int = 1,
                 train_transform: TransformType = None,
                 val_transform: TransformType = None,
                 img_root: str = ""
                 ):
        
        super().__init__()
        
        self.batch_size = batch_size
        self.train_transform = train_transform
        self.val_transform = val_transform
        self.img_root = img_root
        
    def setup(self, stage=None):
        self.train_set = ImageFolder(self.img_root / Path("train"), 
                                     transform=self.train_transform)
        self.val_set = ImageFolder(self.img_root / Path("val"), 
                                     transform=self.val_transform)
        self.test_set = ImageFolder(self.img_root / Path("test"), 
                                     transform=self.val_transform)
        
    def train_dataloader(self):
        return DataLoader(self.train_set, batch_size=self.batch_size,
                          shuffle=True, num_workers=2)
    
    def val_dataloader(self):
        return DataLoader(self.val_set, batch_size=self.batch_size,
                          shuffle=False, num_workers=2)
    
    def test_dataloader(self):
        return DataLoader(self.test_set, batch_size=self.batch_size,
                          shuffle=False, num_workers=2)

# %% Создание трансформаций

transform_train = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(p=.5),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])     
])

transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])     
])

# %% Создание экземляра класса и вывод объектов класса

dm = MyImageDataset(train_transform=transform_train,
                    val_transform=transform_test,
                    img_root="./Datasets/weather-dataset-split")
dm.setup()

print({k:v for k, v in enumerate(dm.test_set.classes)})