import os
import cv2
import glob
import random
from PIL import Image
import numpy as np
import torch
from torchvision.io import read_video, write_jpeg
from torch.utils.data import Dataset
from torchvision import transforms as T
import imgaug.augmenters as iaa
import albumentations as A
from perlin import rand_perlin_2d_np
from nsa import patch_ex
from skimage.segmentation import mark_boundaries


class CustomDataset(Dataset):
    def __init__(self, c, is_train=True, excluded_images=None):
        self.dataset_path = c.data_path
        self.is_train = is_train
        self.cropsize = c.crop_size
        self.train_list = c.train_list
        self.val_list = c.val_list
        # load dataset
        if excluded_images is not None:
            self.x, self.y, self.mask, self.img_types = self.load_dataset_list()
            self.x, self.y, self.mask, self.img_types = excluding_images(self.x, self.y, self.mask, self.img_types, excluded_images)
        else:
            self.x, self.y, self.mask, self.img_types = self.load_dataset_list()
        # set transforms
        if is_train:
            self.transform_x = T.Compose([
                T.Resize(c.img_size, Image.ANTIALIAS),
                #T.RandomRotation(5),
                # T.CenterCrop(c.crop_size),
                T.ToTensor()])
        # test:
        else:
            self.transform_x = T.Compose([
                T.Resize(c.img_size, Image.ANTIALIAS),
                # T.CenterCrop(c.crop_size),
                T.ToTensor()])
        # mask
        self.transform_mask = T.Compose([
            T.Resize(c.img_size, Image.NEAREST),
            # T.CenterCrop(c.crop_size),
            T.ToTensor()])

        self.normalize = T.Compose([T.Normalize(c.norm_mean, c.norm_std)])

    def __getitem__(self, idx):
        img_path, y, mask, img_type = self.x[idx], self.y[idx], self.mask[idx], self.img_types[idx]
        
        x = Image.open(img_path)
        x = self.normalize(self.transform_x(x))
        
        if y == 0:
            mask = torch.zeros([1, self.cropsize[0], self.cropsize[1]])
        else:
            mask = Image.open(mask)
            mask = self.transform_mask(mask)
        return x, y, mask, os.path.basename(img_path[:-4]), img_type

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

    def load_dataset_list(self):
        x, y, mask, types = [], [], [], []
        x_n,y_n,mask_n,types_n = [],[],[],[]
        x_ab,y_ab,mask_ab,types_ab = [],[],[],[]
        
        # [img_path, normal(0) or NG(1), mask_path, type_name]
        if self.is_train == True:
            for t in self.train_list:
                if t[1] == 0:
                    x_n.extend([t[0]])
                    y_n.extend([t[1]])
                    mask_n.extend([t[2]])
                    types_n.extend([t[3]])
            x = x_n
            y = y_n
            mask = mask_n
            types = types_n
        elif self.is_train== False:
            for v in self.val_list:
                if v[1] == 0:
                    x_n.extend([v[0]])
                    y_n.extend([v[1]])
                    mask_n.extend([v[2]])
                    types_n.extend([v[3]])
                if v[1] == 1:
                    x_ab.extend([v[0]])
                    y_ab.extend([v[1]])
                    mask_ab.extend([v[2]])
                    types_ab.extend([v[3]])
            x = x_n + x_ab
            y = y_n + y_ab
            mask = mask_n + mask_ab
            types = types_n + types_ab
        
        return list(x), list(y), list(mask), list(types)
        
        
class CustomFSCopyPasteDataset(Dataset):
    """
    Mvtec train dataset with anomaly samples.
    Anomaly samples: real anomaly samples,
        anomaly samples generated by copy-pasting abnormal regions to normal samples
    """
    def __init__(self, c, is_train=True):
        self.dataset_path = c.data_path
        self.is_train = is_train
        self.cropsize = c.crop_size
        self.anomaly_nums = c.num_anomalies  # number of anomaly samples in each abnormal type
        self.repeat_num = 10  # repeat times for anomaly samples
        self.reuse_times = 5  # real anomaly reuse times
        self.in_fg_region = c.in_fg_region
        
        self.train_list = c.train_list
        # load dataset
        self.n_imgs, self.n_labels, self.n_masks, self.a_imgs, self.a_labels, self.a_masks = self.load_dataset_list()
        self.a_imgs = self.a_imgs * self.repeat_num
        self.a_labels = self.a_labels * self.repeat_num
        self.a_masks = self.a_masks * self.repeat_num
        
        self.labels = np.array(self.n_labels + self.a_labels)
        self.normal_idx = np.argwhere(self.labels == 0).flatten()
        self.anomaly_idx = np.argwhere(self.labels == 1).flatten()
        
        self.img_size = c.img_size
        
        # set transforms
        if is_train:
            self.transform_img = T.Compose([
                T.Resize(c.img_size, Image.ANTIALIAS),
                #T.RandomRotation(5),
                T.CenterCrop(c.crop_size),
                T.ToTensor()])
        # test:
        else:
            self.transform_img = T.Compose([
                T.Resize(c.img_size, Image.ANTIALIAS),
                T.CenterCrop(c.crop_size),
                T.ToTensor()])
        # mask
        self.transform_mask = T.Compose([
            T.Resize(c.img_size, Image.NEAREST),
            T.CenterCrop(c.crop_size),
            T.ToTensor()])
        
        self.augmentors = [A.RandomRotate90(),
                A.Flip(),
                A.Transpose(),
                A.OpticalDistortion(p=1.0, distort_limit=1.0),
                A.OneOf([
                    A.IAAAdditiveGaussianNoise(),
                    A.GaussNoise(),
                ], p=0.2),
                A.OneOf([
                    A.MotionBlur(p=.2),
                    A.MedianBlur(blur_limit=3, p=0.1),
                    A.Blur(blur_limit=3, p=0.1),
                ], p=0.2),
                A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=45, p=0.2),
                A.OneOf([
                    A.OpticalDistortion(p=0.3),
                    A.GridDistortion(p=.1),
                    A.IAAPiecewiseAffine(p=0.3),
                ], p=0.2),
                A.OneOf([
                    A.CLAHE(clip_limit=2),
                    A.IAASharpen(),
                    A.IAAEmboss(),
                    A.RandomBrightnessContrast(),            
                ], p=0.3),
                A.HueSaturationValue(p=0.3)]

        self.normalize = T.Compose([T.Normalize(c.norm_mean, c.norm_std)])
    
    def __len__(self):
        return len(self.n_imgs) + len(self.a_imgs)

    def __getitem__(self, idx):
        if idx >= len(self.n_imgs):  # anomaly samples
            idx_ = idx - len(self.n_imgs)
            img, label, mask = self.a_imgs[idx_], self.a_labels[idx_], self.a_masks[idx_]
            if idx >= len(self.n_imgs) + self.anomaly_nums * self.reuse_times:
                # generating anomaly sample by copy-pasting
                img, mask = self.copy_paste(img, mask, self.img_size)
                img, mask = Image.fromarray(img), Image.fromarray(mask)
                # img.save('aug_imgs/gen_img.jpg')
                # mask.save('aug_imgs/gen_mask.png')
                img = self.normalize(self.transform_img(img))
                mask = self.transform_mask(mask)

                return img, label, mask
        else:  # normal samples
            img, label, mask = self.n_imgs[idx], self.n_labels[idx], self.n_masks[idx]
        img = Image.open(img)
        #
        img = self.normalize(self.transform_img(img))
        #
        if label == 0:
            mask = torch.zeros([1, self.cropsize[0], self.cropsize[1]])
        else:
            mask = Image.open(mask)
            mask = self.transform_mask(mask)
        
        return img, label, mask
    
    def randAugmenter(self):
        aug_ind = np.random.choice(np.arange(len(self.augmentors)), 3, replace=False)
        aug = A.Compose([self.augmentors[aug_ind[0]],
                         self.augmentors[aug_ind[1]],
                         self.augmentors[aug_ind[2]]])
        # aug = A.Compose([self.augmentors[0], A.GridDistortion(p=1.0), self.augmentors[3], self.augmentors[1], self.augmentors[7]])
        return aug

    def copy_paste(self, img, mask, img_size):
        n_idx = np.random.randint(len(self.n_imgs))  # get a random normal sample
        aug = self.randAugmenter()

        image = cv2.imread(img)  # anomaly sample
        image = cv2.resize(image,img_size)
        # temp_img = Image.open(img)
        # temp_img.save("aug_imgs/ano_img.jpg")
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # (900, 900, 3)
        n_image = cv2.imread(self.n_imgs[n_idx])  # normal sample
        n_image = cv2.resize(n_image,img_size)
        n_image = cv2.cvtColor(n_image, cv2.COLOR_BGR2RGB)  # (900, 900, 3)
        # temp = Image.open(self.n_imgs[n_idx])
        # temp.save('aug_imgs/nor_img.jpg')
        img_height, img_width = n_image.shape[0], n_image.shape[1]

        mask = Image.open(mask)
        mask = mask.resize(img_size)
        mask = np.asarray(mask)  # (900, 900)
        
        
        # augmente the abnormal region
        augmentated = aug(image=image, mask=mask)
        aug_image, aug_mask = augmentated['image'], augmentated['mask']
        # temp_img = Image.fromarray(aug_image)
        # temp_img.save("aug_imgs/ano_aug_img.jpg")
        # crop_img = aug_image.copy()
        # crop_img[aug_mask == 0] = 0
        # crop_img = Image.fromarray(crop_img)
        # crop_img.save('aug_imgs/crop_img.jpg')
        # if self.in_fg_region:
        #     n_img_path = self.n_imgs[n_idx]
        #     img_file = n_img_path.split('/')[-1]
        #     fg_path = os.path.join(f'fg_mask/{self.class_name}', img_file)
        #     fg_mask = Image.open(fg_path)
        #     #fg_mask.save("aug_imgs/fg_mask.jpg")
        #     fg_mask = np.asarray(fg_mask)
            
        #     intersect_mask = np.logical_and(fg_mask == 255, aug_mask == 255)
        #     if (np.sum(intersect_mask) > int(2 / 3 * np.sum(aug_mask == 255))):
        #         # when most part of aug_mask is in the fg_mask region 
        #         # copy the augmentated anomaly area to the normal image
        #         n_image[aug_mask == 255, :] = aug_image[aug_mask == 255, :]
        #         return n_image, aug_mask
        #     else:
        #         contours, _ = cv2.findContours(aug_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        #         center_xs, center_ys = [], []
        #         widths, heights = [], []
        #         for i in range(len(contours)):
        #             M = cv2.moments(contours[i])
        #             if M['m00'] == 0:  # error case
        #                 x_min, x_max = np.min(contours[i][:, :, 0]), np.max(contours[i][:, :, 0])
        #                 y_min, y_max = np.min(contours[i][:, :, 1]), np.max(contours[i][:, :, 1])
        #                 center_x = int((x_min + x_max) / 2)
        #                 center_y = int((y_min + y_max) / 2)
        #             else:
        #                 center_x = int(M["m10"] / M["m00"])
        #                 center_y = int(M["m01"] / M["m00"])
        #             center_xs.append(center_x)
        #             center_ys.append(center_y)
        #             x_min, x_max = np.min(contours[i][:, :, 0]), np.max(contours[i][:, :, 0])
        #             y_min, y_max = np.min(contours[i][:, :, 1]), np.max(contours[i][:, :, 1])
        #             width, height = x_max - x_min, y_max - y_min
        #             widths.append(width)
        #             heights.append(height)
        #         if len(widths) == 0 or len(heights) == 0:  # no contours
        #             n_image[aug_mask == 255, :] = aug_image[aug_mask == 255, :]
        #             return n_image, aug_mask
        #         else:
        #             max_width, max_height = np.max(widths), np.max(heights)
        #             center_mask = np.zeros((img_height, img_width), dtype=np.uint8)
        #             center_mask[int(max_height/2):img_height-int(max_height/2), int(max_width/2):img_width-int(max_width/2)] = 255
        #             fg_mask = np.logical_and(fg_mask == 255, center_mask == 255)

        #             x_coord = np.arange(0, img_width)
        #             y_coord = np.arange(0, img_height)
        #             xx, yy = np.meshgrid(x_coord, y_coord)
        #             # coordinates of fg region points
        #             xx_fg = xx[fg_mask]
        #             yy_fg = yy[fg_mask]
        #             xx_yy_fg = np.stack([xx_fg, yy_fg], axis=-1)  # (N, 2)
                    
        #             if xx_yy_fg.shape[0] == 0:  # no fg
        #                 n_image[aug_mask == 255, :] = aug_image[aug_mask == 255, :]
        #                 return n_image, aug_mask

        #             aug_mask_shifted = np.zeros((img_height, img_width), dtype=np.uint8)
        #             for i in range(len(contours)):
        #                 aug_mask_shifted_i = np.zeros((img_height, img_width), dtype=np.uint8)
        #                 new_aug_mask_i = np.zeros((img_height, img_width), dtype=np.uint8)
        #                 # random generate a point in the fg region
        #                 idx = np.random.choice(np.arange(xx_yy_fg.shape[0]), 1)
        #                 rand_xy = xx_yy_fg[idx]
        #                 delta_x, delta_y = center_xs[i] - rand_xy[0, 0], center_ys[i] - rand_xy[0, 1]
                        
        #                 x_min, x_max = np.min(contours[i][:, :, 0]), np.max(contours[i][:, :, 0])
        #                 y_min, y_max = np.min(contours[i][:, :, 1]), np.max(contours[i][:, :, 1])
                        
        #                 # mask for one anomaly region
        #                 aug_mask_i = np.zeros((img_height, img_width), dtype=np.uint8)
        #                 aug_mask_i[y_min:y_max, x_min:x_max] = 255
        #                 aug_mask_i = np.logical_and(aug_mask == 255, aug_mask_i == 255)
                        
        #                 # coordinates of orginal mask points
        #                 xx_ano, yy_ano = xx[aug_mask_i], yy[aug_mask_i]
                        
        #                 # shift the original mask into fg region
        #                 xx_ano_shifted = xx_ano - delta_x
        #                 yy_ano_shifted = yy_ano - delta_y
        #                 outer_points_x = np.logical_or(xx_ano_shifted < 0, xx_ano_shifted >= img_width) 
        #                 outer_points_y = np.logical_or(yy_ano_shifted < 0, yy_ano_shifted >= img_height)
        #                 outer_points = np.logical_or(outer_points_x, outer_points_y)
                        
        #                 # keep points in image
        #                 xx_ano_shifted = xx_ano_shifted[~outer_points]
        #                 yy_ano_shifted = yy_ano_shifted[~outer_points]
        #                 aug_mask_shifted_i[yy_ano_shifted, xx_ano_shifted] = 255
                        
        #                 # original points should be changed
        #                 xx_ano = xx_ano[~outer_points]
        #                 yy_ano = yy_ano[~outer_points]
        #                 new_aug_mask_i[yy_ano, xx_ano] = 255
        #                 # copy the augmentated anomaly area to the normal image
        #                 n_image[aug_mask_shifted_i == 255, :] = aug_image[new_aug_mask_i == 255, :]
        #                 aug_mask_shifted[aug_mask_shifted_i == 255] = 255
        #             return n_image, aug_mask_shifted
        # else:  # no fg restriction
        #     # copy the augmentated anomaly area to the normal image
        n_image[aug_mask == 255, :] = aug_image[aug_mask == 255, :]

        return n_image, aug_mask

    def load_dataset_list(self):
        n_img_paths, n_labels, n_mask_paths = [], [], []  # normal
        a_img_paths, a_labels, a_mask_paths = [], [], []  # abnormal

        for t in self.train_list:
            if t[1] == 0:
                n_img_paths.extend([t[0]])
                n_labels.extend([t[1]])
                n_mask_paths.extend([t[2]])
            if t[1] == 1:
                a_img_paths.extend([t[0]])
                a_labels.extend([t[1]])
                a_mask_paths.extend([t[2]])
        
        
        return n_img_paths, n_labels, n_mask_paths, a_img_paths, a_labels, a_mask_paths



def excluding_images(images, labels, masks, img_types, excluding_images):
    retain_images = []
    retain_labels = []
    retain_masks = []
    retain_img_types = []
    for idx, image in enumerate(images):
        if image in excluding_images:
            continue
        retain_images.append(image)
        retain_labels.append(labels[idx])
        retain_masks.append(masks[idx])
        retain_img_types.append(img_types[idx])
    
    return retain_images, retain_labels, retain_masks, retain_img_types