import csv
import os
from pathlib import Path
from xml.etree import ElementTree as ET

from tqdm import tqdm
from scipy.spatial.transform import Rotation
from skimage.io import imread
import numpy as np

from tadataka.camera import CameraModel, CameraParameters, FOV
from tadataka.dataset.frame import Frame
from tadataka.dataset.base import BaseDataset
from tadataka.pose import Pose


def load_depth(path):
    tree = ET.parse(path)
    root = tree.getroot()
    rows_node, cols_node, dt_node, data_node = root[0]
    height, width = int(rows_node.text), int(cols_node.text)

    depth_text = data_node.text
    depth_text = depth_text.replace('\n', '').strip()
    depth_map = np.fromstring(depth_text, sep=' ')
    return depth_map.reshape(height, width)


def generate_cache(src_dir, cache_dir, src_extension, loader):
    def generate_(subdir):
        os.makedirs(str(Path(cache_dir, subdir)))

        print(f"Generating cache from {subdir}")

        paths = Path(src_dir, subdir).glob("*" + src_extension)
        for path in tqdm(list(paths)):
            filename = path.name.replace(src_extension, ".npy")
            cache_path = Path(cache_dir, subdir, filename)
            array = loader(path)
            np.save(str(cache_path), array)

    generate_("left")
    generate_("right")


def generate_image_cache(image_dir, cache_dir):
    print("Generating image cache")
    generate_cache(image_dir, cache_dir, ".png", imread)


def generate_depth_cache(depth_dir, cache_dir):
    print("Generating depth cache")
    generate_cache(depth_dir, cache_dir, ".xml", load_depth)


def align_coordinate_system(positions, euler_angles):
    # Camera coordinate system and world coordinate system are not aligned
    #
    # Usually camera coordinate system is represented in the format that
    # x: right  y: down  z: forward
    # however, in 'camera_track.txt', they are written in
    # x: right  y: up    z: backward
    #
    # This means the camera coordinate system is
    # rotated 180 degrees around the x-axis from the world coordinate system

    # rotate 180 degrees around the x-axis
    R = Rotation.from_rotvec([np.pi, 0, 0]).as_matrix()
    positions = np.dot(R, positions.T).T

    # Reverse rotations around y and z because axes are flipped
    # (rot_x, rot_y, rot_z) <- (rot_x, -rot_y, -rot_z)
    euler_angles[:, 1:3] = -euler_angles[:, 1:3]
    return positions, euler_angles


def load_poses(pose_path):
    poses = np.loadtxt(pose_path, delimiter=',')
    positions, euler_angles = poses[:, 0:3], poses[:, 3:6]
    positions, euler_angles = align_coordinate_system(positions, euler_angles)
    rotations = Rotation.from_euler('xyz', euler_angles, degrees=True)
    return rotations, positions


def discard_alpha(image):
    return image[:, :, 0:3]


def calc_baseline_offset(rotation, baseline_length):
    local_offset = np.array([baseline_length, 0, 0])
    R = rotation.as_matrix()
    return np.dot(R, local_offset)


# TODO download and set dataset_root automatically
class NewTsukubaDataset(BaseDataset):
    def __init__(self, dataset_root, condition="daylight"):
        self.camera_model = CameraModel(
            CameraParameters(focal_length=[615, 615], offset=[320, 240]),
            distortion_model=None
        )
        groundtruth_dir = Path(dataset_root, "groundtruth")
        illumination_dir = Path(dataset_root, "illumination")

        pose_path = Path(groundtruth_dir, "camera_track.txt")

        self.baseline_length = 10.0
        self.rotations, self.positions = load_poses(pose_path)

        depth_dir = Path(groundtruth_dir, "depth_maps")
        depth_cache_dir = Path(groundtruth_dir, "depth_cache")

        if not depth_cache_dir.exists():
            generate_depth_cache(depth_dir, depth_cache_dir)

        self.depth_L_paths = sorted(Path(depth_cache_dir, "left").glob("*.npy"))
        self.depth_R_paths = sorted(Path(depth_cache_dir, "right").glob("*.npy"))

        image_dir = Path(illumination_dir, condition)
        image_cache_dir = Path(illumination_dir, condition + "_cache")

        if not image_cache_dir.exists():
            generate_image_cache(image_dir, image_cache_dir)

        self.image_L_paths = sorted(Path(image_cache_dir, "left").glob("*.npy"))
        self.image_R_paths = sorted(Path(image_cache_dir, "right").glob("*.npy"))

        assert((len(self.depth_L_paths) == len(self.depth_R_paths) ==
                len(self.image_L_paths) == len(self.image_R_paths) ==
                len(self.rotations) == len(self.positions)))

        for i in range(len(self.positions)):
            DL = self.depth_L_paths[i].name
            DR = self.depth_R_paths[i].name
            IL = self.image_L_paths[i].name
            IR = self.image_R_paths[i].name

            assert(DL[-8:] == DR[-8:] == IL[-8:] == IR[-8:])

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

    def load(self, index):
        image_l = np.load(self.image_L_paths[index])
        image_r = np.load(self.image_R_paths[index])

        image_l = discard_alpha(image_l)
        image_r = discard_alpha(image_r)

        depth_l = np.load(self.depth_L_paths[index])
        depth_r = np.load(self.depth_R_paths[index])

        position_center = self.positions[index]
        rotation = self.rotations[index]

        offset = calc_baseline_offset(rotation, self.baseline_length)
        pose_wl = Pose(rotation, position_center - offset / 2.0)
        pose_wr = Pose(rotation, position_center + offset / 2.0)
        return (
            Frame(self.camera_model, pose_wl, image_l, depth_l),
            Frame(self.camera_model, pose_wr, image_r, depth_r)
        )

[((18, 11, 18, 25), 'xml.etree.ElementTree.parse', 'ET.parse', ({(18, 20, 18, 24): 'path'}, {}), '(path)', True, 'from xml.etree import ElementTree as ET\n'), ((25, 16, 25, 50), 'numpy.fromstring', 'np.fromstring', (), '', True, 'import numpy as np\n'), ((78, 12, 78, 48), 'numpy.loadtxt', 'np.loadtxt', (), '', True, 'import numpy as np\n'), ((81, 16, 81, 70), 'scipy.spatial.transform.Rotation.from_euler', 'Rotation.from_euler', (), '', False, 'from scipy.spatial.transform import Rotation\n'), ((90, 19, 90, 52), 'numpy.array', 'np.array', ({(90, 28, 90, 51): '[baseline_length, 0, 0]'}, {}), '([baseline_length, 0, 0])', True, 'import numpy as np\n'), ((92, 11, 92, 34), 'numpy.dot', 'np.dot', ({(92, 18, 92, 19): 'R', (92, 21, 92, 33): 'local_offset'}, {}), '(R, local_offset)', True, 'import numpy as np\n'), ((69, 16, 69, 38), 'numpy.dot', 'np.dot', ({(69, 23, 69, 24): 'R', (69, 26, 69, 37): 'positions.T'}, {}), '(R, positions.T)', True, 'import numpy as np\n'), ((102, 26, 102, 59), 'pathlib.Path', 'Path', ({(102, 31, 102, 43): 'dataset_root', (102, 45, 102, 58): '"""groundtruth"""'}, {}), "(dataset_root, 'groundtruth')", False, 'from pathlib import Path\n'), ((103, 27, 103, 61), 'pathlib.Path', 'Path', ({(103, 32, 103, 44): 'dataset_root', (103, 46, 103, 60): '"""illumination"""'}, {}), "(dataset_root, 'illumination')", False, 'from pathlib import Path\n'), ((105, 20, 105, 61), 'pathlib.Path', 'Path', ({(105, 25, 105, 40): 'groundtruth_dir', (105, 42, 105, 60): '"""camera_track.txt"""'}, {}), "(groundtruth_dir, 'camera_track.txt')", False, 'from pathlib import Path\n'), ((110, 20, 110, 55), 'pathlib.Path', 'Path', ({(110, 25, 110, 40): 'groundtruth_dir', (110, 42, 110, 54): '"""depth_maps"""'}, {}), "(groundtruth_dir, 'depth_maps')", False, 'from pathlib import Path\n'), ((111, 26, 111, 62), 'pathlib.Path', 'Path', ({(111, 31, 111, 46): 'groundtruth_dir', (111, 48, 111, 61): '"""depth_cache"""'}, {}), "(groundtruth_dir, 'depth_cache')", False, 'from pathlib import Path\n'), ((119, 20, 119, 53), 'pathlib.Path', 'Path', ({(119, 25, 119, 41): 'illumination_dir', (119, 43, 119, 52): 'condition'}, {}), '(illumination_dir, condition)', False, 'from pathlib import Path\n'), ((120, 26, 120, 70), 'pathlib.Path', 'Path', ({(120, 31, 120, 47): 'illumination_dir', (120, 49, 120, 69): "condition + '_cache'"}, {}), "(illumination_dir, condition + '_cache')", False, 'from pathlib import Path\n'), ((144, 18, 144, 52), 'numpy.load', 'np.load', ({(144, 26, 144, 51): 'self.image_L_paths[index]'}, {}), '(self.image_L_paths[index])', True, 'import numpy as np\n'), ((145, 18, 145, 52), 'numpy.load', 'np.load', ({(145, 26, 145, 51): 'self.image_R_paths[index]'}, {}), '(self.image_R_paths[index])', True, 'import numpy as np\n'), ((150, 18, 150, 52), 'numpy.load', 'np.load', ({(150, 26, 150, 51): 'self.depth_L_paths[index]'}, {}), '(self.depth_L_paths[index])', True, 'import numpy as np\n'), ((151, 18, 151, 52), 'numpy.load', 'np.load', ({(151, 26, 151, 51): 'self.depth_R_paths[index]'}, {}), '(self.depth_R_paths[index])', True, 'import numpy as np\n'), ((157, 18, 157, 64), 'tadataka.pose.Pose', 'Pose', ({(157, 23, 157, 31): 'rotation', (157, 33, 157, 63): 'position_center - offset / 2.0'}, {}), '(rotation, position_center - offset / 2.0)', False, 'from tadataka.pose import Pose\n'), ((158, 18, 158, 64), 'tadataka.pose.Pose', 'Pose', ({(158, 23, 158, 31): 'rotation', (158, 33, 158, 63): 'position_center + offset / 2.0'}, {}), '(rotation, position_center + offset / 2.0)', False, 'from tadataka.pose import Pose\n'), ((38, 25, 38, 58), 'pathlib.Path', 'Path', ({(38, 30, 38, 39): 'cache_dir', (38, 41, 38, 47): 'subdir', (38, 49, 38, 57): 'filename'}, {}), '(cache_dir, subdir, filename)', False, 'from pathlib import Path\n'), ((68, 8, 68, 43), 'scipy.spatial.transform.Rotation.from_rotvec', 'Rotation.from_rotvec', ({(68, 29, 68, 42): '[np.pi, 0, 0]'}, {}), '([np.pi, 0, 0])', False, 'from scipy.spatial.transform import Rotation\n'), ((99, 12, 99, 72), 'tadataka.camera.CameraParameters', 'CameraParameters', (), '', False, 'from tadataka.camera import CameraModel, CameraParameters, FOV\n'), ((160, 12, 160, 63), 'tadataka.dataset.frame.Frame', 'Frame', ({(160, 18, 160, 35): 'self.camera_model', (160, 37, 160, 44): 'pose_wl', (160, 46, 160, 53): 'image_l', (160, 55, 160, 62): 'depth_l'}, {}), '(self.camera_model, pose_wl, image_l, depth_l)', False, 'from tadataka.dataset.frame import Frame\n'), ((161, 12, 161, 63), 'tadataka.dataset.frame.Frame', 'Frame', ({(161, 18, 161, 35): 'self.camera_model', (161, 37, 161, 44): 'pose_wr', (161, 46, 161, 53): 'image_r', (161, 55, 161, 62): 'depth_r'}, {}), '(self.camera_model, pose_wr, image_r, depth_r)', False, 'from tadataka.dataset.frame import Frame\n'), ((31, 24, 31, 47), 'pathlib.Path', 'Path', ({(31, 29, 31, 38): 'cache_dir', (31, 40, 31, 46): 'subdir'}, {}), '(cache_dir, subdir)', False, 'from pathlib import Path\n'), ((35, 16, 35, 37), 'pathlib.Path', 'Path', ({(35, 21, 35, 28): 'src_dir', (35, 30, 35, 36): 'subdir'}, {}), '(src_dir, subdir)', False, 'from pathlib import Path\n'), ((116, 36, 116, 65), 'pathlib.Path', 'Path', ({(116, 41, 116, 56): 'depth_cache_dir', (116, 58, 116, 64): '"""left"""'}, {}), "(depth_cache_dir, 'left')", False, 'from pathlib import Path\n'), ((117, 36, 117, 66), 'pathlib.Path', 'Path', ({(117, 41, 117, 56): 'depth_cache_dir', (117, 58, 117, 65): '"""right"""'}, {}), "(depth_cache_dir, 'right')", False, 'from pathlib import Path\n'), ((125, 36, 125, 65), 'pathlib.Path', 'Path', ({(125, 41, 125, 56): 'image_cache_dir', (125, 58, 125, 64): '"""left"""'}, {}), "(image_cache_dir, 'left')", False, 'from pathlib import Path\n'), ((126, 36, 126, 66), 'pathlib.Path', 'Path', ({(126, 41, 126, 56): 'image_cache_dir', (126, 58, 126, 65): '"""right"""'}, {}), "(image_cache_dir, 'right')", False, 'from pathlib import Path\n')]