from __future__ import division
import numpy as np
from random import sample

def generatePoints(n, min_xyz, max_xyz, grid=False, randomZ=True, randFixedZ=False, depth=None, offset=0.5, xoffset=0, yoffset=0, zoffset=0):
  '''
  Generates 3D points either on a grid or randomly. the depth of each point can also be selected
  at random, given fixed depth for all points, or a random fixed depth for the points.

  UPDATE: the order of the points in the end is also randomly shuffled.
  '''
  if randFixedZ: # means all points have the same z but z is chosen at random between max and min z
    z = min_xyz[2] + np.random.random() * (max_xyz[2] - min_xyz[2])
  else: # same depth
    if not isinstance(depth, list) and not depth: # depth is exactly the middle of min and max z
        z = min_xyz[2] + (max_xyz[2] - min_xyz[2]) / 2
    else:
        z = depth
  if not grid: # compute randomly
    xr, yr, zr = max_xyz[0] - min_xyz[0], max_xyz[1] - min_xyz[1], max_xyz[2] - min_xyz[2]

    xr = np.random.rand(1, n)[0] * xr + min_xyz[0]
    yr = np.random.rand(1, n)[0] * yr + min_xyz[1]
    if randomZ:
        zr = np.random.rand(1, n)[0] * zr + min_xyz[2]
    else:
        zr = np.ones((1, n))[0] * z
    return zip(xr, yr, zr)
  else: # compute points on a mXm grid when m = sqrt(n)
    m = int(np.sqrt(n))
    gwx = (max_xyz[0] - min_xyz[0]) / m
    gwy = (max_xyz[1] - min_xyz[1]) / m
    zr = max_xyz[2] - min_xyz[2]
    if randomZ:
        return [(min_xyz[0] + (i+offset) * gwx + xoffset, 
                 min_xyz[1] + (j+offset) * gwy + yoffset,     
                 np.random.random() * zr + min_xyz[2] + zoffset) for i in xrange(m) for j in xrange(m)]
    else:
        if not isinstance(depth, list):
            ret = [(min_xyz[0] + (i+offset) * gwx + xoffset, # offset .5
                     min_xyz[1] + (j+offset) * gwy + yoffset, # offset .5
                     z + zoffset) for i in xrange(m) for j in xrange(m)]
            # return ret
            return sample(ret, len(ret)) # this shuffles the points
        else:
            ret = []
            for dz in depth:
                ret.extend([(min_xyz[0] + (i+offset) * gwx + xoffset, 
                             min_xyz[1] + (j+offset) * gwy + yoffset, 
                             dz + zoffset) for i in xrange(m) for j in xrange(m)])
            # return ret
            return sample(ret, len(ret)) # this shuffles the points


def getSphericalCoords(x, y, z):
    '''
    According to our coordinate system, this returns the 
    spherical coordinates of a 3D vector.
    A vector originating from zero and pointing to the positive Z direction (no X or Y deviation)
    will correspond to (teta, phi) = (0, 90)  (in degrees)
    The coordinate system we are using is similar to https://en.wikipedia.org/wiki/File:3D_Spherical_2.svg

         Y
         |
         |
         |______X
        / 
       /
      /
     Z   

    with a CCW rotation of the axis vectors
    '''
    r = np.sqrt(x*x + y*y + z*z)
    teta = np.arctan(x/z)
    phi = np.arccos(y/r)
    return teta, phi, r

def getAngularDiff(T, E, C):
    '''
    T is the target point
    E is the estimated target
    C is camera center
    Returns angular error

    (using law of cosines: http://mathcentral.uregina.ca/QQ/database/QQ.09.07/h/lucy1.html)
    '''
    t = (E - C).mag
    e = (C - T).mag
    c = (T - E).mag
    return np.degrees(np.arccos((e*e + t*t - c*c)/(2*e*t)))

def getRotationMatrixFromAngles(r):
    '''
    Returns a rotation matrix by combining elemental rotations 
    around x, y', and z''

    It also appends a zero row, so the end result looks like:
    [R_11 R_12 R_13]
    [R_11 R_12 R_13]
    [R_11 R_12 R_13]
    [0    0    0   ] 

    This is convenient since we can insert a translation column
    and get a full transformation matrix for vectors in homogeneous
    coordinates
    '''
    cos = map(np.cos, r)
    sin = map(np.sin, r)
    Rx = np.array([
        [1,      0,       0],
        [0, cos[0], -sin[0]],
        [0, sin[0], cos[0]]])
    Ry = np.array([
        [ cos[1], 0, sin[1]],
        [ 0     , 1,      0],
        [-sin[1], 0, cos[1]]])
    Rz = np.array([
        [cos[2], -sin[2], 0],
        [sin[2],  cos[2], 0],
        [0     ,       0, 1]])
    R = Rz.dot(Ry.dot(Rx))
    return np.concatenate((R, [[0, 0, 0]]))

# import cv2
# def getRotationMatrix(a, b):
#     y = a[1] - b[1]
#     z = a[2] - b[2]
#     x = a[0] - b[0]
#     rotx = np.arctan(y/z)
#     roty = np.arctan(x*np.cos(rotx)/z)
#     rotz = np.arctan(np.cos(rotx)/(np.sin(rotx)*np.sin(roty)))
#     return cv2.Rodrigues(np.array([rotx, roty, rotz]))[0]

def getRotationMatrix(a, b):
    '''
    Computes the rotation matrix that maps unit vector a to unit vector b
    
    It also augments a zero row, so the end result looks like:
    [R_11 R_12 R_13]
    [R_11 R_12 R_13]
    [R_11 R_12 R_13]
    [0    0    0   ] 

    (simply slice the output like R = output[:3] to get only the rotation matrix)

    based on the solution here:
    https://math.stackexchange.com/questions/180418/calculate-rotation-matrix-to-align-vector-a-to-vector-b-in-3d
    '''
    a, b = np.array(a), np.array(b)
    v = np.cross(a, b, axis=0)
    s = np.linalg.norm(v) # sine of angle
    c = a.dot(b) # cosine of angle

    vx = np.array([
        [0    , -v[2],  v[1]],
        [v[2] ,     0, -v[0]],
        [-v[1],  v[0],    0]]) 

    if s == 0: # a == b
        return np.concatenate((np.eye(3), [[0, 0, 0]]))
    if c == 1:
        return np.concatenate((np.eye(3) + vx, [[0, 0, 0]]))
    return np.concatenate((np.eye(3) + vx + vx.dot(vx)*((1-c)/s/s), [[0, 0, 0]]))

# Default FOV angle for a camera
def_fov = np.pi * 2./3

class PinholeCamera:
    '''
    Models a Pinhole Camera with 9 degrees of freedom
    '''
    ################################################################################################
    ## Intrinsic parameters
    f = 1 # focal length
    p = (0, 0) # position of principal point in the image plane
    
    # Extrinsic parameters
    # this rotation corresponds to a camera setting pointing towards (0, 0, -1)
    r = (0, 0, 0) # rotations in x, y', and z'' planes respectively 
    t = (0, 0, 0) # camera center translation w.r.t world coordinate system
    #
    # Using the above parameters we can construct the camera matrix
    #
    #                      [f 0 p.x 0]
    # P = K[R|t] where K = [0 f p.y 0] is the camera calibration matrix (full projection matrix)
    #                      [0 0  1  0]
    #
    # and thus we have x = PX for every point X in the word coordinate system 
    # and its corresponding projection in the camera image plane x
    # NOTE: points are assumed to be represented in homogeneous coordinates

    # Other parameters
    label = ''
    direction = (0, 0, 1) # camera direction vector (viewpoint direction)
    fov = def_fov # field of view (both horizontal and vertical)
    image_width = 2*f*np.tan(fov/2.)
    ################################################################################################
    
    def __init__(self, label, f = 1, r = (0, 0, 0), t = (0, 0, 0), direction = (0, 0, 1), fov=def_fov):
        self.label = label
        self.f = f
        self.r = r
        self.direction = direction
        self.t = t
        self.setFOV(fov)
        self.recomputeCameraMatrix(True, True)

    def recomputeCameraMatrix(self, changedRotation, changedIntrinsic):
        '''
        Updates camera matrix
        '''
        if changedRotation:
            # Computing rotation matrix using elemental rotations
            self.R = getRotationMatrixFromAngles(self.r)
            # by default if rotation is 0 then camera optical axis points to positive Z
            self.direction = np.array([0, 0, 1, 0]).dot(self.R) 

        # Computing the extrinsic matrix
        _t = -self.R.dot(np.array(self.t)) # t = -RC
        self.Rt = np.concatenate((self.R, np.array([[_t[0], _t[1], _t[2], 1]]).T), axis=1)
        # instead of the above, we could also represent translation matrix as [I|-C] and
        #                              [R -RC]
        # then compute Rt as R[I|-C] = [0   1]                [R]      [-RC]
        # but we basically do the same thing by concatenating [0] with [  1]

        if changedIntrinsic:
            # Computing intrinsic matrix
            f, px, py = self.f, self.p[0], self.p[1]
            self.K = np.array([
                [f, 0, px, 0],
                [0, f, py, 0],
                [0, 0, 1,  0]])

        # Full Camera Projection Matrix
        self.P = self.K.dot(self.Rt)
        
    ################################################################################################
    ## Intrinsic parameter setters
    def setF(self, f, auto_adjust = True):
        self.f = f
        if auto_adjust:
            self.setFOV(self.fov)
            self.recomputeCameraMatrix(False, True)
    def setP(self, p, auto_adjust = True):
        self.p = p
        if auto_adjust:
            self.recomputeCameraMatrix(False, True)
    def setFOV(self, fov):
        self.fov = fov
        self.image_width = 2*self.f*np.tan(fov/2.)
    
    ################################################################################################
    ## Extrinsic parameter setters
    def setT(self, t, auto_adjust = True):
        self.t = t
        if auto_adjust:
            self.recomputeCameraMatrix(False, False)
    def setR(self, r, auto_adjust = True):
      self.r = r
      if auto_adjust:
          self.recomputeCameraMatrix(True, False)
    
    ################################################################################################
    ## Main methods
    def project(self, p):
        '''
        Computes projection of a point p in world coordinate system
        using its homogeneous coordinates [x, y, z, 1]
        '''
        if len(p) < 4: p = (p[0], p[1], p[2], 1)
        projection = self.P.dot(np.array(p))
        # dividing by the Z value to get a 2D point from homogeneous coordinates
        return np.array((projection[0], projection[1]))/projection[2]

    def getNormalizedPts(self, pts):
        '''
        Returns normalized x and y coordinates in range [0, 1]
        '''
        px, py = self.p[0], self.p[1]
        return map(lambda p:np.array([p[0] - px + self.image_width/2, 
                                      p[1] - py + self.image_width/2]) / self.image_width, pts)
    def getDenormalizedPts(self, pts):
        '''
        Returns original points in the camera image coordinate plane from normalized points
        '''
        px, py = self.p[0], self.p[1]
        offset = np.array([self.image_width/2-px, self.image_width/2-py])
        return map(lambda p:(p*self.image_width)-offset, pts)

class Camera(PinholeCamera):
    '''
    A wrapper around abstract Pinhole camera object with further
    methods for visualization.
    '''
    default_radius = 2.0

    def updateShape(self):
        c = v(self.t)
        # Update camera center
        self.center.pos = c
        # Update the arrow
        self.dir.pos = c
        self.dir.axis = v(self.direction) * self.f
        # Update image plane
        self.img_plane.pos = c + self.dir.axis
        self.img_plane.length = self.image_width
        self.img_plane.height = self.image_width
        # TODO: handle rotation of image plane