import glob

import h5py
from PIL import Image
from torchvision.transforms import RandomCrop
from torchvision.transforms.functional import to_tensor
from tqdm import tqdm

from Dataloader import ImageAugment

patch_size = 128
shrink_size = 2
noise_level = 1
patches_per_img = 20
images = glob.glob("dataset/train/*")

database = h5py.File("train_images.hdf5", "w")

dat_group = database.create_group("shrink_2_noise_level_1_downsample_random_rgb")
# del database['shrink_2_noise_level_1_downsample_random']
storage_lr = dat_group.create_dataset(
    "train_lr",
    shape=(
        patches_per_img * len(images),
        3,
        patch_size // shrink_size,
        patch_size // shrink_size,
    ),
    dtype="float32",
    # compression='lzf',
)
storage_hr = dat_group.create_dataset(
    "train_hr",
    shape=(patches_per_img * len(images), 3, patch_size, patch_size),
    # compression='lzf',
    dtype="float32",
)

random_cropper = RandomCrop(size=patch_size)
img_augmenter = ImageAugment(shrink_size, noise_level, down_sample_method=None)


def get_img_patches(img_pil):
    img_patch = random_cropper(img_pil)
    lr_hr_patches = img_augmenter.process(img_patch)
    return lr_hr_patches


counter = 0
for img in tqdm(images):
    img_pil = Image.open(img).convert("RGB")
    for i in range(patches_per_img):
        patch = get_img_patches(img_pil)
        storage_lr[counter] = to_tensor(patch[0].convert("RGB")).numpy()
        storage_hr[counter] = to_tensor(patch[1].convert("RGB")).numpy()
        counter += 1
database.close()