# Data loading based on https://github.com/NVIDIA/flownet2-pytorch

import numpy as np
import torch
import torch.utils.data as data
import torch.nn.functional as F

import os
import cv2
import math
import random
from glob import glob
import os.path as osp

from utils import frame_utils
from utils.augmentor import FlowAugmentor, FlowAugmentorKITTI


class CombinedDataset(data.Dataset):
    def __init__(self, datasets):
        self.datasets = datasets

    def __len__(self):
        length = 0 
        for i in range(len(self.datasets)):
            length += len(self.datsaets[i])
        return length

    def __getitem__(self, index):
        i = 0
        for j in range(len(self.datasets)):
            if i + len(self.datasets[j]) >= index:
                yield self.datasets[j][index-i]
                break
            i += len(self.datasets[j])

    def __add__(self, other):
        self.datasets.append(other)
        return self

class FlowDataset(data.Dataset):
    def __init__(self, args, image_size=None, do_augument=False):
        self.image_size = image_size
        self.do_augument = do_augument

        if self.do_augument:
            self.augumentor = FlowAugmentor(self.image_size)

        self.flow_list = []
        self.image_list = []

        self.init_seed = False

    def __getitem__(self, index):

        if not self.init_seed:
            worker_info = torch.utils.data.get_worker_info()
            if worker_info is not None:
                torch.manual_seed(worker_info.id)
                np.random.seed(worker_info.id)
                random.seed(worker_info.id)
                self.init_seed = True

        index = index % len(self.image_list)
        flow = frame_utils.read_gen(self.flow_list[index])
        img1 = frame_utils.read_gen(self.image_list[index][0])
        img2 = frame_utils.read_gen(self.image_list[index][1])

        img1 = np.array(img1).astype(np.uint8)[..., :3]
        img2 = np.array(img2).astype(np.uint8)[..., :3]
        flow = np.array(flow).astype(np.float32)

        if self.do_augument:
            img1, img2, flow = self.augumentor(img1, img2, flow)

        img1 = torch.from_numpy(img1).permute(2, 0, 1).float()
        img2 = torch.from_numpy(img2).permute(2, 0, 1).float()
        flow = torch.from_numpy(flow).permute(2, 0, 1).float()
        valid = torch.ones_like(flow[0])

        return img1, img2, flow, valid

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

    def __add(self, other):
        return CombinedDataset([self, other])


class MpiSintelTest(FlowDataset):
    def __init__(self, args, root='datasets/Sintel/test', dstype='clean'):
        super(MpiSintelTest, self).__init__(args, image_size=None, do_augument=False)

        self.root = root
        self.dstype = dstype

        image_dir = osp.join(self.root, dstype)
        all_sequences = os.listdir(image_dir)

        self.image_list = []
        for sequence in all_sequences:
            frames = sorted(glob(osp.join(image_dir, sequence, '*.png')))
            for i in range(len(frames)-1):
                self.image_list += [[frames[i], frames[i+1], sequence, i]]

    def __getitem__(self, index):
        img1 = frame_utils.read_gen(self.image_list[index][0])
        img2 = frame_utils.read_gen(self.image_list[index][1])
        sequence = self.image_list[index][2]
        frame = self.image_list[index][3]

        img1 = np.array(img1).astype(np.uint8)[..., :3]
        img2 = np.array(img2).astype(np.uint8)[..., :3]
        img1 = torch.from_numpy(img1).permute(2, 0, 1).float()
        img2 = torch.from_numpy(img2).permute(2, 0, 1).float()
        return img1, img2, sequence, frame


class MpiSintel(FlowDataset):
    def __init__(self, args, image_size=None, do_augument=True, root='datasets/Sintel/training', dstype='clean'):
        super(MpiSintel, self).__init__(args, image_size, do_augument)
        if do_augument:
            self.augumentor.min_scale = -0.2
            self.augumentor.max_scale = 0.7

        self.root = root
        self.dstype = dstype

        flow_root = osp.join(root, 'flow')
        image_root = osp.join(root, dstype)

        file_list = sorted(glob(osp.join(flow_root, '*/*.flo')))
        for flo in file_list:
            fbase = flo[len(flow_root)+1:]
            fprefix = fbase[:-8]
            fnum = int(fbase[-8:-4])

            img1 = osp.join(image_root, fprefix + "%04d"%(fnum+0) + '.png')
            img2 = osp.join(image_root, fprefix + "%04d"%(fnum+1) + '.png')

            if not osp.isfile(img1) or not osp.isfile(img2) or not osp.isfile(flo):
                continue

            self.image_list.append((img1, img2))
            self.flow_list.append(flo)


class FlyingChairs(FlowDataset):
    def __init__(self, args, image_size=None, do_augument=True, root='datasets/FlyingChairs_release/data'):
        super(FlyingChairs, self).__init__(args, image_size, do_augument)
        self.root = root
        self.augumentor.min_scale = -0.2
        self.augumentor.max_scale = 1.0

        images = sorted(glob(osp.join(root, '*.ppm')))
        self.flow_list = sorted(glob(osp.join(root, '*.flo')))
        assert (len(images)//2 == len(self.flow_list))

        self.image_list = []
        for i in range(len(self.flow_list)):
            im1 = images[2*i]
            im2 = images[2*i + 1]
            self.image_list.append([im1, im2])


class SceneFlow(FlowDataset):
    def __init__(self, args, image_size, do_augument=True, root='datasets',
            dstype='frames_cleanpass', use_flyingthings=True, use_monkaa=False, use_driving=False):
        
        super(SceneFlow, self).__init__(args, image_size, do_augument)
        self.root = root
        self.dstype = dstype

        self.augumentor.min_scale = -0.2
        self.augumentor.max_scale = 0.8

        if use_flyingthings:
            self.add_flyingthings()
        
        if use_monkaa:
            self.add_monkaa()

        if use_driving:
            self.add_driving()

    def add_flyingthings(self):
        root = osp.join(self.root, 'FlyingThings3D')

        for cam in ['left']:
            for direction in ['into_future', 'into_past']:
                image_dirs = sorted(glob(osp.join(root, self.dstype, 'TRAIN/*/*')))
                image_dirs = sorted([osp.join(f, cam) for f in image_dirs])

                flow_dirs = sorted(glob(osp.join(root, 'optical_flow/TRAIN/*/*')))
                flow_dirs = sorted([osp.join(f, direction, cam) for f in flow_dirs])

                for idir, fdir in zip(image_dirs, flow_dirs):
                    images = sorted(glob(osp.join(idir, '*.png')) )
                    flows = sorted(glob(osp.join(fdir, '*.pfm')) )
                    for i in range(len(flows)-1):
                        if direction == 'into_future':
                            self.image_list += [[images[i], images[i+1]]]
                            self.flow_list += [flows[i]]
                        elif direction == 'into_past':
                            self.image_list += [[images[i+1], images[i]]]
                            self.flow_list += [flows[i+1]]
      
    def add_monkaa(self):
        pass # we don't use monkaa

    def add_driving(self):
        pass # we don't use driving


class KITTI(FlowDataset):
    def __init__(self, args, image_size=None, do_augument=True, is_test=False, is_val=False, do_pad=False, split=True, root='datasets/KITTI'):
        super(KITTI, self).__init__(args, image_size, do_augument)
        self.root = root
        self.is_test = is_test
        self.is_val = is_val
        self.do_pad = do_pad

        if self.do_augument:
            self.augumentor = FlowAugmentorKITTI(self.image_size, min_scale=-0.2, max_scale=0.5)

        if self.is_test:
            images1 = sorted(glob(os.path.join(root, 'testing', 'image_2/*_10.png')))
            images2 = sorted(glob(os.path.join(root, 'testing', 'image_2/*_11.png')))
            for i in range(len(images1)):
                self.image_list += [[images1[i], images2[i]]]

        else:
            flows = sorted(glob(os.path.join(root, 'training', 'flow_occ/*_10.png')))
            images1 = sorted(glob(os.path.join(root, 'training', 'image_2/*_10.png')))
            images2 = sorted(glob(os.path.join(root, 'training', 'image_2/*_11.png')))

            for i in range(len(flows)):
                self.flow_list += [flows[i]]
                self.image_list += [[images1[i], images2[i]]]


    def __getitem__(self, index):

        if self.is_test:
            frame_id = self.image_list[index][0]
            frame_id = frame_id.split('/')[-1]

            img1 = frame_utils.read_gen(self.image_list[index][0])
            img2 = frame_utils.read_gen(self.image_list[index][1])

            img1 = np.array(img1).astype(np.uint8)[..., :3]
            img2 = np.array(img2).astype(np.uint8)[..., :3]

            img1 = torch.from_numpy(img1).permute(2, 0, 1).float()
            img2 = torch.from_numpy(img2).permute(2, 0, 1).float()
            return img1, img2, frame_id


        else:
            if not self.init_seed:
                worker_info = torch.utils.data.get_worker_info()
                if worker_info is not None:
                    np.random.seed(worker_info.id)
                    random.seed(worker_info.id)
                    self.init_seed = True

            index = index % len(self.image_list)
            frame_id = self.image_list[index][0]
            frame_id = frame_id.split('/')[-1]

            img1 = frame_utils.read_gen(self.image_list[index][0])
            img2 = frame_utils.read_gen(self.image_list[index][1])
            flow, valid = frame_utils.readFlowKITTI(self.flow_list[index])

            img1 = np.array(img1).astype(np.uint8)[..., :3]
            img2 = np.array(img2).astype(np.uint8)[..., :3]

            if self.do_augument:
                img1, img2, flow, valid = self.augumentor(img1, img2, flow, valid)

            img1 = torch.from_numpy(img1).permute(2, 0, 1).float()
            img2 = torch.from_numpy(img2).permute(2, 0, 1).float()
            flow = torch.from_numpy(flow).permute(2, 0, 1).float()
            valid = torch.from_numpy(valid).float()

            if self.do_pad:
                ht, wd = img1.shape[1:]
                pad_ht = (((ht // 8) + 1) * 8 - ht) % 8
                pad_wd = (((wd // 8) + 1) * 8 - wd) % 8
                pad_ht1 = [0, pad_ht]
                pad_wd1 = [pad_wd//2, pad_wd - pad_wd//2]
                pad = pad_wd1 + pad_ht1

                img1 = img1.view(1, 3, ht, wd)
                img2 = img2.view(1, 3, ht, wd)
                flow = flow.view(1, 2, ht, wd)
                valid = valid.view(1, 1, ht, wd)

                img1 = torch.nn.functional.pad(img1, pad, mode='replicate')
                img2 = torch.nn.functional.pad(img2, pad, mode='replicate')
                flow = torch.nn.functional.pad(flow, pad, mode='constant', value=0)
                valid = torch.nn.functional.pad(valid, pad, mode='replicate', value=0)

                img1 = img1.view(3, ht+pad_ht, wd+pad_wd)
                img2 = img2.view(3, ht+pad_ht, wd+pad_wd)
                flow = flow.view(2, ht+pad_ht, wd+pad_wd)
                valid = valid.view(ht+pad_ht, wd+pad_wd)

            if self.is_test:
                return img1, img2, flow, valid, frame_id

            return img1, img2, flow, valid