# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, Hao Xiang <haxiang@g.ucla.edu>,
# License: TDG-Attribution-NonCommercial-NoDistrib


"""
Transformation utils
"""

import numpy as np
from scipy.spatial.transform import Rotation as R


# def x_to_world(pose):
#     """
#     The transformation matrix from x-coordinate system to carla world system

#     Parameters
#     ----------
#     pose : list
#         [x, y, z, roll, yaw, pitch]

#     Returns
#     -------
#     matrix : np.ndarray
#         The transformation matrix.
#     """
#     x, y, z, roll, yaw, pitch = pose[:]

#     # used for rotation matrix
#     c_y = np.cos(np.radians(yaw))
#     s_y = np.sin(np.radians(yaw))
#     c_r = np.cos(np.radians(roll))
#     s_r = np.sin(np.radians(roll))
#     c_p = np.cos(np.radians(pitch))
#     s_p = np.sin(np.radians(pitch))

#     matrix = np.identity(4)
#     # translation matrix
#     matrix[0, 3] = x
#     matrix[1, 3] = y
#     matrix[2, 3] = z

#     # rotation matrix
#     matrix[0, 0] = c_p * c_y
#     matrix[0, 1] = c_y * s_p * s_r - s_y * c_r
#     matrix[0, 2] = -c_y * s_p * c_r - s_y * s_r
#     matrix[1, 0] = s_y * c_p
#     matrix[1, 1] = s_y * s_p * s_r + c_y * c_r
#     matrix[1, 2] = -s_y * s_p * c_r + c_y * s_r
#     matrix[2, 0] = s_p
#     matrix[2, 1] = -c_p * s_r
#     matrix[2, 2] = c_p * c_r

#     return matrix

def x_to_world(pose):
    """
    The transformation matrix from x-coordinate system to carla world system

    Parameters
    ----------
    pose : list
        [x, y, z, roll, yaw, pitch]

    Returns
    -------
    matrix : np.ndarray
        The transformation matrix.
    """
    x, y, z, roll, yaw, pitch = pose[:]
    matrix = np.eye(4)
    matrix[:3, :3] = R.from_euler('xzy', [roll, yaw, pitch], degrees=True).as_matrix()
    matrix[:3, 3] = [x, y, z]

    return matrix

def x1_to_x2(x1, x2):
    """
    Transformation matrix from x1 to x2.

    Parameters
    ----------
    x1 : list
        The pose of x1 under world coordinates.
    x2 : list
        The pose of x2 under world coordinates.

    Returns
    -------
    transformation_matrix : np.ndarray
        The transformation matrix.

    """
    x1_to_world = x_to_world(x1)
    x2_to_world = x_to_world(x2)
    world_to_x2 = np.linalg.inv(x2_to_world)

    transformation_matrix = np.dot(world_to_x2, x1_to_world)
    return transformation_matrix


def dist_to_continuous(p_dist, displacement_dist, res, downsample_rate):
    """
    Convert points discretized format to continuous space for BEV representation.
    Parameters
    ----------
    p_dist : numpy.array
        Points in discretized coorindates.

    displacement_dist : numpy.array
        Discretized coordinates of bottom left origin.

    res : float
        Discretization resolution.

    downsample_rate : int
        Dowmsamping rate.

    Returns
    -------
    p_continuous : numpy.array
        Points in continuous coorindates.

    """
    p_dist = np.copy(p_dist)
    p_dist = p_dist + displacement_dist
    p_continuous = p_dist * res * downsample_rate
    return p_continuous
