from __future__ import print_function
import sys
sys.path.append('.')
import os
from typing import Optional, Union

import cv2
import numpy as np
import PIL.Image as Image

import pickle
import torch
from torch.utils import data

__all__ = ["TAO"]

class TAO(data.Dataset):
    r"""A torch Dataset for loading in `the TAO VOS dataset <https://www.vision.rwth-aachen.de/page/taovos/>`_. Will fetch sequences of
    rgb images, instance segmentation labels, SuperPoint features (optional).

    Example of sequence creation from frames with `seqlen=4`, `dilation=1`, `stride=3`, and `start=2`:

    .. code-block::


                                            sequence0
                        ┎───────────────┲───────────────┲───────────────┒
                        |               |               |               |
        frame0  frame1  frame2  frame3  frame4  frame5  frame6  frame7  frame8  frame9  frame10  frame11 ...
                                                |               |               |                |
                                                └───────────────┵───────────────┵────────────────┚
                                                                    sequence1


    Args:
        basedir (str): Path to the base directory containing the directories from TAO.
        videos (str or tuple of str): Videos to use from sequences (used for creating train/val/test splits). Can
            be path to a `.txt` file where each line is a Video Seqeunce name, a tuple of scene names.
        seqlen (int): Number of frames to use for each sequence of frames. Default: 4
        dilation (int or None): Number of (original video's) frames to skip between two consecutive
            frames in the extracted sequence. See above example if unsure.
            If None, will set `dilation = 0`. Default: None
        stride (int or None): Number of frames between the first frames of two consecutive extracted sequences.
            See above example if unsure. If None, will set `stride = seqlen * (dilation + 1)`
            (non-overlapping sequences). Default: None
        start (int or None): Index of the frame from which to start extracting sequences for every video.
            If None, will start from the first frame. Default: None
        end (int): Index of the frame at which to stop extracting sequences for every video.
            If None, will continue extracting frames until the end of the video. Default: None
        height (int): Spatial height to resize frames to. Default: 480
        width (int): Spatial width to resize frames to. Default: 640
        return_seg (bool): Determines whether to return instance segmentation labels. Default: True
        return_points (bool): Determines whether to return SuperPoint Features. Default: False
        return_videonames (bool): Determines whether to return videonames for the sequences. Default: False

    """

    def __init__(
        self,
        basedir: str,
        videos: Union[tuple, str, None],
        seqlen: int = 4,
        dilation: Optional[int] = None,
        stride: Optional[int] = None,
        start: Optional[int] = None,
        end: Optional[int] = None,
        height: int = 480,
        width: int = 640,
        *,
        return_img: bool = True,
        return_seg: bool = True,
        return_points: bool = False,
        return_videonames: bool = False,
    ):
        super(TAO, self).__init__()

        self.basedir = os.path.normpath(basedir)
        if not os.path.isdir(self.basedir):
            raise ValueError("Base Directory: {} doesn't exist".format(basedir))

        self.height = height
        self.width = width

        self.return_img = return_img
        self.return_seg = return_seg
        self.return_points = return_points
        self.return_videonames = return_videonames

        if not isinstance(seqlen, int):
            raise TypeError("seqlen must be int. Got {0}.".format(type(seqlen)))
        if not (isinstance(stride, int) or stride is None):
            raise TypeError("stride must be int or None. Got {0}.".format(type(stride)))
        if not (isinstance(dilation, int) or dilation is None):
            raise TypeError(
                "dilation must be int or None. Got {0}.".format(type(dilation))
            )
        dilation = dilation if dilation is not None else 0
        stride = stride if stride is not None else seqlen * (dilation + 1)
        self.seqlen = seqlen
        self.stride = stride
        self.dilation = dilation
        if seqlen < 0:
            raise ValueError("seqlen must be positive. Got {0}.".format(seqlen))
        if dilation < 0:
            raise ValueError('"dilation" must be positive. Got {0}.'.format(dilation))
        if stride < 0:
            raise ValueError("stride must be positive. Got {0}.".format(stride))

        if not (isinstance(start, int) or start is None):
            raise TypeError("start must be int or None. Got {0}.".format(type(start)))
        if not (isinstance(end, int) or end is None):
            raise TypeError("end must be int or None. Got {0}.".format(type(end)))
        start = start if start is not None else 0
        self.start = start
        self.end = end
        if start < 0:
            raise ValueError("start must be positive. Got {0}.".format(stride))
        if not (end is None or end > start):
            raise ValueError(
                "end ({0}) must be None or greater than start ({1})".format(end, start)
            )

        # videos should be a tuple
        if isinstance(videos, str):
            if os.path.isfile(videos):
                with open(videos, "r") as f:
                    videos = tuple(f.read().split("\n"))
            else:
                raise ValueError("incorrect filename: {} doesn't exist".format(videos))
        elif not (isinstance(videos, tuple)):
            msg = "videos should either be path to split.txt or tuple of videos, but was of type %r instead"
            raise TypeError(msg % type(videos))

        self.RGB_data = []
        self.Seg_data = []
        self.Points_data = []
        self.Videonames_data = []

        idx = np.arange(self.seqlen) * (self.dilation + 1)
        
        rgbdir = os.path.join(self.basedir, 'JPEGImages/')
        pointsdir = os.path.join(self.basedir, 'points/')
        segdir = os.path.join(self.basedir, 'Annotations/')

        for video in videos:
            file_names = [f for f in sorted(os.listdir(os.path.join(rgbdir, video))) if f.endswith('.jpg')]
            rgb_list = [os.path.join(os.path.join(rgbdir, video), x) for x in file_names]
            if self.return_points:
                points_list = [os.path.join(os.path.join(pointsdir, video), x.replace('.jpg','.pkl')) for x in file_names]
            if self.return_seg:
                seg_list = [os.path.join(os.path.join(segdir, video), x.replace('.jpg','.png')) for x in file_names]

            video_len = len(rgb_list)

            for start_index in range(self.start, video_len, self.stride):
                if start_index + idx[-1] >= video_len:
                    break
                inds = start_index + idx
                self.RGB_data.append([rgb_list[ind] for ind in inds])

                if self.return_seg:
                    self.Seg_data.append([seg_list[ind] for ind in inds])

                if self.return_points:
                    self.Points_data.append([points_list[ind] for ind in inds])

                if self.return_videonames:
                    self.Videonames_data.append(video)

        self.num_sequences = len(self.RGB_data)

    def __len__(self):
        r"""Returns the length of the dataset. """
        return self.num_sequences

    def __getitem__(self, idx: int):
        r"""Returns the data from the sequence at index idx.

        Returns:
            color_seq (torch.Tensor): Sequence of grayscale rgb images of each frame
            seg_seq (torch.Tensor): Sequence of instance segmentation labels for objects present in the frames
            points_seq (torch.Tensor): Sequence of SuperPoint Features
            videoname (str): Videoname of Sequence

        Shape:
            - color_seq: :math:`(L, 3, H, W)` where `L` denotes sequence length
            - seg_seq: : "math: List of per frame instance segmentations with length `L`
            - points_seq: "math: List of SuperPoint Features with length `L`
        """

        # Read in the color info.
        if self.return_img:
            color_seq_path = self.RGB_data[idx]
        if self.return_seg:
            seg_seq_path = self.Seg_data[idx]
        if self.return_points:
            points_seq_path = self.Points_data[idx]

        color_seq, seg_seq, points_seq = [], [], []

        for i in range(self.seqlen):
          
            if self.return_img:

                image = cv2.imread(color_seq_path[i])
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

                image = torch.from_numpy(image).type(torch.float16)
                image = image.permute(2,0,1)
                image /= 255
                color_seq.append(image)

            if self.return_seg:
                instance_img = np.array(Image.open(seg_seq_path[i]))
                obj_ids = np.unique(instance_img)
                obj_ids = obj_ids[~np.isin(obj_ids, [0])]
                
                frame_ann = []
                for obj_id in obj_ids:
                    ann = {}
                    ann['obj_id'] = obj_id
                    ann_mask = np.isin(instance_img, obj_id).astype(int)
                    ann['ann_mask'] = ann_mask
                    frame_ann.append(ann)
                seg_seq.append(frame_ann)

            if self.return_points:
                with open(points_seq_path[i],'rb') as fp:
                    points = pickle.load(fp)
                points_seq.append(points)

        output = []
        if self.return_img:
            color_seq = torch.stack(color_seq, 0).float()
            output.append(color_seq)

        if self.return_seg:
            output.append(seg_seq)

        if self.return_points:
            output.append(points_seq)

        if self.return_videonames:
            output.append(self.Videonames_data[idx])

        return tuple(output)
[((3, 0, 3, 20), 'sys.path.append', 'sys.path.append', ({(3, 16, 3, 19): '"""."""'}, {}), "('.')", False, 'import sys\n'), ((77, 23, 77, 48), 'os.path.normpath', 'os.path.normpath', ({(77, 40, 77, 47): 'basedir'}, {}), '(basedir)', False, 'import os\n'), ((141, 17, 141, 58), 'os.path.join', 'os.path.join', ({(141, 30, 141, 42): 'self.basedir', (141, 44, 141, 57): '"""JPEGImages/"""'}, {}), "(self.basedir, 'JPEGImages/')", False, 'import os\n'), ((142, 20, 142, 57), 'os.path.join', 'os.path.join', ({(142, 33, 142, 45): 'self.basedir', (142, 47, 142, 56): '"""points/"""'}, {}), "(self.basedir, 'points/')", False, 'import os\n'), ((143, 17, 143, 59), 'os.path.join', 'os.path.join', ({(143, 30, 143, 42): 'self.basedir', (143, 44, 143, 58): '"""Annotations/"""'}, {}), "(self.basedir, 'Annotations/')", False, 'import os\n'), ((78, 15, 78, 42), 'os.path.isdir', 'os.path.isdir', ({(78, 29, 78, 41): 'self.basedir'}, {}), '(self.basedir)', False, 'import os\n'), ((125, 15, 125, 37), 'os.path.isfile', 'os.path.isfile', ({(125, 30, 125, 36): 'videos'}, {}), '(videos)', False, 'import os\n'), ((139, 14, 139, 36), 'numpy.arange', 'np.arange', ({(139, 24, 139, 35): 'self.seqlen'}, {}), '(self.seqlen)', True, 'import numpy as np\n'), ((205, 24, 205, 53), 'cv2.imread', 'cv2.imread', ({(205, 35, 205, 52): 'color_seq_path[i]'}, {}), '(color_seq_path[i])', False, 'import cv2\n'), ((206, 24, 206, 62), 'cv2.cvtColor', 'cv2.cvtColor', ({(206, 37, 206, 42): 'image', (206, 44, 206, 61): 'cv2.COLOR_BGR2RGB'}, {}), '(image, cv2.COLOR_BGR2RGB)', False, 'import cv2\n'), ((215, 26, 215, 49), 'numpy.unique', 'np.unique', ({(215, 36, 215, 48): 'instance_img'}, {}), '(instance_img)', True, 'import numpy as np\n'), ((147, 37, 147, 64), 'os.path.join', 'os.path.join', ({(147, 50, 147, 56): 'rgbdir', (147, 58, 147, 63): 'video'}, {}), '(rgbdir, video)', False, 'import os\n'), ((214, 40, 214, 67), 'PIL.Image.open', 'Image.open', ({(214, 51, 214, 66): 'seg_seq_path[i]'}, {}), '(seg_seq_path[i])', True, 'import PIL.Image as Image\n'), ((229, 29, 229, 44), 'pickle.load', 'pickle.load', ({(229, 41, 229, 43): 'fp'}, {}), '(fp)', False, 'import pickle\n'), ((234, 24, 234, 49), 'torch.stack', 'torch.stack', ({(234, 36, 234, 45): 'color_seq', (234, 47, 234, 48): '0'}, {}), '(color_seq, 0)', False, 'import torch\n'), ((149, 44, 149, 74), 'os.path.join', 'os.path.join', ({(149, 57, 149, 66): 'pointsdir', (149, 68, 149, 73): 'video'}, {}), '(pointsdir, video)', False, 'import os\n'), ((151, 41, 151, 68), 'os.path.join', 'os.path.join', ({(151, 54, 151, 60): 'segdir', (151, 62, 151, 67): 'video'}, {}), '(segdir, video)', False, 'import os\n'), ((208, 24, 208, 47), 'torch.from_numpy', 'torch.from_numpy', ({(208, 41, 208, 46): 'image'}, {}), '(image)', False, 'import torch\n'), ((146, 55, 146, 82), 'os.path.join', 'os.path.join', ({(146, 68, 146, 74): 'rgbdir', (146, 76, 146, 81): 'video'}, {}), '(rgbdir, video)', False, 'import os\n'), ((216, 35, 216, 56), 'numpy.isin', 'np.isin', ({(216, 43, 216, 50): 'obj_ids', (216, 52, 216, 55): '[0]'}, {}), '(obj_ids, [0])', True, 'import numpy as np\n'), ((222, 31, 222, 60), 'numpy.isin', 'np.isin', ({(222, 39, 222, 51): 'instance_img', (222, 53, 222, 59): 'obj_id'}, {}), '(instance_img, obj_id)', True, 'import numpy as np\n')]