import torch
import torch.utils.data as data
import torchvision.transforms as transforms

import os
import numpy as np
from PIL import Image
from glob import glob
from skimage.segmentation import slic
import torchvision.transforms.functional as TF
from scipy.io import loadmat
import random
import cv2

import sys
sys.path.append("../..")

def label2one_hot_torch(labels, C=14):
    """ Converts an integer label torch.autograd.Variable to a one-hot Variable.

    Args:
      labels(tensor) : segmentation label
      C (integer) : number of classes in labels

    Returns:
      target (tensor) : one-hot vector of the input label

    Shape:
      labels: (B, 1, H, W)
      target: (B, N, H, W)
    """
    b,_, h, w = labels.shape
    one_hot = torch.zeros(b, C, h, w, dtype=torch.long).to(labels)
    target = one_hot.scatter_(1, labels.type(torch.long).data, 1) #require long type

    return target.type(torch.float32)

class Dataset(data.Dataset):
    def __init__(self, data_dir, crop_size = 128, test=False,
                 sp_num = 256, slic = True, preprocess_name = False,
                 gt_label = False, label_path = None, test_time = False,
                 img_path = None):
        super(Dataset, self).__init__()
        ext = ["*.jpg"]
        dl = []
        self.test = test
        self.test_time = test_time

        [dl.extend(glob(data_dir + '/**/' + e, recursive=True)) for e in ext]

        data_list = sorted(dl)
        self.data_list = data_list
        self.sp_num = sp_num
        self.slic = slic

        self.crop = transforms.CenterCrop(size = (crop_size, crop_size))
        self.crop_size = crop_size
        self.test = test

        self.gt_label = gt_label
        if gt_label:
            self.label_path = label_path

        self.img_path = img_path

    def preprocess_label(self, seg):
        segs = label2one_hot_torch(seg.unsqueeze(0), C = seg.max() + 1)
        new_seg = []
        for cnt in range(seg.max() + 1):
            if segs[0, cnt].sum() > 0:
                new_seg.append(segs[0, cnt])
        new_seg = torch.stack(new_seg)
        return torch.argmax(new_seg, dim = 0)

    def __getitem__(self, index):
        if self.img_path is None:
            data_path = self.data_list[index]
        else:
            data_path = self.img_path
        rgb_img = Image.open(data_path)
        imgH, imgW = rgb_img.size

        if self.gt_label:
            img_name = data_path.split("/")[-1].split("_")[0]
            mat_path = os.path.join(self.label_path, data_path.split('/')[-2], img_name.replace('.jpg', '.mat'))
            mat = loadmat(mat_path)
            max_label_num = 0
            final_seg = None
            for i in range(len(mat['groundTruth'][0])):
                seg = mat['groundTruth'][0][i][0][0][0]
                if len(np.unique(seg)) > max_label_num:
                    max_label_num = len(np.unique(seg))
                    final_seg = seg
            seg = torch.from_numpy(final_seg.astype(np.float32))
            segs = seg.long().unsqueeze(0)

        if self.img_path is None:
            i, j, h, w = transforms.RandomCrop.get_params(rgb_img, output_size=(self.crop_size, self.crop_size))
        else:
            i = 40; j = 40; h = self.crop_size; w = self.crop_size
        rgb_img = TF.crop(rgb_img, i, j, h, w)
        if self.gt_label:
            segs = TF.crop(segs, i, j, h, w)
            segs = self.preprocess_label(segs)

        if self.slic:
            sp_num = self.sp_num
            # compute superpixel
            slic_i = slic(np.array(rgb_img), n_segments=sp_num, compactness=10, start_label=0, min_size_factor=0.3)
            slic_i = torch.from_numpy(slic_i)
            slic_i[slic_i >= sp_num] = sp_num - 1
            oh = label2one_hot_torch(slic_i.unsqueeze(0).unsqueeze(0), C = sp_num).squeeze()

        rgb_img = TF.to_tensor(rgb_img)
        if rgb_img.shape[0] == 1:
            rgb_img = rgb_img.repeat(3, 1, 1)
        rgb_img = rgb_img[:3, :, :]

        rets = []
        rets.append(rgb_img)
        if self.slic:
            rets.append(oh)
        rets.append(data_path.split("/")[-1])
        rets.append(index)
        if self.gt_label:
            rets.append(segs.view(1, segs.shape[-2], segs.shape[-1]))
        return rets

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

if __name__ == '__main__':
    import torchvision.utils as vutils
    dataset = Dataset('/home/xtli/DATA/texture_data/',
                      sampled_num=3000)
    loader_ = torch.utils.data.DataLoader(dataset     = dataset,
                                         batch_size  = 1,
                                         shuffle     = True,
                                         num_workers = 1,
                                         drop_last   = True)
    loader = iter(loader_)
    img, points, pixs = loader.next()

    crop_size = 128
    canvas = torch.zeros((1, 3, crop_size, crop_size))
    for i in range(points.shape[-2]):
        p = (points[0, i] + 1) / 2.0 * (crop_size - 1)
        canvas[0, :, int(p[0]), int(p[1])] = pixs[0, :, i]
    vutils.save_image(canvas, 'canvas.png')
    vutils.save_image(img, 'img.png')