#!/usr/bin/python3

import os,sys,math
import numpy as np
import matplotlib.pyplot as plt

__all__ = ["geodetic_rhoz_to_latalt",
    "geodetic_rhoz_to_latalt",
    "grs80_lonlatalt_to_xyz",
    "grs80_xyz_to_lonlatalt",
    "grs80_uvw_coordsys",
    "grs80_xyzvectors_to_uvw",
    "grs80_uvwvectors_to_xyz"
]

pi = np.pi
cos = np.cos
sin = np.sin
tan = np.tan
atan = np.atan
atan2 = np.atan2
sqrt = np.sqrt
abs = np.abs

def arg(x,y):
    x = np.array(x)
    y = np.array(y)

    single = False
    if(len(x.shape)==0):
        single = True
        x = np.reshape(x,[1])
        y = np.reshape(y,[1])

    th = np.atan2(y,x)
    if(th<0.0):
        th = th + 2*pi

    if(single):
        th = th[0]
        y = y[0]
    
    return th

def azel(x,y,z):
    x = np.array(x)
    y = np.array(y)
    z = np.array(z)

    single = False
    if(len(x.shape)==0):
        single = True
        x = np.reshape(x,[1])
        y = np.reshape(y,[1])
        z = np.reshape(z,[1])

    th = arg(x,y)
    rho = np.sqrt(x**2+y**2)
    el = np.atan(z/rho)

    if(single):
        th = th[0]
        el = el[0]
    
    return [th,el]

def geodetic_rhoz_to_latalt(p, z):
    """
    [lat, alt] = geodetic_rhoz_to_latalt(p, z):
    Inputs: 
        p - radial sqrt(x**2+y**2) coordinate [km]
        z - altitude [km]
    Outputs:
        geodetic latitude [deg]
        geodetic altitude [km]
    """

    #GRS-80 ellipsoid parameters
    a = 6378.137
    finv = 298.257222101
    b = a*(1.0-1.0/finv)
    e2 = (a**2-b**2)/(a**2)
    e2p = (a**2-b**2)/(b**2)

    #Input conditioning
    p = np.array(p)
    z = np.array(z)

    single = False
    if(len(p.shape)==0):
        single = True
        p = p.reshape([1])
        z = z.reshape([1])

    """
    //Ferari's Direct Quartic Solution for k
    //float ksi,rho2,s,t,u,v,w,k;
    // ksi = (1.0f-e*e)*(z/a)*(z/a);
    // rho2 = 1.0f/6.0f*((rho/a)*(rho/a)+ksi-e*e*e*e);
    // s = e*e*e*e*ksi*(rho/a)*(rho/a)/(4.0f*rho2*rho2*rho2);
    // t = powf(1+s+::sqrt(s*(s+2)),1.0f/3.0f);
    // u = rho2*(t+1+1/t);
    // v = ::sqrtf(u*u+e*e*e*e*ksi);
    // w = e*e*(u+v-ksi)/(2.0f*v);

    // k = 1+e*e*(::sqrtf(u+v+w*w)+w)/(u+v);

    // ret = ::atanf(k*z/rho);

    //given further down the wikipedia page
    // https://en.wikipedia.org/wiki/Geographic_coordinate_conversion#From_ECEF_to_geodetic_coordinates
    """

    p = np.abs(p)
    F = 54.0*b*b*z*z
    G = p*p+(1-e2)*z*z-e2*(a*a-b*b)
    c = (e2*e2*F*p*p)/(G*G*G)
    s = (1.0 + c + np.sqrt(c*c+2*c))**(1.0/3.0)
    k = s + 1.0 + 1.0/s
    P = F/(3.0*k*k*G*G)
    Q = np.sqrt(1.0+2.0*e2*e2*P)
    r0 = -(P*e2*p)/(1.0+Q) + np.sqrt(0.5*a*a*(1.0+1.0/Q)-(P*(1.0-e2)*z*z)/(Q*(1.0+Q))-0.5*P*p*p)
    U = np.sqrt((p-e2*r0)*(p-e2*r0)+z*z)
    V = np.sqrt((p-e2*r0)*(p-e2*r0)+(1.0-e2)*z*z)
    z0 = b*b*z/(a*V)
    alt = U*(1.0-b*b/(a*V))
    lat = np.atan((z+e2p*z0)/p)  

    lat = lat * 180.0/pi

    #Output conditioning
    if(single):
        lat = lat[0]
        alt = alt[0]

    return [lat,alt]

def geodetic_latalt_to_rhoz(lat,alt):
    """
    [lat, alt] = geodetic_latalt_to_rhoz(lat, alt):
    Inputs: 
        geodetic latitude [deg]
        geodetic altitude [km]]
    Outputs:
        p - radial sqrt(x**2+y**2) coordinate [km]
        z - altitude [km]
    """

    #GRS-80 ellipsoid parameters
    a = 6378.137
    finv = 298.257222101
    b = a*(1.0-1.0/finv)
    e2 = (a**2-b**2)/(a**2)
    e2p = (a**2-b**2)/(b**2)

    #Input Conditioning
    lat = np.array(lat)
    alt = np.array(alt)

    single = False
    if(len(lat.shape)==0):
        single = True
        lat = lat.reshape([1])
        alt = alt.reshape([1])
    
    lat = lat * pi/180.0
    # lat[lat>90.0] = 90.0-lat[lat>90.0]
    # lat[lat<-90.0] = -90.0 - lat[lat<-90.0]
    N = a**2/np.sqrt(a**2*cos(lat)**2+b**2*sin(lat)**2)
    p = np.abs((N+alt)*np.cos(lat))
    z = ((b/a)**2*N+alt)*np.sin(lat)

    #Output Conditioning
    if(single):
        p = p[0]
        z = z[0]

    return [p,z]

def grs80_lonlatalt_to_xyz(lon,lat,alt):
    """
    [x,y,z] = grs80_lonlatalt_to_xyz(lon,lat,alt)
    Inputs:
        lon [deg]
        lat [deg]
        alt [km]
    Outputs:
        x,y,z [km] ECEF coordinates
    """
    lon = np.array(lon)
    lat = np.array(lat)
    alt = np.array(alt)

    single = False
    if(len(lon.shape)==0):
        single = True
        lon = lon.reshape([1])
        lat = lat.reshape([1])
        alt = alt.reshape([1])

    [rho,z] = geodetic_latalt_to_rhoz(lat,alt)
    x = rho*np.cos(lon*pi/180.0)
    y = rho*np.sin(lon*pi/180.0)

    if(single):
        x = x[0]
        y = y[0]
        z = z[0]

    return [x,y,z]

def grs80_xyz_to_lonlatalt(x,y,z):
    """
    [lon,lat,alt] = grs80_xyz_to_lonlatalt(x,y,z)
    Inputs:
        x,y,z [km] ECEF coordinates
    Outputs:
        lon [deg]
        lat [deg]
        alt [km]
    """

    x = np.array(x); y = np.array(y); z = np.array(z)
    single = False
    if(len(x.shape)==0):
        single = True
        x = x.reshape([1]); y = y.reshape([1]); z = z.reshape([1])
    
    rho = np.sqrt(x**2+y**2)
    [lat,alt] = geodetic_rhoz_to_latalt(rho,z)
    lon = np.atan2(y,x)*180.0/pi

    if(single):
        lon = lon[0]
        lat = lat[0]
        alt = alt[0]

    return [lon,lat,alt]

def grs80_uvw_coordsys(x,y,z):
    """
    [u,v,w] = grs80_uvw_coordsys(x,y,z)
    Args:
        x,y,z - [km] vectors of ECEF coordinates
    Returns:
        [u,v,w] - each [3,N] unit vectors pointing along zonal, meridional, and vertical directions
    """

    x = np.array(x); y = np.array(y); z = np.array(z)
    single = False
    if(len(x.shape)==0):
        single = True
        x = x.reshape([1]); y = y.reshape([1]); z = z.reshape([1])

    N = x.shape[0]
    u = np.zeros([3,N])
    v = np.zeros([3,N])
    w = np.zeros([3,N])

    rho = np.sqrt(x**2+y**2)
    [lat,alt] = geodetic_rhoz_to_latalt(rho,z)
    lat = lat * pi/180.0
    lon = np.atan2(y,x)

    u[0,:] = -sin(lon)
    u[1,:] = cos(lon)
    u[2,:] = 0.0
    v[0,:] = -cos(lon)*sin(lat)
    v[1,:] = -sin(lon)*sin(lat)
    v[2,:] = cos(lat)
    w[0,:] = cos(lon)*cos(lat)
    w[1,:] = sin(lon)*cos(lat)
    w[2,:] = sin(lat)

    return [u,v,w]

def _local_vdot(v1,v2):

    v1 = np.array(v1)
    v2 = np.array(v2)

    single = False
    if(len(v1.shape)==1):
        single = True
        v1 = v1.reshape([3,1])
        v2 = v2.reshape([3,1])

    dots = np.sum(v1*v2,axis=0)

    if(single):
        dots = dots[0]

    return dots

def grs80_xyzvectors_to_uvw(xyzpos,xyzvect):
    """
    [lonlatalt, uvwvect] = grs80_xyzvectors_to_uvw(xyzpos,xyzvect)
    Inputs:
        xyzpos [km] - [3,N] array of vector ECEF positions
        xyzvect - [3,N] array of vectors
    Outputs:
        lonlatalt [deg, deg, km] - [3,N] array of vector geodetic coordinates
        uvwvect - [3,N] array of vectors in terms of local zonal, meridional, and vertical components
    """
    xyzpos = np.array(xyzpos)
    xyzvect = np.array(xyzvect)

    single = False
    if(len(xyzpos.shape)==1):
        single = True
        xyzpos = xyzpos.reshape([3,1])
        xyzvect = xyzvect.reshape([3,1])

    N = xyzpos.shape[1]

    [lon,lat,alt] = grs80_xyz_to_lonlatalt(xyzpos[0,:],xyzpos[1,:],xyzpos[2,:])
    lonlatalt = np.stack([lon,lat,alt],axis=0)
    [u,v,w] = grs80_uvw_coordsys(xyzpos[0,:],xyzpos[1,:],xyzpos[2,:])
    
    uc = _local_vdot(xyzvect,u)
    vc = _local_vdot(xyzvect,v)
    wc = _local_vdot(xyzvect,w)
    
    uvwvect = np.stack([uc,vc,wc],axis=0)

    if(single):
        lonlatalt = np.squeeze(lonlatalt)
        uvwvect = np.squeeze(uvwvect)

    return [lonlatalt, uvwvect]

def grs80_uvwvectors_to_xyz(lonlatalt, uvwvect):
    """
    [xyzpos, xyzvect] = grs80_uvwvectors_to_xyz(lonlatalt, uvwvect)
    Inputs:
        lonlatalt [deg, deg, km] - [3,N] array of vector geodetic coordinates
        uvwvect - [3,N] array of vectors in terms of local zonal, meridional, and vertical components
    Outputs:
        xyzpos [km] - [3,N] array of vector ECEF positions
        xyzvect - [3,N] array of vectors
    """
    lonlatalt = np.array(lonlatalt)
    uvwvect = np.array(uvwvect)

    single = False
    if(len(lonlatalt.shape)==1):
        single = True
        lonlatalt = lonlatalt.reshape([3,1])
        uvwvect = uvwvect.reshape([3,1])

    [x,y,z] = grs80_lonlatalt_to_xyz(lonlatalt[0,:],lonlatalt[1,:],lonlatalt[2,:])
    xyzpos = np.stack([x,y,z],axis=0)

    N = lonlatalt.shape[1]
    xyzvect = np.zeros([3,N])

    [u,v,w] = grs80_uvw_coordsys(x,y,z)

    xyzvect[0,:] = u[0]*uvwvect[0,:] + v[0]*uvwvect[1,:] + w[0]*uvwvect[2,:]
    xyzvect[1,:] = u[1]*uvwvect[0,:] + v[1]*uvwvect[1,:] + w[1]*uvwvect[2,:]
    xyzvect[2,:] = u[2]*uvwvect[0,:] + v[2]*uvwvect[1,:] + w[2]*uvwvect[2,:]

    if(single==True):
        xyzpos = np.squeeze(xyzpos)
        xyzvect = np.squeeze(xyzvect)

    return [xyzpos, xyzvect]

###########
## Tests ##
###########

def test_geodetic_latitude():

    th = np.linspace(0,2*pi,100)
    
    a = 6378.137
    p = cos(th)*a
    z = sin(th)*a

    [lat,alt] = geodetic_rhoz_to_latalt(p,z)
    [p2,z2] = geodetic_latalt_to_rhoz(lat,alt)
    [lat2,alt2] = geodetic_rhoz_to_latalt(p2,z2)

    err1 = np.abs(np.abs(p2)-np.abs(p))
    err2 = np.abs(z2-z)
    err3 = np.abs(lat2-lat)
    err4 = np.abs(alt2-alt)

    for I in range(0,len(th)):
        ln = "{:1.3f}: {:1.3f},{:1.3f}: {:1.3f},{:1.3f} : {:1.3f} {:1.3f} {:1.3f} {:1.3f}".format(th[I],p[I],z[I],lat[I],alt[I],err1[I],err2[I],err3[I],err4[I])
        print(ln)

    a = 360000.0
    p = cos(th)*a
    z = sin(th)*a

    [lat,alt] = geodetic_rhoz_to_latalt(p,z)
    [p2,z2] = geodetic_latalt_to_rhoz(lat,alt)
    [lat2,alt2] = geodetic_rhoz_to_latalt(p2,z2)

    err1 = np.abs(np.abs(p2)-np.abs(p))
    err2 = np.abs(z2-z)
    err3 = np.abs(lat2-lat)
    err4 = np.abs(alt2-alt)

    for I in range(0,len(th)):
        ln = "{:1.3f}: {:1.3f},{:1.3f}: {:1.3f},{:1.3f} : {:1.3f} {:1.3f} {:1.3f} {:1.3f}".format(th[I],p[I],z[I],lat[I],alt[I],err1[I],err2[I],err3[I],err4[I])
        print(ln)

    a = 100.0
    p = cos(th)*a
    z = sin(th)*a

    [lat,alt] = geodetic_rhoz_to_latalt(p,z)
    [p2,z2] = geodetic_latalt_to_rhoz(lat,alt)
    [lat2,alt2] = geodetic_rhoz_to_latalt(p2,z2)

    err1 = np.abs(np.abs(p2)-np.abs(p))
    err2 = np.abs(z2-z)
    err3 = np.abs(lat2-lat)
    err4 = np.abs(alt2-alt)

    for I in range(0,len(th)):
        ln = "{:1.3f}: {:1.3f},{:1.3f}: {:1.3f},{:1.3f} : {:1.3f} {:1.3f} {:1.3f} {:1.3f}".format(th[I],p[I],z[I],lat[I],alt[I],err1[I],err2[I],err3[I],err4[I])
        print(ln)


    return

def test_grs80_1():

    [x,y,z] = grs80_lonlatalt_to_xyz(45,0,0)
    print(x,y,z)
    [x,y,z] = grs80_lonlatalt_to_xyz(0,45,0)
    print(x,y,z)

    [lon,lat,alt] = grs80_xyz_to_lonlatalt(6371,0,0)
    print(lon,lat,alt)

    lon = np.linspace(-180,180,50)
    lat = np.linspace(-90,90,25)
    [lon2,lat2] = np.meshgrid(lon,lat,indexing='ij')

    [x,y,z] = grs80_lonlatalt_to_xyz(lon2,lat2,0)
    [x2,y2,z2] = grs80_lonlatalt_to_xyz(lon2,lat2,-6600)

    plt.figure().add_subplot(projection="3d")
    plt.plot(x,y,z,'k.',markersize=1)
    plt.plot(x2,y2,z2,'r.',markersize=1)
    plt.show()

    return


if(__name__=="__main__"):

    #test_geodetic_latitude()
    test_grs80_1()

    exit(0)