from __future__ import absolute_import, division, print_function

import os
import skimage.transform
import numpy as np
import PIL.Image as pil

from kitti_utils import generate_depth_map
from .mono_dataset import MonoDataset


class EsophagusDataset(MonoDataset):
    def __init__(self, *args, **kwargs):
        super(EsophagusDataset, self).__init__(*args, **kwargs)

        self.K = np.array([[229.909354, 0, 232.628838, 0],
                           [0, 223.048278, 177.914264, 0],
                           [0, 0, 1, 0],
                           [0, 0, 0, 1]], dtype=np.float32)

        self.full_res_shape = (480, 352)

    def check_depth(self):
        line = self.filenames[0].split()
        scene_name = line[0]
        frame_index = int(line[1])

        velo_filename = os.path.join(
            self.data_path,
            scene_name,
            "velodyne_points/data/{:010d}.bin".format(int(frame_index)))

        return os.path.isfile(velo_filename)


    def get_color(self, folder, frame_index, seq_num, do_flip):
        color = self.loader(self.get_image_path(folder, frame_index, seq_num))

        if do_flip:
            color = color.transpose(pil.FLIP_LEFT_RIGHT)

        return color


class EsophagusRAWDataset(EsophagusDataset):

    def __init__(self, *args, **kwargs):
        super(EsophagusRAWDataset, self).__init__(*args, **kwargs)

    def get_image_path(self, folder, frame_index, seq_num):
        f_str = "{:08d}{}".format(frame_index, self.img_ext)
        image_path = os.path.join(self.data_path, folder, "seq_{}".format(seq_num), f_str)
        return image_path

    # ここは適当
    def get_depth(self, folder, frame_index, side, do_flip):
        calib_path = os.path.join(self.data_path, folder.split("/")[0])

        velo_filename = os.path.join(
            self.data_path,
            folder,
            "velodyne_points/data/{:010d}.bin".format(int(frame_index)))

        depth_gt = generate_depth_map(calib_path, velo_filename, self.side_map[side])
        depth_gt = skimage.transform.resize(
            depth_gt, self.full_res_shape[::-1], order=0, preserve_range=True, mode='constant')

        if do_flip:
            depth_gt = np.fliplr(depth_gt)

        return depth_gt