#!/usr/bin/env python
"""
Library handling projection operations of a standard camera model.

Note: the image coordiante system has its origin at the center of the top left
pixel.

"""
from __future__ import division, print_function
import numpy as np
import matplotlib.pyplot as plt

# ROS imports
from tf.transformations import quaternion_multiply, quaternion_matrix, \
    quaternion_from_euler, quaternion_inverse, quaternion_from_matrix, \
    euler_from_matrix


def generate_camera_quaternion(pan, tilt, ins_quaternion=None):
    """Tool to generate a reasonable camera quaternion.

    When pan=tilt=0, the camera is level and points along world y.

    :param pan: Pan in degrees

    :parma tilt: Tilt in degrees

    :return: ROS quaternion
    :rtype: 4-array
    """
    if ins_quaternion is None:
        ins_quaternion = np.array([1/np.sqrt(2),1/np.sqrt(2),0,0])

    pan = pan/180*np.pi
    tilt = tilt/180*np.pi

    # ROS quaternion representing level camera pointing along y in the world.
    quat_n2b = [-np.sqrt(2)/2,0,0,np.sqrt(2)/2]
    quat_b2c = quaternion_from_euler(0, pan, tilt, axes='rzyx')

    # The ROS convention for the rotation is that it transforms the
    # coordinate system. The computer vision community convention is that
    # the operator transforms a vector to get it into the moving coordinate
    # system. So, we invert each quaternion.
    world_quaternion = quaternion_multiply(quaternion_inverse(quat_b2c),
                                           quaternion_inverse(quat_n2b))

    R = quaternion_matrix(world_quaternion)[:3,:3]
    print('Camera axes in world coordinates:')
    print('Camera x-axis:', repr(tuple(R[0])))
    print('Camera y-axis:', repr(tuple(R[1])))
    print('Camera z-axis:', repr(tuple(R[2])),'\n')

    camera_quaternion = quaternion_multiply(world_quaternion, ins_quaternion)

    R = quaternion_matrix(camera_quaternion)[:3,:3]
    print('Camera axes in INS coordinates:')
    print('Camera x-axis:', repr(tuple(R[0])))
    print('Camera y-axis:', repr(tuple(R[1])))
    print('Camera z-axis:', repr(tuple(R[2])))

    camera_quaternion = quaternion_inverse(camera_quaternion)
    camera_quaternion /= camera_quaternion[-1]
    camera_quaternion /= np.linalg.norm(camera_quaternion)
    return  camera_quaternion


def generate_base_quat(x, y, z):
    """Tool to generate a reasonable PTZ camera base quaternion.

    :param x: Vector along the direction of the PTZ camera base coordinate
        x-axis system in INS coordinates.
    :type x: 3-array

    :param y: Vector along the direction of the PTZ camera base coordinate
        y-axis system in INS coordinates.
    :type y: 3-array

    :param z: Vector along the direction of the PTZ camera base coordinate
        z-axis system in INS coordinates.
    :type z: 3-array

    :return: ROS quaternion
    :rtype: 4-array

    """
    R = np.identity(4)
    x = np.array(x)
    y = np.array(y)
    z = np.array(z)
    R[:3,0] = x/np.linalg.norm(x)
    R[:3,1] = y/np.linalg.norm(y)
    R[:3,2] = z/np.linalg.norm(z)
    base_quaternion = quaternion_from_matrix(R)
    return  base_quaternion


def analyze_base_quat(base_quaternion):
    """Analyze the PTZ camera's base quaternion.

    :param x: Vector along the direction of the PTZ camera base coordinate
        x-axis system in INS coordinates.
    :type x: 3-array

    :param y: Vector along the direction of the PTZ camera base coordinate
        y-axis system in INS coordinates.
    :type y: 3-array

    :param z: Vector along the direction of the PTZ camera base coordinate
        z-axis system in INS coordinates.
    :type z: 3-array

    :return: ROS quaternion
    :rtype: 4-array

    """
    R = quaternion_matrix(base_quaternion)[:3,:3].T
    print('PTZ camera base coordinate systems in INS coordinates:')
    print('Camera x-axis:', repr(tuple(R[0])))
    print('Camera y-axis:', repr(tuple(R[1])))
    print('Camera z-axis:', repr(tuple(R[2])))


def rotation_between_quats(quat1, quat2):
    """Return the axis angle representation of the rotation from quat1 to quat2.
    """
    # quat2 = dq*quat1
    dq = quaternion_multiply(quat2, quaternion_inverse(quat1))
    axis = dq[:3]
    axis = axis/np.linalg.norm(axis)
    angle = np.arccos(np.clip(dq[3], -1, 1))*2
    return axis, angle


def get_ellipsoid_from_cov(cov):
    """Return 1-sigma ellipsoid axes from a covariance matrix.

    Arguements:
    cov : 2x2 or 3x3

    Returns:
    ellipsoid_axes : 2x2 or 3x3
        The columns are vectors along the primary axes of the ellipsoid.  The
        length of each vector is the 1-sigma point (meters).

    """
    variance, ellipsoid_axes = np.linalg.eig(cov)

    for i in range(len(variance)):
        ellipsoid_axes[:,i] *= np.sqrt(variance[i])

    return ellipsoid_axes


def get_2d_ellipsoid_from_3d_cov(cov):
    """Return 1-sigma 2-D ellipse axes from a 3-D ellipsoid covariance matrix.

    Arguements:
    cov : 3x3

    Returns:
    ellipsoid_axes : 2x2
        The columns are vectors along the primary axes of the ellipsoid.  The
        length of each vector is the 1-sigma point (meters).

    """
    variance, ellipsoid_axes = np.linalg.eig(cov)

    # Columns of axes are vectors along the primary axes of the ellipsoid
    # with length equal to the 1-sigma uncertainty (meters). This is a 3-D
    # ellipsoid, but we want a 2-D ellipsoid as seen from above. So, we
    # want to project the ellipse onto the <1,0,0> / <0,1,0> plane.

    # In lieu of something much more rigorous, we will leverage the fact that
    # the camera coordinate system is very close to upright. So, one principal
    # axis will be very close to vertical. Therefore, we simply need to
    # identify that axis and select the other two, which will be in the x/y
    # plane.

    # This is the index for the axis that corresponds to vertical.
    ind = np.argmax(np.abs(ellipsoid_axes[2]))

    # We no longer care about the z-component of the axes
    ellipsoid_axes = ellipsoid_axes[:2]

    # Remove the vertical axis
    ellipsoid_axes = np.delete(ellipsoid_axes, ind, 1)
    variance = np.delete(variance, ind)

    for i in range(len(variance)):
        ellipsoid_axes[:,i] *= np.sqrt(variance[i])

    # np.linalg.eig returns axes in increasing order of variance, so let's
    # flip that.
    ellipsoid_axes = np.fliplr(ellipsoid_axes)

    return ellipsoid_axes


def get_ellipse2d_pts(center, axes, N, plot_results=False):
    """Return a closed curve for the perimeter of a 2-d ellipse.

    Note: I'm 90% sure this produces a true ellipse perimeter.  If not, it
    produces something very close in appearance.

    Arguements:
    center : 2-array
        Coordinates for the center of the ellipse.
    axes : 2x2 array
        Column vectors representing the orthogonal axes of the ellipse, such as
        the output from get_ellipsoid_from_cov.
    N : int
        Number of points along the perimeter.  Note that the last point will be
        identical to the first so that the plot will produce a closed curve.

    :return: Points on the ellipse.
    :rtype: 2xN array

    """
    theta = np.linspace(0, 2*np.pi, N)
    pts = np.zeros((2,N))
    pts[0] = np.cos(theta)*axes[0,0] + np.sin(theta)*axes[0,1] + center[0]
    pts[1] = np.cos(theta)*axes[1,0] + np.sin(theta)*axes[1,1] + center[1]

    if plot_results:
        plt.figure()
        plt.plot(pts[0], pts[1], 'r-')
        #plt.plot([center[0],center[0]+axes[0,0]], [center[1],center[1]+axes[1,0]], 'k-')
        #plt.plot([center[0],center[0]+axes[0,1]], [center[1],center[1]+axes[1,1]], 'k-')
        plt.axis('image')

    return pts


def draw_ellipse(center, axes, N):
    """Draw a 2-D ellipse.

    Arguements:
    center : 2-array
        Coordinates for the center of the ellipse.
    axes : 2x2 array
        Vectors (columns) representing the orthogonal axes of the ellipse, such as the output from
        get_ellipsoid_from_cov.

    """
    get_ellipse2d_pts(center, axes, N)

    plt.plot(x, y, 'r-')
    #plt.plot([center[0],center[0]+axes[0,0]], [center[1],center[1]+axes[1,0]], 'k-')
    #plt.plot([center[0],center[0]+axes[0,1]], [center[1],center[1]+axes[1,1]], 'k-')
    plt.axis('image')


def draw_moving_coordinate_system(positions, rotation_matrices, arrow_scale=1):
    """Draw coordinate system pose as a function of time.

    :param positions: Array of position 3-vectors.
    :type positions: N x 3

    :param rotation_matrices: Rotation matrix. The rows are the unit vectors
        pointing along the x, y, and z axes of the moving coordinate system
        defined within the base coordinate system.
    :type rotation_matrices: list of 3x3 arrays

    """
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    pmax = positions.max(axis=0)
    pmin = positions.min(axis=0)
    pcenter = (pmax + pmin)/2
    prange = pmax - pmin
    side_radius = prange.max()/2

    # Size of the arrows showing the axes.
    r = side_radius/20*arrow_scale

    plt.plot(positions[:, 0], positions[:, 1], positions[:, 2], 'k-')
    plt.plot(positions[:, 0], positions[:, 1], positions[:, 2], 'ro')

    for i in range(len(rotation_matrices)):
        pos = positions[i]
        R = rotation_matrices[i]

        s = ['r-','g-','b-']
        for i in range(3):
            plt.plot([pos[0],pos[0]+R[i][0]*r], [pos[1],pos[1]+R[i][1]*r],
                     [pos[2],pos[2]+R[i][2]*r], s[i], linewidth=3)

    ax.axes.set_xlim3d(left=pcenter[0] - side_radius,
                       right=pcenter[0] + side_radius)
    ax.axes.set_ylim3d(bottom=pcenter[1] - side_radius,
                       top=pcenter[1] + side_radius)
    ax.axes.set_zlim3d(bottom=pcenter[2] - side_radius,
                       top=pcenter[2] + side_radius)

    plt.xlabel('X-Axis')
    plt.ylabel('Y-Axis')

    plt.figure()
    plt.subplot(321)
    plt.plot(positions[:, 0])
    plt.ylabel('X-Position (m)')
    plt.subplot(323)
    plt.plot(positions[:, 1])
    plt.ylabel('Y-Position (m)')
    plt.subplot(325)
    plt.plot(positions[:, 2])
    plt.ylabel('Z-Position (m)')

    ypr = np.array([euler_from_matrix(rmat.T, axes='rzyx')
                    for rmat in rotation_matrices])
    ypr *= 180/np.pi

    plt.subplot(322)
    plt.plot(ypr[:, 0])
    plt.ylabel('Yaw (deg)')
    plt.subplot(324)
    plt.plot(ypr[:, 1])
    plt.ylabel('Pitch (deg)')
    plt.subplot(326)
    plt.plot(ypr[:, 2])
    plt.ylabel('Roll (deg)')


def horn(P, Q, fit_translation=True):
    """
    :param P: The "from" vectors.
    :type P: 3 x N array-like

    :param Q: The "to" vectors.
    :type Q: 3 x N array-like


    """
    if P.shape != Q.shape:
        raise Exception()

    if fit_translation:
        centroids_P = np.mean(P, axis=1)
        centroids_Q = np.mean(Q, axis=1)
        A = P - np.outer(centroids_P, np.ones(P.shape[1]))
        B = Q - np.outer(centroids_Q, np.ones(Q.shape[1]))
    else:
        A = P
        B = Q

    C = np.dot(A, B.transpose())
    U, S, V = np.linalg.svd(C)
    R = np.dot(V.transpose(), U.transpose())
    if(np.linalg.det(R) < 0):
        L = np.eye(3)
        L[2][2] *= -1
        R = np.dot(V.transpose(), np.dot(L, U.transpose()))

    if fit_translation:
        t = np.dot(-R, centroids_P) + centroids_Q
        return R, t
    else:
        return R