import numpy as np
import logging

MONITOR_POSITION_OLD = np.array([17.0, 0.0, 0.0])
MONITOR_POSITION_NEW = np.array([11.86, 8.62, 3.16])
MONITOR_ROTATIONS = np.array([0.0, 0.0, 0.0])

CAMERA_POSITION_OLD = np.array([13.0, 0, 0])
CAMERA_POSITION_NEW = np.array([10.28, 7.47, 2.74])
CAMERA_ROTATIONS_OLD = np.array([0.0, 0.0, 13.1*np.pi/180])
CAMERA_ROTATIONS_NEW = np.array([0.0, 0.0, 2.8*np.pi/180])

LED_POSITION_ORIGINAL = np.array([26.51, -3.93, 0.1])
LED_POSITION_OLD = np.array([25.89, -6.12, 3.21])
LED_POSITION_NEW = np.array([24.6, 9.23, 5.26])

EYE_RADIUS = 0.1682  # in cm
CM_PER_PIXEL = 10.2/10000.0

class EyeCalibration(object):
    '''Class for performing eye-tracking calibration.

    Provides methods for estimating the position of the pupil in
    3D space and projecting the gaze onto the monitor in both
    3D space and monitor space given the experimental geometry.

    Parameters
    ----------
    monitor_position : numpy.ndarray
        [x,y,z] position of monitor in cm.
    monitor_rotations : numpy.ndarray
        [x,y,z] rotations of monitor in radians.
    led_position : numpy.ndarray
        [x,y,z] position of LED in cm.
    camera_position : numpy.ndarray
        [x,y,z] position of camera in cm.
    camera_rotations : numpy.ndarray
        [x,y,z] rotations for camera in radians. X and Y must be 0.
    eye_radius : float
        Radius of the eye in cm.
    cm_per_pixel : float
        Pixel size of eye-tracking camera.
    '''
    def __init__(self, monitor_position=MONITOR_POSITION_NEW,
                 monitor_rotations=MONITOR_ROTATIONS,
                 led_position=LED_POSITION_OLD,
                 camera_position=CAMERA_POSITION_OLD,
                 camera_rotations=CAMERA_ROTATIONS_OLD,
                 eye_radius=EYE_RADIUS,
                 cm_per_pixel=CM_PER_PIXEL):
        '''Constructor.'''

        self.eye_radius = eye_radius
        self.cm_per_pixel = cm_per_pixel

        self.monitor_position = monitor_position
        self.led_position = led_position
        self.camera_position = camera_position

        self.cr = self.cr_position_in_mouse_eye_coordinates(led_position,
                                                            eye_radius)

        self.monitor_rotations = monitor_rotations
        if camera_rotations[0] != 0 or camera_rotations[1] != 0:
            logging.warning("Got nonzero x=%s,y=%s rotations for camera",
                            camera_rotations[0], camera_rotations[1])
        self.camera_rotation = camera_rotations[2]

    def pupil_position_in_mouse_eye_coordinates(self, pupil_parameters,
                                                cr_parameters):
        '''Compute the 3D pupil position in mouse eye coordinates.

        Parameters
        ----------
        pupil_parameters :  numpy.ndarray
            Array of pupil parameters for each eye tracking frame.
        cr_paramaeters : numpy.ndarray
            Array of corneal reflection parameters for each eye
            tracking frame.

        Returns
        -------
        numpy.ndarray
            Pupil position estimates in eye coordinates.
        '''
        # x, y are in screen coordinates, with y increasing towards the top
        delta_px = (pupil_parameters.T[0] - cr_parameters.T[0]) * \
            self.cm_per_pixel
        delta_py = (cr_parameters.T[1] - pupil_parameters.T[1]) * \
            self.cm_per_pixel  # +y is down on image

        R_cam_to_eye = base_object_to_eye_rotation_matrix(
            self.camera_position)
        # camera frame is passed to us pointed at the eye, but the image
        # appears as if the camera were rotated 180 degrees about its y-axis
        R_cam = object_rotation_matrix(0, np.pi, self.camera_rotation)
        cr_cam = np.dot(R_cam, np.dot(R_cam_to_eye.T, self.cr))

        px_cam = cr_cam[0] + delta_px
        py_cam = cr_cam[1] + delta_py
        pz_cam = np.sqrt(self.eye_radius**2 - px_cam**2 - py_cam**2)

        # estimating a position outside the eyeball is impossible, bad data
        bad_idx = np.sqrt(px_cam**2 + py_cam**2) > self.eye_radius
        px_cam[bad_idx] = np.nan
        py_cam[bad_idx] = np.nan
        pz_cam[bad_idx] = np.nan

        p_cam = np.vstack([px_cam, py_cam, pz_cam])

        # rotate estimates 
        return np.dot(R_cam_to_eye, np.dot(R_cam.T, p_cam)).T

    @staticmethod
    def cr_position_in_mouse_eye_coordinates(led_position, eye_radius):
        '''Determine the 3D position of the corneal reflection.

        The eye is modeled as a spherical mirror, so the reflection
        appears to be half the radius of the eye from the origin along
        the eye-LED axis.

        Parameters
        ----------
        led_position : numpy.ndarray
            [x,y,z] position of the LED in eye coordinates.
        eye_radius : float
            Radius of the eye in centimeters.

        Returns
        -------
        numpy.ndarray
            [x,y,z] location of the corneal reflection in eye coordinates.
        '''
        return (eye_radius/(2*np.linalg.norm(led_position))) * led_position

    def pupil_position_on_monitor_in_cm(self, pupil_parameters,
                                        cr_parameters):
        '''Compute the pupil position on the monitor in cm.

        Parameters
        ----------
        pupil_parameters :  numpy.ndarray
            Array of pupil parameters for each eye tracking frame.
        cr_paramaeters : numpy.ndarray
            Array of corneal reflection parameters for each eye
            tracking frame.

        Returns
        -------
        numpy.ndarray
            Pupil position estimates in eye coordinates.
        '''
        pupil_positions = self.pupil_position_in_mouse_eye_coordinates(
            pupil_parameters, cr_parameters)

        monitor_normal = object_norm_eye_coordinates(
            self.monitor_position, self.monitor_rotations[0],
            self.monitor_rotations[1], self.monitor_rotations[2])

        projected_positions = project_to_plane(monitor_normal,
                                               self.monitor_position,
                                               pupil_positions)

        monitor_positions = projected_positions - self.monitor_position

        R_monitor_to_eye = base_object_to_eye_rotation_matrix(
            self.monitor_position)
        R_monitor = object_rotation_matrix(self.monitor_rotations[0],
                                           self.monitor_rotations[1],
                                           self.monitor_rotations[2])

        result = np.dot(R_monitor.T,
                        np.dot(R_monitor_to_eye.T, monitor_positions.T))
        return result[:2].T

    def pupil_position_on_monitor_in_degrees(self, pupil_parameters,
                                             cr_parameters):
        '''Get pupil position on monitor measured in visual degrees.

        Parameters
        ----------
        pupil_parameters :  numpy.ndarray
            Array of pupil parameters for each eye tracking frame.
        cr_paramaeters : numpy.ndarray
            Array of corneal reflection parameters for each eye
            tracking frame.

        Returns
        -------
        numpy.ndarray
            Pupil position estimate in visual degrees.
        '''

        mag = np.sqrt(np.sum(self.monitor_position**2))

        pupil_pos = self.pupil_position_on_monitor_in_cm(pupil_parameters,
                                                         cr_parameters)

        x = pupil_pos.T[0]
        y = pupil_pos.T[1]

        meridian = np.arctan(x/mag)*180/np.pi
        elevation = np.arctan(y/np.sqrt(mag**2 + x**2))*180/np.pi

        angles = np.vstack([meridian, elevation]).T

        return angles

    def compute_area(self, pupil_parameters):
        '''Compute the area of the pupil.

        Assume the pupil is a circle, and that as it moves off-axis
        with the camera the observed ellipse major axis remains the
        diameter of the circle.

        Parameters
        ----------
        pupil_parameters : numpy.ndarray
            [nx5] array of pupil parameters.

        Returns
        -------
        numpy.ndarray
            [nx1] array of pupil areas in estimated pixels.
        '''
        r = np.maximum(pupil_parameters.T[3], pupil_parameters.T[4])
        return np.pi*r*r


def project_to_plane(plane_normal, plane_point, points):
    '''Project from the origin through points onto a plane.

    Parameters
    ----------
    plane_normal : numpy.ndarray
        [x, y, z] normal unit vector to the plane.
    plane_point : numpy.ndarray
        [x, y, z] point on the plane.
    points : numpy.ndarray
        [nx3] points in space through which to project.

    Returns
    -------
    numpy.ndarray
        [nx3] points projected on the plane.
    '''
    factor = np.sum(plane_normal*plane_point) / \
             np.sum(plane_normal*points, axis=1)
    return (factor*points.T).T


def object_norm_eye_coordinates(object_position, x_rotation,
                                y_rotation, z_rotation):
    '''Get the normal vector for the object plane in eye coordinates.

    Parameters
    ----------
    object_position : numpy.ndarray
        [x, y, z] location of the object in eye coordinates.
    x_rotation : float
        Rotation about the x-axis in radians.
    y_rotation : float
        Rotation about the y-axis in radians.
    z_rotation : float
        Rotation about the z-axis in radians.

    Returns
    -------
    numpy.ndarray
        Endpoint of the object plane vector in eye coordinates.
    '''
    R_object_to_eye = base_object_to_eye_rotation_matrix(object_position)
    R_object_frame = object_rotation_matrix(x_rotation, y_rotation,
                                            z_rotation)
    return np.dot(R_object_to_eye, np.dot(R_object_frame, [0, 0, 1]))


def base_object_to_eye_rotation_matrix(object_position):
    '''Rotation matrix to rotate base object frame to eye coordinates.

    By convention, any other object's coordinate frame before rotations
    is set with positive Z pointing from the object's position back
    to the origin of the eye coordinate system, with X parallel to the
    eye X-Y plane.

    Parameters
    ----------
    object_position : np.ndarray
        [x, y, z] position of object in eye coordinates.

    Returns
    -------
    numpy.ndarray
        [3x3] rotation matrix.
    '''
    eye_norm = -object_position/np.linalg.norm(object_position)

    # rotate about eye-z to align eye-x to object-x
    theta_z = -(np.pi/2 + np.arctan2(eye_norm[1], eye_norm[0]))
    Rz = np.array([[np.cos(theta_z), -np.sin(theta_z), 0],
                   [np.sin(theta_z), np.cos(theta_z), 0],
                   [0, 0, 1]])
    eye_norm_about_z = np.dot(Rz, eye_norm)

    # rotate about x' to align eye-z to object-z
    theta_x = np.pi/2 - np.arctan2(eye_norm_about_z[2], eye_norm_about_z[1])
    Rx = np.array([[1, 0, 0],
                   [0, np.cos(theta_x), -np.sin(theta_x)],
                   [0, np.sin(theta_x), np.cos(theta_x)]])

    R = np.dot(Rx, Rz).T
    return R


def object_rotation_matrix(x_rotation, y_rotation, z_rotation):
    '''Rotation matrix in object coordinate frame.

    The rotation matrix for rotating the object coordinate frame from
    the initial position. This is done by rotating around x, then
    around y', then around z''.

    Parameters
    ----------
    x_rotation : float
        Rotation about x axis in radians.
    y_rotation : float
        Rotation about y axis in radians.
    z_rotation : float
        Rotation about z axis in radians.

    Returns
    -------
    numpy.ndarray
        [3x3] rotation matrix.
    '''
    Rx = np.array([[1, 0, 0],
                   [0, np.cos(x_rotation), -np.sin(x_rotation)],
                   [0, np.sin(x_rotation), np.cos(x_rotation)]])
    Ry = np.array([[np.cos(y_rotation), 0, np.sin(y_rotation)],
                   [0, 1, 0],
                   [-np.sin(y_rotation), 0, np.cos(y_rotation)]])
    Rz = np.array([[np.cos(z_rotation), -np.sin(z_rotation), 0],
                   [np.sin(z_rotation), np.cos(z_rotation), 0],
                   [0, 0, 1]])
    result = np.dot(Rz, np.dot(Ry, Rx))
    return result
