import os
import math
import numpy as np
from glob import glob

from random import shuffle
from PIL import Image, ImageFilter

import torch
import torchvision.transforms.functional as F
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader



class InpaintingData(Dataset):
    def __init__(self, args):
        super(Dataset, self).__init__()     # 繼承Dataset的父類的初始化函數
        self.w = self.h = args.image_size   # 通過args傳入新的屬性---圖像的w和h
        self.mask_type = args.mask_type     # 通過args傳入新的屬性---mask_type

        # image and mask 
        self.image_path = []    #創建image_path的數組
        for ext in ['*.jpg', '*.png']:  # 獲取每一個后綴為.jpg或者.png的圖片，為ext
            # 將dir_image、data_train和ext拼接作為圖片的路徑，并將其存入到數組image_path之中,glob()獲取一個lsit集合
            self.image_path.extend(glob(os.path.join(args.dir_image, args.data_train, ext)))
        self.mask_path = glob(os.path.join(args.dir_mask, args.mask_type, '*.png')) #拼接dir_mask、mask_type和路徑下所有的.png作為mask_path

        # augmentation 
        self.img_trans = transforms.Compose([   #接收一個 transforms方法的list為參數，將這些操作組合到一起，返回一個新的tranforms
            transforms.RandomResizedCrop(args.image_size),  #隨機隨機長寬比裁剪，大小為image_size
            transforms.RandomHorizontalFlip(), #隨機水平翻轉
            transforms.ColorJitter(0.05, 0.05, 0.05, 0.05), #改變圖像的亮度、對比度、飽和度和色調。
            transforms.ToTensor()])     # 轉為tensor，并歸一化至[0-1]
        self.mask_trans = transforms.Compose([
            transforms.Resize(args.image_size, interpolation=transforms.InterpolationMode.NEAREST), #將輸入圖像調整為給定的大小，interpolation是插值方式，此處是默認值NEAREST
            transforms.RandomHorizontalFlip(),  #隨機水平翻轉
            transforms.RandomRotation(  #隨機旋轉
                (0, 45), interpolation=transforms.InterpolationMode.NEAREST), #(0, 45)是角度
        ])

    def __len__(self):  # __len__和__getitem__DataSet類必須實現的靜態方法
        return len(self.image_path)

    def __getitem__(self, index):
        # load image
        image = Image.open(self.image_path[index]).convert('RGB') #獲取圖像，并將其轉化為RGB（3x8位像素）模式
        filename = os.path.basename(self.image_path[index]) #獲取圖片的路徑

        if self.mask_type == 'pconv': #如果mask_type為pconv
            index = np.random.randint(0, len(self.mask_path)) #隨機從mask_path中獲取一個下標
            mask = Image.open(self.mask_path[index])    #根據下標獲取mask圖片
            mask = mask.convert('L')    #將mask圖片轉化為L(8位像素的黑白圖片，0表示黑，255表示白)模式
        else:   # 構造mask，有mask數據集的話就運行不到這里
            mask = np.zeros((self.h, self.w)).astype(np.uint8) #構造與h和w一樣大的圖片，都用0填充，并將其轉換為uint8
            mask[self.h // 4:self.h // 4 * 3, self.w // 4:self.w // 4 * 3] = 1
            mask = Image.fromarray(m).convert('L')

        # augment
        image = self.img_trans(image) * 2. - 1. # 數據標準化，將輸出限定在一定的范圍
        mask = F.to_tensor(self.mask_trans(mask))   # 將轉化后的mask圖像轉化為tensor

        return image, mask, filename    #返回


if __name__ == '__main__':
    from attrdict import AttrDict

    args = {
        'dir_image': '../../examples/logos',
        'data_train': 'image',
        'dir_mask': '../../examples/logos/mask',
        'mask_type': 'pconv',
        'image_size': 512
    }
    args = AttrDict(args) # 將上面定義的參數傳入AttrDict()作為新參數

    data = InpaintingData(args)     #創建InpaintingData對象
    print(len(data), len(data.mask_path))   #輸出data的長度，mask的長度
    img, mask, filename = data[0]   # 獲取第一張圖片
    print(img.size(), mask.size(), filename)    #打印上述信息