from matplotlib import pyplot as plt
import numpy as np
from random import random


fx = 0.0176  # focal length[m]
fy = 0.0176  # focal length[m]
Nu = 1920  # number of horizontal[pixels]
Nv = 1200  # number of vertical[pixels]
ppx = 5.86e-6  # horizontal pixel pitch[m / pixel]
ppy = ppx  # vertical pixel pitch[m / pixel]
fpx = fx / ppx  # horizontal focal length[pixels]
fpy = fy / ppy  # vertical focal length[pixels]

K = [[fpx, 0, Nu / 2],
     [0, fpy, Nv / 2],
     [0, 0, 1]]
K = np.array(K)

def quat2dcm(q):
    # Computing direction cosine matrix from quaternion
    # adapted from PyNav

    # normalizing quaternion
    q = q / np.linalg.norm(q)

    q0 = q[0]
    q1 = q[1]
    q2 = q[2]
    q3 = q[3]

    dcm = np.zeros((3, 3))

    dcm[0, 0] = 2 * q0 ** 2 - 1 + 2 * q1 ** 2
    dcm[1, 1] = 2 * q0 ** 2 - 1 + 2 * q2 ** 2
    dcm[2, 2] = 2 * q0 ** 2 - 1 + 2 * q3 ** 2

    dcm[0, 1] = 2 * q1 * q2 + 2 * q0 * q3
    dcm[0, 2] = 2 * q1 * q3 - 2 * q0 * q2

    dcm[1, 0] = 2 * q1 * q2 - 2 * q0 * q3
    dcm[1, 2] = 2 * q2 * q3 + 2 * q0 * q1

    dcm[2, 0] = 2 * q1 * q3 + 2 * q0 * q2
    dcm[2, 1] = 2 * q2 * q3 - 2 * q0 * q1

    return dcm

# Rotate points in body frame
def rotate(points, q):
    k = np.transpose(quat2dcm(q)).dot(np.transpose(points))
    return np.transpose(np.transpose(quat2dcm(q)).dot(np.transpose(points)))

# translate in body frame
def translate(points, r):
    return points + r

def ortho(points, q):
    X = np.transpose(quat2dcm(q)).dot(np.transpose(points)[:3])
    x, y = (X[0], X[1])
    return x, y

def pinhole(points, q, r):

    hom_points = np.hstack([points, np.ones((points.shape[0], 1))])
    points_body = np.transpose(hom_points)
#     points_body[:-1, :] = np.multiply(points_body[:-1, :], scale)  # scaling homogeneous coordinates

    # transformation to camera frame
    poseMat = np.hstack((np.transpose(quat2dcm(q)), np.expand_dims(r, 1)))
    p_cam = np.dot(poseMat, points_body)

    # getting homogeneous coordinates
    X = p_cam / p_cam[2]

    # projection to image plane
    X = K.dot(X)
    x, y = (X[0], X[1])
    return x, y