from __future__ import print_function
import torch.utils.data as data
from PIL import Image
import os, time
import os.path
from multiprocessing import Pool
from functools import partial
from gibson2.core.render.profiler import Profiler
import errno
import torch
import json
import codecs
import cv2
import numpy as np
import ctypes as ct
import sys
from tqdm import tqdm
import torchvision.transforms as transforms
import argparse
import json
from numpy.linalg import inv
import pickle
import gibson2

IMG_EXTENSIONS = [
    '.jpg',
    '.JPG',
    '.jpeg',
    '.JPEG',
    '.png',
    '.PNG',
    '.ppm',
    '.PPM',
    '.bmp',
    '.BMP',
]


def is_image_file(filename):
    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)


def default_loader(path):
    ## Heavy usage
    img = cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB)    #.convert('RGB')
    #img = Image.open(path)
    return img


def depth_loader(path):
    ## Heavy usage
    ## TODO: Image.open for depth image is main data loading bottleneck
    #img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)#.convert('I')
    img = Image.open(path)
    return img





def get_item_fn(inds, select, root, loader, transform, off_3d, target_transform, depth_trans,
                off_pc_render, dll, train, require_rgb):
    """ Functional programming version of Dataset.__getitem__
    The advantage is that it is pickle-friendly and supports python multiprocessing
    
    Argument: 
        inds: tuple of scene index and output index
    """
    index, out_i = inds
    scene = select[index][0][0]
    uuids = [item[1] for item in select[index]]
    paths = ([
        os.path.join(root, scene, 'pano', 'rgb',
                     "point_" + item + "_view_equirectangular_domain_rgb.png") for item in uuids
    ])
    mist_paths = ([
        os.path.join(root, scene, 'pano', 'mist',
                     "point_" + item + "_view_equirectangular_domain_mist.png") for item in uuids
    ])
    normal_paths = ([
        os.path.join(root, scene, 'pano', 'normal',
                     "point_" + item + "_view_equirectangular_domain_normal.png") for item in uuids
    ])
    pose_paths = ([
        os.path.join(root, scene, 'pano', 'points', "point_" + item + ".json") for item in uuids
    ])
    semantic_paths = ([
        os.path.join(root, scene, 'pano', 'semantic',
                     "point_" + item + "_view_equirectangular_domain_semantic.png")
        for item in uuids
    ])
    poses = []
    for i, item in enumerate(pose_paths):
        f = open(item)
        pose_dict = json.load(f)
        p = np.concatenate(np.array(pose_dict[1][u'camera_rt_matrix'])).astype(np.float32).reshape(
            (4, 4))
        rotation = np.array([[0, 1, 0, 0], [0, 0, 1, 0], [-1, 0, 0, 0], [0, 0, 0, 1]])
        p = np.dot(p, rotation)
        poses.append(p)
        f.close()

    img_paths = paths[1:]
    target_path = paths[0]
    img_poses = poses[1:]
    target_pose = poses[0]

    mist_img_paths = mist_paths[1:]
    mist_target_path = mist_paths[0]

    normal_img_paths = normal_paths[1:]
    normal_target_path = normal_paths[0]

    semantic_img_paths = semantic_paths[1:]
    semantic_target_path = semantic_paths[0]
    poses_relative = []

    semantic_imgs = None
    semantic_target = None

    normal_imgs = None
    normal_target = None

    mist_imgs = None
    mist_target = None

    imgs = None
    target = None

    for pose_i, item in enumerate(img_poses):
        pose_i = pose_i + 1
        relative = np.dot(inv(target_pose), item)
        poses_relative.append(torch.from_numpy(relative))

    if require_rgb:
        imgs = [loader(item) for item in img_paths]
        target = loader(target_path)
        org_img = imgs[0].copy()

    if not off_3d and require_rgb:
        mist_imgs = [depth_loader(item) for item in mist_img_paths]
        mist_target = depth_loader(mist_target_path)
        if train:
            normal_imgs = [loader(item) for item in normal_img_paths]
            normal_target = loader(normal_target_path)

    if not transform is None and require_rgb:
        imgs = [transform(item) for item in imgs]
    if not target_transform is None and require_rgb:
        target = target_transform(target)

    if not off_3d and require_rgb:
        mist_imgs = [
            np.expand_dims(np.array(item).astype(np.float32) / 65536.0, 2) for item in mist_imgs
        ]
        org_mist = mist_imgs[0][:, :, 0].copy()
        mist_target = np.expand_dims(np.array(mist_target).astype(np.float32) / 65536.0, 2)

        if not depth_trans is None:
            mist_imgs = [depth_trans(item) for item in mist_imgs]
        if not depth_trans is None:
            mist_target = depth_trans(mist_target)

        if train:
            if not transform is None:
                normal_imgs = [transform(item) for item in normal_imgs]
            if not target_transform is None:
                normal_target = target_transform(normal_target)

    if not off_pc_render and require_rgb:
        img = np.array(org_img)
        h, w, _ = img.shape
        render = np.zeros((h, w, 3), dtype='uint8')
        target_depth = np.zeros((h, w)).astype(np.float32)
        depth = org_mist
        pose = poses_relative[0].numpy()
        dll.render(ct.c_int(img.shape[0]), ct.c_int(img.shape[1]), img.ctypes.data_as(ct.c_void_p),
                   depth.ctypes.data_as(ct.c_void_p), pose.ctypes.data_as(ct.c_void_p),
                   render.ctypes.data_as(ct.c_void_p), target_depth.ctypes.data_as(ct.c_void_p))
        if not transform is None:
            render = transform(Image.fromarray(render))
        if not depth_trans is None:
            target_depth = depth_trans(np.expand_dims(target_depth, 2))

    if off_3d:
        out = (imgs, target, poses_relative)
    elif off_pc_render:
        out = (imgs, target, mist_imgs, mist_target, normal_imgs, normal_target, poses_relative)
    else:
        out = (imgs, target, mist_imgs, mist_target, normal_imgs, normal_target, poses_relative,
               render, target_depth)
    return (out_i, out)


class ViewDataSet3D(data.Dataset):
    def __init__(self,
                 root=None,
                 train=False,
                 transform=None,
                 mist_transform=None,
                 loader=default_loader,
                 seqlen=5,
                 debug=False,
                 dist_filter=None,
                 off_3d=True,
                 off_pc_render=True,
                 overwrite_fofn=False,
                 semantic_transform=np.array,
                 env=None,
                 only_load=None):
        print('Processing the data:')
        if not root:
            self.root = gibson2.dataset_path
        else:
            self.root = root
        self.train = train
        self.env = env
        self.loader = loader
        self.seqlen = seqlen
        self.transform = transform
        self.target_transform = transform
        self.depth_trans = mist_transform
        self.semantic_trans = semantic_transform
        self._require_semantics = "SEMANTICS" in self.env.config["ui_components"]
        self._require_rgb = "RGB_FILLED" in self.env.config[
            "ui_components"] or "RGB_PREFILLED" in self.env.config[
                "ui_components"] or "rgb_filled" in self.env.config[
                    "output"] or "rgb_prefill" in self.env.config["output"]
        self.off_3d = off_3d
        self.select = []
        self.fofn = self.root + '_fofn' + str(int(train)) + '.pkl'
        self.off_pc_render = off_pc_render
        self.dll = None

        if not self.off_pc_render:
            self.dll = np.ctypeslib.load_library('render', '.')

        if overwrite_fofn or not os.path.isfile(self.fofn):

            if only_load is None:
                self.scenes = sorted([
                    d for d in (os.listdir(self.root)) if os.path.isdir(os.path.join(self.root, d))
                    and os.path.isfile(os.path.join(self.root, d, 'camera_poses.csv'))
                    and os.path.isdir(os.path.join(self.root, d, 'pano'))
                ])

                num_scenes = len(self.scenes)
                num_train = int(num_scenes * 0.9)

            else:

                self.scenes = sorted([only_load])
                num_scenes = 1
                num_train = 0

            print("Total %d scenes %d train %d test" %
                  (num_scenes, num_train, num_scenes - num_train))
            if train:
                self.scenes = self.scenes[:num_train]

            self.meta = {}

            last = len(self.scenes)

            for scene in self.scenes[:last]:
                posefile = os.path.join(self.root, scene, 'camera_poses.csv')
                with open(posefile) as f:
                    for line in f:
                        l = line.strip().split(',')
                        uuid = l[0]
                        xyz = list(map(float, l[1:4]))
                        quat = list(map(float, l[4:8]))
                        if not scene in self.meta:
                            self.meta[scene] = {}
                        metadata = (uuid, xyz, quat)
                        # print(uuid, xyz)
                        if os.path.isfile(
                                os.path.join(self.root, scene, 'pano', 'points',
                                             'point_' + uuid + '.json')):
                            if np.linalg.norm(np.array(xyz) - np.array([0, 0, 0])
                                              ) > 1e-5:    #remove scans that are not registered
                                self.meta[scene][uuid] = metadata
            print("Indexing")

            for scene, meta in tqdm(list(self.meta.items())):
                if len(meta) < self.seqlen:
                    continue
                for uuid, v in list(meta.items()):
                    dist_list = [(uuid2, np.linalg.norm(np.array(v2[1]) - np.array(v[1])))
                                 for uuid2, v2 in list(meta.items())]
                    dist_list = sorted(dist_list, key=lambda x: x[-1])

                    if not dist_filter is None:
                        if dist_list[1][-1] < dist_filter:
                            self.select.append([[scene, dist_list[i][0], dist_list[i][1]]
                                                for i in range(self.seqlen)])

                    else:
                        self.select.append([[scene, dist_list[i][0], dist_list[i][1]]
                                            for i in range(self.seqlen)])

            with open(self.fofn, 'wb') as fp:
                pickle.dump([self.scenes, self.meta, self.select, num_scenes, num_train], fp)

        else:
            with open(self.fofn, 'rb') as fp:
                self.scenes, self.meta, self.select, num_scenes, num_train = pickle.load(fp)
                print("Total %d scenes %d train %d test" %
                      (num_scenes, num_train, num_scenes - num_train))

    def get_scene_info(self, index):
        scene = self.scenes[index]
        data = [(i, item) for i, item in enumerate(self.select) if item[0][0] == scene]
        uuids = ([(item[1][0][1], item[0]) for item in data])

        pose_paths = ([
            os.path.join(self.root, scene, 'pano', 'points', "point_" + item[0] + ".json")
            for item in uuids
        ])
        poses = []
        for item in pose_paths:
            f = open(item)
            pose_dict = json.load(f)
            p = np.concatenate(np.array(pose_dict[1][u'camera_rt_matrix'])).astype(
                np.float32).reshape((4, 4))

            rotation = np.array([[0, 1, 0, 0], [0, 0, 1, 0], [-1, 0, 0, 0], [0, 0, 0, 1]])

            p = np.dot(p, rotation)
            poses.append(p)
            f.close()

        return uuids, poses

    def __getitem__(self, index):
        scene = self.select[index][0][0]
        uuids = [item[1] for item in self.select[index]]
        paths = ([
            os.path.join(self.root, scene, 'pano', 'rgb',
                         "point_" + item + "_view_equirectangular_domain_rgb.png") for item in uuids
        ])
        mist_paths = ([
            os.path.join(self.root, scene, 'pano', 'mist',
                         "point_" + item + "_view_equirectangular_domain_mist.png")
            for item in uuids
        ])
        normal_paths = ([
            os.path.join(self.root, scene, 'pano', 'normal',
                         "point_" + item + "_view_equirectangular_domain_normal.png")
            for item in uuids
        ])
        pose_paths = ([
            os.path.join(self.root, scene, 'pano', 'points', "point_" + item + ".json")
            for item in uuids
        ])
        semantic_paths = ([
            os.path.join(self.root, scene, 'pano', 'semantic',
                         "point_" + item + "_view_equirectangular_domain_semantic.png")
            for item in uuids
        ])
        poses = []
        for i, item in enumerate(pose_paths):
            f = open(item)
            pose_dict = json.load(f)
            p = np.concatenate(np.array(pose_dict[1][u'camera_rt_matrix'])).astype(
                np.float32).reshape((4, 4))
            rotation = np.array([[0, 1, 0, 0], [0, 0, 1, 0], [-1, 0, 0, 0], [0, 0, 0, 1]])
            p = np.dot(p, rotation)
            poses.append(p)
            f.close()

        img_paths = paths[1:]
        target_path = paths[0]
        img_poses = poses[1:]
        target_pose = poses[0]

        mist_img_paths = mist_paths[1:]
        mist_target_path = mist_paths[0]

        normal_img_paths = normal_paths[1:]
        normal_target_path = normal_paths[0]

        semantic_img_paths = semantic_paths[1:]
        semantic_target_path = semantic_paths[0]
        poses_relative = []

        semantic_imgs = None
        semantic_target = None

        normal_imgs = None
        normal_target = None

        mist_imgs = None
        mist_target = None

        imgs, target = None, None

        for pose_i, item in enumerate(img_poses):
            pose_i = pose_i + 1
            relative = np.dot(inv(target_pose), item)
            poses_relative.append(torch.from_numpy(relative))

        if self._require_rgb:
            imgs = [self.loader(item) for item in img_paths]
            target = self.loader(target_path)

        if not self.off_3d and self._require_rgb:
            mist_imgs = [depth_loader(item) for item in mist_img_paths]
            mist_target = depth_loader(mist_target_path)
            if self.train:    # Optimize
                normal_imgs = [self.loader(item) for item in normal_img_paths]
                normal_target = self.loader(normal_target_path)

        if not self.off_pc_render and self._require_rgb:
            org_img = imgs[0].copy()
            if not self.transform is None:
                imgs = [self.transform(item) for item in imgs]
            if not self.target_transform is None:
                target = self.target_transform(target)

        if not self.off_3d and self._require_rgb:
            mist_imgs = [
                np.expand_dims(np.array(item).astype(np.float32) / 65536.0, 2) for item in mist_imgs
            ]
            if not self.off_pc_render:
                org_mist = mist_imgs[0][:, :, 0].copy()
            mist_target = np.expand_dims(np.array(mist_target).astype(np.float32) / 65536.0, 2)

            if not self.depth_trans is None:
                mist_imgs = [self.depth_trans(item) for item in mist_imgs]
            if not self.depth_trans is None:
                mist_target = self.depth_trans(mist_target)

            if self.train:
                if not self.transform is None:
                    normal_imgs = [self.transform(item) for item in normal_imgs]
                if not self.target_transform is None:
                    normal_target = self.target_transform(normal_target)

        if not self.off_pc_render and self._require_rgb:
            img = np.array(org_img)
            h, w, _ = img.shape
            render = np.zeros((h, w, 3), dtype='uint8')
            target_depth = np.zeros((h, w)).astype(np.float32)
            depth = org_mist
            pose = poses_relative[0].numpy()
            self.dll.render(ct.c_int(img.shape[0]), ct.c_int(img.shape[1]),
                            img.ctypes.data_as(ct.c_void_p), depth.ctypes.data_as(ct.c_void_p),
                            pose.ctypes.data_as(ct.c_void_p), render.ctypes.data_as(ct.c_void_p),
                            target_depth.ctypes.data_as(ct.c_void_p))
            if not self.transform is None:
                render = self.transform(Image.fromarray(render))
            if not self.depth_trans is None:
                target_depth = self.depth_trans(np.expand_dims(target_depth, 2))

        if self.off_3d:
            return imgs, target, poses_relative
        elif self.off_pc_render:
            return imgs, target, mist_imgs, mist_target, normal_imgs, normal_target, poses_relative
        else:
            return imgs, target, mist_imgs, mist_target, normal_imgs, normal_target, poses_relative, render, target_depth

    def get_multi_index(self, uuids):
        indices = range(len(uuids))
        p = Pool(16)
        partial_fn = partial(get_item_fn,
                             select=self.select,
                             root=self.root,
                             loader=self.loader,
                             transform=self.transform,
                             off_3d=self.off_3d,
                             target_transform=self.target_transform,
                             depth_trans=self.depth_trans,
                             off_pc_render=self.off_pc_render,
                             dll=self.dll,
                             train=self.train,
                             require_rgb=self._require_rgb)
        mapped_pairs = list(tqdm(p.imap(partial_fn, list(zip(uuids, indices))), total=len(uuids)))
        sorted_pairs = sorted(mapped_pairs, key=lambda x: x[0])
        out_data = [key_pair[1] for key_pair in sorted_pairs]
        p.close()
        p.join()
        return out_data

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


########### BELOW THIS POINT: Legacy code #################
########### KEEPING ONLY FOR REFERENCE ####################


class Places365Dataset(data.Dataset):
    def __init__(self, root, train=True, transform=None, loader=default_loader):
        self.root = root.rstrip('/')
        self.train = train
        self.fns = []
        self.fofn = os.path.basename(root) + '_fofn' + str(int(train)) + '.pkl'
        self.loader = loader
        self.transform = transform
        if not os.path.isfile(self.fofn):
            for subdir, dirs, files in os.walk(self.root):
                if self.train:
                    files = files[:len(files) / 10 * 9]
                else:
                    files = files[len(files) / 10 * 9:]
                print(subdir)
                for file in files:
                    self.fns.append(os.path.join(subdir, file))
            with open(self.fofn, 'wb') as fp:
                pickle.dump(self.fns, fp)
        else:
            with open(self.fofn, 'rb') as fp:
                self.fns = pickle.load(fp)

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

    def __getitem__(self, index):
        path = self.fns[index]
        img = self.loader(path)
        if not self.transform is None:
            img = self.transform(img)
        return img


class PairDataset(data.Dataset):
    def __init__(self, root, train=True, transform=None, mist_transform=None, loader=np.load):
        self.root = root.rstrip('/')
        self.train = train
        self.fns = []
        self.fofn = os.path.basename(root) + '_fofn' + str(int(train)) + '.pkl'
        self.loader = loader
        self.transform = transform
        self.mist_transform = mist_transform
        if not os.path.isfile(self.fofn):
            for subdir, dirs, files in os.walk(self.root):
                if self.train:
                    files = files[:len(files) / 10 * 9]
                else:
                    files = files[len(files) / 10 * 9:]
                print(subdir)
                for file in files:
                    if file[-3:] == 'npz':
                        self.fns.append(os.path.join(subdir, file))
            with open(self.fofn, 'wb') as fp:
                pickle.dump(self.fns, fp)
        else:
            with open(self.fofn, 'rb') as fp:
                self.fns = pickle.load(fp)

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

    def __getitem__(self, index):
        path = self.fns[index]
        data = self.loader(path)

        try:
            source, depth, target = data['source'], data['depth'], data['target']
            # print(source.shape, depth.shape, target.shape)
        except:
            source = np.zeros((1024, 2048, 3)).astype(np.uint8)
            target = np.zeros((1024, 2048, 3)).astype(np.uint8)
            depth = np.zeros((1024, 2048)).astype(np.float32)

        if not self.transform is None:
            source = self.transform(source)
            target = self.transform(target)
            # depth = self.mist_transform(depth)
            depth = torch.from_numpy(depth.astype(np.float32))
        return source, depth, target


class RenderPairDataset(data.Dataset):
    def __init__(self, root, train=True, transform=None):
        self.root = root.rstrip('/')
        self.train = train
        self.fns = []
        self.fofn = os.path.basename(root) + '_fofn' + str(int(train)) + '.pkl'
        self.transform = transform
        if not os.path.isfile(self.fofn):
            for subdir, dirs, files in os.walk(self.root):
                if self.train:
                    files = files[:len(files) // 10 * 9]
                else:
                    files = files[len(files) // 10 * 9:]
                print(subdir)
                for file in files:
                    if file[-3:] == 'png' and file[:3] == 'img':
                        self.fns.append(
                            (os.path.join(subdir, file), os.path.join(subdir, 'render' + file[3:])))
            with open(self.fofn, 'wb') as fp:
                pickle.dump(self.fns, fp)
        else:
            with open(self.fofn, 'rb') as fp:
                self.fns = pickle.load(fp)

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

    def __getitem__(self, index):
        img_path, render_path = self.fns[index]

        source = Image.open(render_path).convert('RGB')
        target = Image.open(img_path).convert('RGB')

        source = self.transform(source)
        target = self.transform(target)
        return source, target


if __name__ == '__main__':
    print('test')
    parser = argparse.ArgumentParser()
    parser.add_argument('--debug', action='store_true', help='debug mode')
    parser.add_argument('--dataroot', required=True, help='path to dataset')
    parser.add_argument('--dataset', required=True, help='dataset type')
    opt = parser.parse_args()

    if opt.dataset == 'view3d':
        d = ViewDataSet3D(root=opt.dataroot,
                          debug=opt.debug,
                          seqlen=2,
                          dist_filter=0.8,
                          off_3d=False,
                          off_pc_render=False)
        print(len(d))
        sample = (d[1])
        print(sample)
        if sample is not None:
            print('3d test passed')

        uuids, xyzs, poses = d.get_scene_info(0)
        print(uuids, xyzs, poses)

    elif opt.dataset == 'places365':
        d = Places365Dataset(root=opt.dataroot)
        print(len(d))
        sample = d[0]
        print(sample)
        if sample is not None:
            print('places 365 test passed')

    elif opt.dataset == 'pair':
        d = PairDataset(root=opt.dataroot)
        print(len(d))
        sample = d[0]
        print(sample)
        if sample is not None:
            print('pair test passed')

    elif opt.dataset == 'render_pair':
        d = RenderPairDataset(root=opt.dataroot,
                              transform=transforms.Compose(
                                  [transforms.Resize(512),
                                   transforms.ToTensor()]))
        print(len(d))
        sample = d[0]
        print(sample[0].size(), sample[1].size())
