# Path: ibrnet/data_loaders/data_utils.py
# def deepvoxels_parse_intrinsics(filepath, trgt_sidelength, invert_y=False):
#     # Get camera intrinsics
#     with open(filepath, 'r') as file:
#         f, cx, cy = list(map(float, file.readline().split()))[:3]
#         grid_barycenter = torch.Tensor(list(map(float, file.readline().split())))
#         near_plane = float(file.readline())
#         scale = float(file.readline())
#         height, width = map(float, file.readline().split())
# 
#         try:
#             world2cam_poses = int(file.readline())
#         except ValueError:
#             world2cam_poses = None
# 
#     if world2cam_poses is None:
#         world2cam_poses = False
# 
#     world2cam_poses = bool(world2cam_poses)
# 
#     cx = cx / width * trgt_sidelength
#     cy = cy / height * trgt_sidelength
#     f = trgt_sidelength / height * f
# 
#     fx = f
#     if invert_y:
#         fy = -f
#     else:
#         fy = f
# 
#     # Build the intrinsic matrices
#     full_intrinsic = np.array([[fx, 0., cx, 0.],
#                                [0., fy, cy, 0],
#                                [0., 0, 1, 0],
#                                [0, 0, 0, 1]])
# 
#     return full_intrinsic, grid_barycenter, scale, near_plane, world2cam_poses
# 
# def get_nearest_pose_ids(tar_pose, ref_poses, num_select, tar_id=-1, angular_dist_method='vector',
#                          scene_center=(0, 0, 0)):
#     '''
#     Args:
#         tar_pose: target pose [3, 3]
#         ref_poses: reference poses [N, 3, 3]
#         num_select: the number of nearest views to select
#     Returns: the selected indices
#     '''
#     num_cams = len(ref_poses)
#     num_select = min(num_select, num_cams-1)
#     batched_tar_pose = tar_pose[None, ...].repeat(num_cams, 0)
# 
#     if angular_dist_method == 'matrix':
#         dists = batched_angular_dist_rot_matrix(batched_tar_pose[:, :3, :3], ref_poses[:, :3, :3])
#     elif angular_dist_method == 'vector':
#         tar_cam_locs = batched_tar_pose[:, :3, 3]
#         ref_cam_locs = ref_poses[:, :3, 3]
#         scene_center = np.array(scene_center)[None, ...]
#         tar_vectors = tar_cam_locs - scene_center
#         ref_vectors = ref_cam_locs - scene_center
#         dists = angular_dist_between_2_vectors(tar_vectors, ref_vectors)
#     elif angular_dist_method == 'dist':
#         tar_cam_locs = batched_tar_pose[:, :3, 3]
#         ref_cam_locs = ref_poses[:, :3, 3]
#         dists = np.linalg.norm(tar_cam_locs - ref_cam_locs, axis=1)
#     else:
#         raise Exception('unknown angular distance calculation method!')
# 
#     if tar_id >= 0:
#         assert tar_id < num_cams
#         dists[tar_id] = 1e3  # make sure not to select the target id itself
# 
#     sorted_ids = np.argsort(dists)
#     selected_ids = sorted_ids[:num_select]
#     # print(angular_dists[selected_ids] * 180 / np.pi)
#     return selected_ids
# 
# def rectify_inplane_rotation(src_pose, tar_pose, src_img, th=40):
#     relative = np.linalg.inv(tar_pose).dot(src_pose)
#     relative_rot = relative[:3, :3]
#     r = R.from_matrix(relative_rot)
#     euler = r.as_euler('zxy', degrees=True)
#     euler_z = euler[0]
#     if np.abs(euler_z) < th:
#         return src_pose, src_img
# 
#     R_rectify = R.from_euler('z', -euler_z, degrees=True).as_matrix()
#     src_R_rectified = src_pose[:3, :3].dot(R_rectify)
#     out_pose = np.eye(4)
#     out_pose[:3, :3] = src_R_rectified
#     out_pose[:3, 3:4] = src_pose[:3, 3:4]
#     h, w = src_img.shape[:2]
#     center = ((w - 1.) / 2., (h - 1.) / 2.)
#     M = cv2.getRotationMatrix2D(center, -euler_z, 1)
#     src_img = np.clip((255*src_img).astype(np.uint8), a_max=255, a_min=0)
#     rotated = cv2.warpAffine(src_img, M, (w, h), borderValue=(255, 255, 255), flags=cv2.INTER_LANCZOS4)
#     rotated = rotated.astype(np.float32) / 255.
#     return out_pose, rotated


import os
import numpy as np
import imageio
import torch
import glob
import sys
from torch.utils.data import Dataset
from .data_utils import deepvoxels_parse_intrinsics, get_nearest_pose_ids, rectify_inplane_rotation
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

sys.path.append('../')


class DeepVoxelsDataset(Dataset):
    def __init__(self, args, subset,
                 scenes='vase',  # string or list
                 **kwargs):

        self.folder_path = os.path.join(args.rootdir, 'data/deepvoxels/')

# Path: ibrnet/data_loaders/data_utils.py
# def deepvoxels_parse_intrinsics(filepath, trgt_sidelength, invert_y=False):
#     # Get camera intrinsics
#     with open(filepath, 'r') as file:
#         f, cx, cy = list(map(float, file.readline().split()))[:3]
#         grid_barycenter = torch.Tensor(list(map(float, file.readline().split())))
#         near_plane = float(file.readline())
#         scale = float(file.readline())
#         height, width = map(float, file.readline().split())
# 
#         try:
#             world2cam_poses = int(file.readline())
#         except ValueError:
#             world2cam_poses = None
# 
#     if world2cam_poses is None:
#         world2cam_poses = False
# 
#     world2cam_poses = bool(world2cam_poses)
# 
#     cx = cx / width * trgt_sidelength
#     cy = cy / height * trgt_sidelength
#     f = trgt_sidelength / height * f
# 
#     fx = f
#     if invert_y:
#         fy = -f
#     else:
#         fy = f
# 
#     # Build the intrinsic matrices
#     full_intrinsic = np.array([[fx, 0., cx, 0.],
#                                [0., fy, cy, 0],
#                                [0., 0, 1, 0],
#                                [0, 0, 0, 1]])
# 
#     return full_intrinsic, grid_barycenter, scale, near_plane, world2cam_poses
# 
# def get_nearest_pose_ids(tar_pose, ref_poses, num_select, tar_id=-1, angular_dist_method='vector',
#                          scene_center=(0, 0, 0)):
#     '''
#     Args:
#         tar_pose: target pose [3, 3]
#         ref_poses: reference poses [N, 3, 3]
#         num_select: the number of nearest views to select
#     Returns: the selected indices
#     '''
#     num_cams = len(ref_poses)
#     num_select = min(num_select, num_cams-1)
#     batched_tar_pose = tar_pose[None, ...].repeat(num_cams, 0)
# 
#     if angular_dist_method == 'matrix':
#         dists = batched_angular_dist_rot_matrix(batched_tar_pose[:, :3, :3], ref_poses[:, :3, :3])
#     elif angular_dist_method == 'vector':
#         tar_cam_locs = batched_tar_pose[:, :3, 3]
#         ref_cam_locs = ref_poses[:, :3, 3]
#         scene_center = np.array(scene_center)[None, ...]
#         tar_vectors = tar_cam_locs - scene_center
#         ref_vectors = ref_cam_locs - scene_center
#         dists = angular_dist_between_2_vectors(tar_vectors, ref_vectors)
#     elif angular_dist_method == 'dist':
#         tar_cam_locs = batched_tar_pose[:, :3, 3]
#         ref_cam_locs = ref_poses[:, :3, 3]
#         dists = np.linalg.norm(tar_cam_locs - ref_cam_locs, axis=1)
#     else:
#         raise Exception('unknown angular distance calculation method!')
# 
#     if tar_id >= 0:
#         assert tar_id < num_cams
#         dists[tar_id] = 1e3  # make sure not to select the target id itself
# 
#     sorted_ids = np.argsort(dists)
#     selected_ids = sorted_ids[:num_select]
#     # print(angular_dists[selected_ids] * 180 / np.pi)
#     return selected_ids
# 
# def rectify_inplane_rotation(src_pose, tar_pose, src_img, th=40):
#     relative = np.linalg.inv(tar_pose).dot(src_pose)
#     relative_rot = relative[:3, :3]
#     r = R.from_matrix(relative_rot)
#     euler = r.as_euler('zxy', degrees=True)
#     euler_z = euler[0]
#     if np.abs(euler_z) < th:
#         return src_pose, src_img
# 
#     R_rectify = R.from_euler('z', -euler_z, degrees=True).as_matrix()
#     src_R_rectified = src_pose[:3, :3].dot(R_rectify)
#     out_pose = np.eye(4)
#     out_pose[:3, :3] = src_R_rectified
#     out_pose[:3, 3:4] = src_pose[:3, 3:4]
#     h, w = src_img.shape[:2]
#     center = ((w - 1.) / 2., (h - 1.) / 2.)
#     M = cv2.getRotationMatrix2D(center, -euler_z, 1)
#     src_img = np.clip((255*src_img).astype(np.uint8), a_max=255, a_min=0)
#     rotated = cv2.warpAffine(src_img, M, (w, h), borderValue=(255, 255, 255), flags=cv2.INTER_LANCZOS4)
#     rotated = rotated.astype(np.float32) / 255.
#     return out_pose, rotated


# Path: ibrnet/data_loaders/deepvoxels.py
import os
import numpy as np
import imageio
import torch
import glob
import sys
from torch.utils.data import Dataset
from .data_utils import deepvoxels_parse_intrinsics, get_nearest_pose_ids, rectify_inplane_rotation

# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

sys.path.append('../')


class DeepVoxelsDataset(Dataset):
    def __init__(self, args, subset,
                 scenes='vase',  # string or list
                 **kwargs):

        self.folder_path = os.path.join(args.rootdir, 'data/deepvoxels/')

        self.rectify_inplane_rotation = args.rectify_inplane_rotation