import os
import random
import sys

import numpy as np
from PIL import Image
import torch
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from pathlib import Path
from common.color import PIL_CONVERT_COLOR
from common.utils import pil2numpy

image_extensions = ['.jpg', '.png']
def load_images_cached(images_dir_path, color_model, reset_cache):
    if not Path(images_dir_path).exists():
        raise Exception(f"{images_dir_path} does not exist.")
    image_paths = sorted([f for f in Path(images_dir_path).glob("*") if f.suffix.lower() in image_extensions])
    cache_path = Path(images_dir_path).parent / f"{Path(images_dir_path).stem}_{color_model}_cache.npy"
    cache_path = cache_path.resolve()
    if not Path(cache_path).exists() or reset_cache:
        print("Caching to:", cache_path)
        value = {f:pil2numpy(PIL_CONVERT_COLOR[color_model](Image.open(f))) for f in image_paths}
        np.save(cache_path, value, allow_pickle=True)
    else:
        value = np.load(cache_path, allow_pickle=True).item()
        print("Loaded cache from:", cache_path)
    return list(value.keys()), list(value.values())

class SRTrainDataset(Dataset):
    def __init__(self, hr_dir_path, lr_dir_path, patch_size, color_model = "RGB", rigid_aug=True, reset_cache = False):
        super(SRTrainDataset, self).__init__()
        self.sz = patch_size
        self.rigid_aug = rigid_aug
        self.hr_image_names, self.hr_images = load_images_cached(hr_dir_path, color_model=color_model, reset_cache=reset_cache) 
        self.lr_image_names, self.lr_images = load_images_cached(lr_dir_path, color_model=color_model, reset_cache=reset_cache) 
        assert len(self.hr_images) == len(self.lr_images)
        
    def __getitem__(self, idx):
        if isinstance(idx, slice):
            batch_hr, batch_lr = [], []
            for i in range(idx.start, idx.stop):
                hr_patch, lr_patch = self.__getitem__(i)
                batch_hr.append(hr_patch)
                batch_lr.append(lr_patch)
            return batch_hr, batch_lr
        else:
            idx = idx % len(self.hr_images)
            hr_image = self.hr_images[idx]
            lr_image = self.lr_images[idx]
            scale = hr_image.shape[0]//lr_image.shape[0]

            i = random.randint(0, lr_image.shape[0] - self.sz)
            j = random.randint(0, lr_image.shape[1] - self.sz)

            if len(hr_image.shape) == 3:
                hr_patch = hr_image[
                    (i*scale):(i*scale + self.sz*scale),
                    (j*scale):(j*scale + self.sz*scale), 
                    :
                ]
                lr_patch = lr_image[
                    i:(i + self.sz), 
                    j:(j + self.sz), 
                    :
                ]

            if self.rigid_aug:
                if random.uniform(0, 1) < 0.5:
                    hr_patch = np.fliplr(hr_patch)
                    lr_patch = np.fliplr(lr_patch)

                if random.uniform(0, 1) < 0.5:
                    hr_patch = np.flipud(hr_patch)
                    lr_patch = np.flipud(lr_patch)

                k = random.choice([0, 1, 2, 3])
                hr_patch = np.rot90(hr_patch, k)
                lr_patch = np.rot90(lr_patch, k)
            
            hr_patch = torch.tensor(hr_patch.copy()).type(torch.float32)
            lr_patch = torch.tensor(lr_patch.copy()).type(torch.float32)
            hr_patch = hr_patch.permute(2,0,1)
            lr_patch = lr_patch.permute(2,0,1)

            return hr_patch, lr_patch

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

class SRTestDataset(Dataset):
    def __init__(self, hr_dir_path, lr_dir_path, color_model, reset_cache):
        super(SRTestDataset, self).__init__()
        self.hr_image_paths, self.hr_images = load_images_cached(hr_dir_path, color_model=color_model, reset_cache=reset_cache) 
        self.lr_image_paths, self.lr_images = load_images_cached(lr_dir_path, color_model=color_model, reset_cache=reset_cache) 
        assert len(self.hr_images) == len(self.lr_images)
        
    def __getitem__(self, idx):
        if isinstance(idx, slice):
            batch_hr, batch_lr = [], []
            for i in range(idx.start, idx.stop):
                hr_image, lr_image = self.__getitem__(i)
                batch_hr.append(hr_image)
                batch_lr.append(lr_image)
            return batch_hr, batch_lr
        else:
            idx = idx % len(self.hr_images)
            return self.hr_images[idx], self.lr_images[idx], self.hr_image_paths[idx], self.lr_image_paths[idx]

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

    def __iter__(self):
        for i in range(0, len(self.hr_images)):
            yield self.__getitem__(i)
