import glob
import os

import numpy as np
import torch.utils.data.dataset
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torchvision import transforms
from PIL import Image

import ex4

IMG_SIZE = 100


class ImageDataset(Dataset):
    def __init__(self, image_dir, precision: np.float32 or np.float64):
        self.image_files = sorted(glob.glob(os.path.join(image_dir, "**", "*.jpg"), recursive=True))
        self.precision = precision

    def __getitem__(self, index):
        # Open image file, convert to numpy array and scale to [0, 1]
        target_image = Image.open(self.image_files[index])

        target_image = preprocess(target_image, self.precision)

        # calculate image with black grid
        doomed_image = ex4.ex4(target_image, (5, 5), (4, 4))

        # convert image to grayscale
        # target_image = rgb2gray(target_image)  # todo look if gray image makes sense

        return doomed_image[0], np.transpose(target_image, (2, 0, 1))

    def __len__(self):
        return len(self.image_files)


def preprocess(input: np.array, precision: np.float32 or np.float64) -> np.array:
    # image = np.array(Image.open(self.image_files[index]), dtype=np.float32) / 255
    resize_transforms = transforms.Compose([
        transforms.Resize(size=IMG_SIZE),
        transforms.CenterCrop(size=(IMG_SIZE, IMG_SIZE)),
    ])
    input = resize_transforms(input)

    # normalize image from 0-1
    target_image = np.array(input, dtype=precision) / 255.0

    # Perform normalization for each channel
    # image = (image - self.norm_mean) / self.norm_std

    return target_image


# postprecess should be the inverese function of preprocess!
def postprocess(input: np.array) -> np.array:
    target_image = (input * 255.0).astype(np.uint8)
    return target_image


def get_image_loader(path: str, precision: np.float32 or np.float64):
    image_dataset = ImageDataset(path, precision)
    totlen = len(image_dataset)

    test_set_size = .1
    trains, tests = torch.utils.data.dataset.random_split(image_dataset, lengths=(totlen - int(totlen * test_set_size),
                                                                                  int(totlen * test_set_size)),
                                                          generator=torch.Generator().manual_seed(0))

    train_loader = DataLoader(
        trains,
        shuffle=True,  # shuffle the order of our samples
        batch_size=25,  # stack 4 samples to a minibatch
        num_workers=4  # no background workers (see comment below)
    )

    test_loader = DataLoader(
        tests,
        shuffle=True,  # shuffle the order of our samples
        batch_size=5,  # stack 4 samples to a minibatch
        num_workers=0  # no background workers (see comment below)
    )

    return train_loader, test_loader