#!/usr/bin/env python3
# Developed by Xieyuanli Chen and Thomas Läbe
# This file is covered by the LICENSE file in the root of this project.
# Brief: This script generate the overlap and orientation combined mapping file.

from tools.utils.utils import *
import matplotlib.pyplot as plt

velodatatype = np.dtype({
    'x': ('<u2', 0),
    'y': ('<u2', 2),
    'z': ('<u2', 4),
    'i': ('u1', 6),
    'l': ('u1', 7)})
velodatasize = 8



def data2xyzi(data, flip=True):
    xyzil = data.view(velodatatype)
    xyz = np.hstack(
        [xyzil[axis].reshape([-1, 1]) for axis in ['x', 'y', 'z']])
    xyz = xyz * 0.005 - 100.0

    if flip:
        R = np.eye(3)
        R[2, 2] = -1
        # xyz = xyz @ R
        xyz = np.matmul(xyz, R)

    return xyz, xyzil['i']

def get_velo(velofile):
    return data2xyzi(np.fromfile(velofile))



def com_overlap_yaw(scan_paths, poses, frame_idx, leg_output_width=360, viz=False):
  """compute the overlap and yaw ground truth from the ground truth poses,
     which is used for OverlapNet training and testing.
     Args:
       scan_paths: paths of all raw LiDAR scans
       poses: ground-truth poses either given by the dataset or generated by SLAM or odometry
       frame_idx: the current frame index
     Returns:
       ground_truth_mapping: the ground truth overlap and yaw used for training OverlapNet,
                             where each row contains [current_frame_idx, reference_frame_idx, overlap, yaw]
  """
  # init ground truth overlap and yaw
  print('Start to compute ground truth overlap and yaw ...')
  overlaps = []
  yaw_idxs = []
  yaw_resolution = leg_output_width
  
  # we calculate the ground truth for one given frame only
  # generate range projection for the given frame
  current_points = get_velo(scan_paths[frame_idx])[0]
  fov_up = 30.67
  fov_down = -10.67
  proj_H = 32
  proj_W = 900
  lowest = 0.1
  highest = 6
  current_range, project_points, _ = range_projection(current_points,
                                                fov_up=fov_up,
                                                fov_down=fov_down,
                                                proj_H=proj_H,
                                                proj_W=proj_W,
                                                max_range=80,
                                                cut_z=False,
                                                low=lowest,
                                                high=highest)
  visible_points = project_points[current_range > 0]
  valid_num = len(visible_points)
  current_pose = poses[frame_idx]

  tau_ = 1.2

  for reference_idx in range(len(scan_paths)):
    # generate range projection for the reference frame
    reference_pose = poses[reference_idx]
    reference_points_tmp = get_velo(scan_paths[reference_idx])[0]

    reference_points = np.ones((reference_points_tmp.shape[0], reference_points_tmp.shape[1] + 1))
    reference_points[:, :-1] = reference_points_tmp

    reference_points_world = reference_pose.dot(reference_points.T).T
    reference_points_in_current = np.linalg.inv(current_pose).dot(reference_points_world.T).T
    reference_range, _, _ = range_projection(reference_points_in_current,
                                                        fov_up=fov_up,
                                                        fov_down=fov_down,
                                                        proj_H=proj_H,
                                                        proj_W=proj_W,
                                                        max_range=80,
                                                        cut_z=False,
                                                        low=lowest,
                                                        high=highest)

    dist = np.linalg.norm(reference_pose[:3, -1].reshape(3,) - current_pose[:3, -1].reshape(3,))

    # calculate overlap
    overlap = np.count_nonzero(
      abs(reference_range[reference_range > 0] - current_range[reference_range > 0]) < tau_) / valid_num
    overlaps.append(overlap)

    
    # calculate yaw angle
    relative_transform = np.linalg.inv(current_pose).dot(reference_pose)
    relative_rotation = relative_transform[:3, :3]
    _, _, yaw = euler_angles_from_rotation_matrix(relative_rotation)

    # discretize yaw angle and shift the 0 degree to the center to make the network easy to lean
    yaw_element_idx = int(- (yaw / np.pi) * yaw_resolution//2 + yaw_resolution//2)
    yaw_idxs.append(yaw_element_idx)

  # ground truth format: each row contains [current_frame_idx, reference_frame_idx, overlap, yaw]
  ground_truth_mapping = np.zeros((len(scan_paths), 4))
  ground_truth_mapping[:, 0] = np.ones(len(scan_paths)) * frame_idx
  ground_truth_mapping[:, 1] = np.arange(len(scan_paths))
  ground_truth_mapping[:, 2] = overlaps
  ground_truth_mapping[:, 3] = yaw_idxs
  
  print('Finish generating ground_truth_mapping!')
  
  return ground_truth_mapping


def com_overlap_yaw_idx(scan_paths, poses, frame_idx, leg_output_width=360, viz=False, idx_in_range=None):

    print('Start to compute ground truth overlap and yaw ...')
    overlaps = []
    yaw_idxs = []
    yaw_resolution = leg_output_width

    # we calculate the ground truth for one given frame only
    # generate range projection for the given frame
    current_points = get_velo(scan_paths[frame_idx])[0]
    fov_up = 30.67
    fov_down = -10.67
    proj_H = 32
    proj_W = 900
    lowest = 0.1
    highest = 6
    current_range, project_points, _ = range_projection(current_points,
                                                        fov_up=fov_up,
                                                        fov_down=fov_down,
                                                        proj_H=proj_H,
                                                        proj_W=proj_W,
                                                        max_range=80,
                                                        cut_z=False,
                                                        low=lowest,
                                                        high=highest)

    plt.imshow(current_range)
    plt.show()


    visible_points = project_points[current_range > 0]
    valid_num = len(visible_points)
    current_pose = poses[frame_idx]

    tau_ = 1.2
    print("tau: ", tau_)

    for reference_idx in range(len(scan_paths)):
        # generate range projection for the reference frame
        reference_pose = poses[reference_idx]
        reference_points_tmp = get_velo(scan_paths[reference_idx])[0]

        reference_points = np.ones((reference_points_tmp.shape[0], reference_points_tmp.shape[1] + 1))
        reference_points[:, :-1] = reference_points_tmp

        reference_points_world = reference_pose.dot(reference_points.T).T
        reference_points_in_current = np.linalg.inv(current_pose).dot(reference_points_world.T).T
        reference_range, _, _ = range_projection(reference_points_in_current,
                                                 fov_up=fov_up,
                                                 fov_down=fov_down,
                                                 proj_H=proj_H,
                                                 proj_W=proj_W,
                                                 max_range=80,
                                                 cut_z=False,
                                                 low=lowest,
                                                 high=highest)
        # if reference_idx == 0:
        dist = np.linalg.norm(reference_pose[:3, -1].reshape(3, ) - current_pose[:3, -1].reshape(3, ))

        # calculate overlap
        overlap = np.count_nonzero(
            abs(reference_range[reference_range > 0] - current_range[reference_range > 0]) < tau_) / valid_num
        # abs(reference_range[reference_range > 0] - current_range[reference_range > 0]) < 2) / valid_num
        overlaps.append(overlap)

        # calculate yaw angle
        relative_transform = np.linalg.inv(current_pose).dot(reference_pose)
        relative_rotation = relative_transform[:3, :3]
        _, _, yaw = euler_angles_from_rotation_matrix(relative_rotation)

        # discretize yaw angle and shift the 0 degree to the center to make the network easy to lean
        yaw_element_idx = int(- (yaw / np.pi) * yaw_resolution // 2 + yaw_resolution // 2)
        yaw_idxs.append(yaw_element_idx)

        # print('finished pair id: ', reference_idx)

    # ground truth format: each row contains [current_frame_idx, reference_frame_idx, overlap, yaw]
    ground_truth_mapping = np.zeros((len(scan_paths), 4))
    ground_truth_mapping[:, 0] = np.ones(len(scan_paths)) * frame_idx
    ground_truth_mapping[:, 1] = np.arange(len(scan_paths))
    ground_truth_mapping[:, 2] = overlaps
    ground_truth_mapping[:, 3] = yaw_idxs

    print('Finish generating ground_truth_mapping!')

    return ground_truth_mapping
