import numpy as np
import math

R0 = 6.378137e6 #Equatorial radius of earth (meters)
ECCEN = 0.0818191908426 #Earth's eccentricity

#convert vector to quaternion (insert a leading 0)
def vec2quat(vec_input):
    """vec_input must be ndarray of shape (3, 1)"""
    quat_result = np.zeros(4)
    quat_result[0] = 0.0

    for idx in range(1, 4):
        quat_result[idx] = vec_input[(idx - 1)]

    return quat_result


#convert quaternion to vector (strip off leading 0)
def quat2vec(quat_input):
    """quat_input must be ndarray of shape (4, 1)"""
    vec_result = np.zeros(3)

    for idx in range(3):
        vec_result[idx] = quat_input[(idx + 1)]

    return vec_result


#quaternion multiplication
def quat_multiply(quat_A, quat_B):
    """quat_A and quat_B must be ndarrays of shape (4, 1)"""
    p0, q0 = quat_A[0], quat_B[0]
    p, q = quat_A[1:4], quat_B[1:4]

    quat_result = np.zeros(4)
    quat_result[0] = (p0 * q0) - np.dot(p, q)
    vector_part = (p0 * q) + (q0 * p) + np.cross(p, q)

    for idx in range(1, 4):
        quat_result[idx] = vector_part[idx-1]

    return quat_result


#compute the norm of a quaternion
def quat_norm(quat_input):
    norm_result = 0.0

    for quat_element in quat_input:
        norm_result = norm_result + (quat_element ** 2)

    return norm_result


#compute the inverse quaternion
def quat_inv(quat_input):
    scalar_part = quat_input[0]
    vec_part = quat2vec(quat_input)
    vec_part = -1.0 * vec_part
    inv_quaternion = vec2quat(vec_part)
    inv_quaternion[0] = scalar_part
    inv_quaternion = (1.0 / quat_norm(quat_input)) * inv_quaternion

    return inv_quaternion


#rotate a vector by a quaternion
def quat_rotate(vec_input, quat_input):
    quat_result = quat_multiply(quat_inv(quat_input), quat_multiply(vec2quat(vec_input), quat_input))
    vec_result = quat2vec(quat_result)

    return vec_result


#Derive Euler angles from (ENU) quaternion assuming a "ZYX" rotation sequence
def quat2angle(quaternion_input):
    q0 = quaternion_input[0]
    q1 = quaternion_input[1]
    q2 = quaternion_input[2]
    q3 = quaternion_input[3]

    roll_angle = (180.0 / math.pi) * math.atan2((2 * ((q0 * q1) + (q2 * q3))), (1.0 - 2.0 *((q1 ** 2) + (q2 ** 2))))
    pitch_angle = (180.0 / math.pi) * math.asin(2 * ((q0 * q2) - (q1 * q3)))
    yaw_angle = (180.0 / math.pi) * math.atan2((2 * ((q0 * q3) + (q1 * q2))), (1.0 - 2.0 *((q2 ** 2) + (q3 ** 2))))

    Euler_Angles = np.zeros(3)
    Euler_Angles[0] = roll_angle - 180.0
    Euler_Angles[1] = (-1.0) * pitch_angle
    Euler_Angles[2] = (-1.0) * (yaw_angle - 90.0)

    return Euler_Angles

#Derive NED to FRD quaternion from Euler angles assuming a "ZYX" rotation sequence
def angle2quat(angle_input):
    psi_rad = math.radians(angle_input[2])
    theta_rad = math.radians(angle_input[1])
    phi_rad = math.radians(angle_input[0])

    quat_psi = np.array([math.cos(psi_rad / 2), 0, 0, math.sin(psi_rad / 2)]);
    quat_theta = np.array([math.cos(theta_rad / 2), 0, math.sin(theta_rad / 2), 0]);
    quat_phi = np.array([math.cos(phi_rad / 2), math.sin(phi_rad / 2), 0, 0]);

    quat_result_i = quat_multiply(quat_psi, quat_multiply(quat_theta, quat_phi));
    quat_result = (1.0 / quat_norm(quat_result_i)) * quat_result_i

    return quat_result


#Calculate the East-West radius of curvature at latitude
def Re_calc(lat_deg):
    lat_rad = math.radians(lat_deg)
    Re_result = R0 / (math.sqrt(1.0 - ((ECCEN ** 2) * (math.sin(lat_rad) ** 2))))

    return Re_result


#Convert Curvilinear position to Cartesian coordinates
def curv2cart(curv_vec):
    lat_deg = curv_vec[0]
    lat_rad = math.radians(lat_deg)

    lon_deg = curv_vec[1]
    lon_rad = math.radians(lon_deg)

    h = curv_vec[2] #height (meters)

    Re = Re_calc(lat_deg)

    x = (Re + h) * math.cos(lat_rad) * math.cos(lon_rad)
    y = (Re + h) * math.cos(lat_rad) * math.sin(lon_rad)
    z = (((1 - (ECCEN ** 2)) * Re) + h) * math.sin(lat_rad)

    cart_posn = np.zeros(3)
    cart_posn[0] = x
    cart_posn[1] = y
    cart_posn[2] = z

    return cart_posn


#Convert Cartesian position to Curvilinear coordinates
def cart2curv(cart_vec):
    x = cart_vec[0]
    y = cart_vec[1]
    z = cart_vec[2]

    #calculate longitude
    lon_rad = math.atan2(y, x)
    lon_deg = math.degrees(lon_rad)

    #initialise the latitude solution to the geocentric latitude
    geo_lat_rad = math.atan(z / (math.sqrt((x ** 2) + (y ** 2))))
    geo_lat_deg = math.degrees(geo_lat_rad)
    lat_deg = geo_lat_deg
    lat_deg_minus = (1 / 100) * math.pi

    while (math.fabs(lat_deg - lat_deg_minus) > 1.3889e-04): #latitude should be correct to within 1 second
        lat_deg_minus = lat_deg
        lat_rad_minus = math.radians(lat_deg_minus)

        #Calculate Re
        Re = Re_calc(lat_deg_minus)

        #Calculate previous height
        h_minus = (math.sqrt((x ** 2) + (y ** 2)) / math.cos(lat_rad_minus)) - Re

        #Calculate updated latitude
        lat_rad = math.atan((z * (Re + h_minus)) / (math.sqrt((x ** 2) + (y ** 2)) * (((1.0 - (ECCEN ** 2)) * Re) + h_minus)))
        lat_deg = math.degrees(lat_rad)
        h = (math.sqrt((x ** 2) + (y ** 2)) / math.cos(lat_rad)) - Re

    curv_posn = np.zeros(3)
    curv_posn[0] = lat_deg
    curv_posn[1] = lon_deg
    curv_posn[2] = h

    return curv_posn

#Convert (lat, lon, alt) to NED
def lla2ned(posn_lla, posn_lla0):
    lat_deg = posn_lla[0]
    lat_deg0 = posn_lla0[0]
    posn_north = (lat_deg - lat_deg0) * 60 * 1855 #north position (meters)

    lon_deg = posn_lla[1]
    lon_deg0 = posn_lla0[1]
    posn_east = math.radians(lon_deg - lon_deg0) * Re_calc(lat_deg) #east position (meters)

    posn_down = (-1.0) * (posn_lla[2] - posn_lla0[2])

    result = np.zeros(3)
    result[0] = posn_north
    result[1] = posn_east
    result[2] = posn_down

    return result

#Convert NED to (lat, lon, alt)
def ned2lla(xNED, lla0):
    xNorth = xNED[0]
    xEast = xNED[1]
    xDown = xNED[2]

    yLat = math.degrees(xNorth / R0) + lla0[0]
    yLon = math.degrees(xEast / Re_calc(yLat)) + lla0[1]
    yAlt = -xDown + lla0[2]

    result = np.array([yLat, yLon, yAlt])

    return result