# Python tools to manage netCDF files.
# L. Fita, CIMA. Mrch 2019
# More information at: http://www.xn--llusfb-5va.cat/python/PyNCplot
#
# pyNCplot and its component geometry_tools.py comes with ABSOLUTELY NO WARRANTY. 
# This work is licendes under a Creative Commons 
#   Attribution-ShareAlike 4.0 International License (http://creativecommons.org/licenses/by-sa/4.0)
#
## Script for geometry calculations and operations as well as definition of different 
###    standard objects and shapes

import numpy as np
import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt

errormsg = 'ERROR -- error -- ERROR -- error'

####### Contents:
# deg_deci: Function to pass from degrees [deg, minute, sec] to decimal angles [rad]
# dist_points: Function to provide the distance between two points
# multi_rotate_2D: Function to rotate multiple vectors by a certain angle in the plain
# position_sphere: Function to tranform fom a point in lon, lat deg coordinates to 
#   cartesian coordinates over an sphere
# rotate_2D: Function to rotate a vector by a certain angle in the plain
# rotate_line2D: Function to rotate a line given by 2 pairs of x,y coordinates by a 
#   certain angle in the plain
# rotate_lines2D: Function to rotate multiple lines given by mulitple pars of x,y 
#   coordinates by a certain angle in the plain
# spheric_line: Function to transform a series of locations in lon, lat coordinates 
#   to x,y,z over an 3D spaceFunction to provide coordinates of a line  on a 3D space

## Shapes/objects
# ellipse_polar: Function to determine an ellipse from its center and polar coordinates
# surface_sphere: Function to provide an sphere as matrix of x,y,z coordinates

## Plotting
# plot_sphere: Function to plot an sphere and determine which standard lines will be 
#   also drawn

def deg_deci(angle):
    """ Function to pass from degrees [deg, minute, sec] to decimal angles [rad]
      angle: list of [deg, minute, sec] to pass
    >>> deg_deci([41., 58., 34.])
    0.732621346072
    """
    fname = 'deg_deci'

    deg = np.abs(angle[0]) + np.abs(angle[1])/60. + np.abs(angle[2])/3600.

    if angle[0] < 0.: deg = -deg*np.pi/180.
    else: deg = deg*np.pi/180.

    return deg

def position_sphere(radii, alpha, beta):
    """ Function to tranform fom a point in lon, lat deg coordinates to cartesian  
          coordinates over an sphere
      radii: radii of the sphere
      alpha: longitude of the point
      beta: latitude of the point
    >>> position_sphere(10., 30., 45.)
    (0.81031678432964027, -5.1903473778327376, 8.5090352453411846
    """
    fname = 'position_sphere'

    xpt = radii*np.cos(beta)*np.cos(alpha)
    ypt = radii*np.cos(beta)*np.sin(alpha)
    zpt = radii*np.sin(beta)

    return xpt, ypt, zpt

def spheric_line(radii,lon,lat):
    """ Function to transform a series of locations in lon, lat coordinates to x,y,z 
          over an 3D space
      radii: radius of the sphere
      lon: array of angles along longitudes
      lat: array of angles along latitudes
    """
    fname = 'spheric_line'

    Lint = lon.shape[0]
    coords = np.zeros((Lint,3), dtype=np.float)

    for iv in range(Lint):
        coords[iv,:] = position_sphere(radii, lon[iv], lat[iv])

    return coords

def rotate_2D(vector, angle):
    """ Function to rotate a vector by a certain angle in the plain
      vector= vector to rotate [y, x]
      angle= angle to rotate [rad]
    >>> rotate_2D(np.array([1.,0.]), np.pi/4.)
    [ 0.70710678 -0.70710678]
    """
    fname = 'rotate_2D'

    rotmat = np.zeros((2,2), dtype=np.float)

    rotmat[0,0] = np.cos(angle)
    rotmat[0,1] = -np.sin(angle)
    rotmat[1,0] = np.sin(angle)
    rotmat[1,1] = np.cos(angle)

    rotvector = np.zeros((2), dtype=np.float)

    vecv = np.zeros((2), dtype=np.float)

    # Unifying vector
    modvec = vector[0]**2+vector[1]**2
    if modvec != 0: 
        vecv[0] = vector[1]/modvec
        vecv[1] = vector[0]/modvec

        rotvec = np.matmul(rotmat, vecv)
        rotvec = np.where(np.abs(rotvec) < 1.e-7, 0., rotvec)

        rotvector[0] = rotvec[1]*modvec
        rotvector[1] = rotvec[0]*modvec

    return rotvector

def multi_rotate_2D(vectors, angle):
    """ Function to rotate multiple vectors by a certain angle in the plain
      line= matrix of vectors to rotate
      angle= angle to rotate [rad]
    >>> square = np.zeros((4,2), dtype=np.float)
    >>> square[0,:] = [-0.5,-0.5]
    >>> square[1,:] = [0.5,-0.5]
    >>> square[2,:] = [0.5,0.5]
    >>> square[3,:] = [-0.5,0.5]
    >>> multi_rotate_2D(square, np.pi/4.)
    [[-0.70710678  0.        ]
     [ 0.         -0.70710678]
     [ 0.70710678  0.        ]
     [ 0.          0.70710678]]
    """
    fname = 'multi_rotate_2D'

    rotvecs = np.zeros(vectors.shape, dtype=np.float)

    Nvecs = vectors.shape[0]
    for iv in range(Nvecs):
        rotvecs[iv,:] = rotate_2D(vectors[iv,:], angle)

    return rotvecs

def rotate_line2D(line, angle):
    """ Function to rotate a line given by 2 pairs of x,y coordinates by a certain 
          angle in the plain
      line= line to rotate as couple of points [[y0,x0], [y1,x1]]
      angle= angle to rotate [rad]
    >>> rotate_line2D(np.array([[0.,0.], [1.,0.]]), np.pi/4.)
    [[ 0.          0.        ]
     [0.70710678  -0.70710678]]
    """
    fname = 'rotate_2D'

    rotline = np.zeros((2,2), dtype=np.float)
    rotline[0,:] = rotate_2D(line[0,:], angle)
    rotline[1,:] = rotate_2D(line[1,:], angle)

    return rotline

def rotate_lines2D(lines, angle):
    """ Function to rotate multiple lines given by mulitple pars of x,y coordinates  
          by a certain angle in the plain
      line= matrix of N couples of points [N, [y0,x0], [y1,x1]]
      angle= angle to rotate [rad]
    >>> square = np.zeros((4,2,2), dtype=np.float)
    >>> square[0,0,:] = [-0.5,-0.5]
    >>> square[0,1,:] = [0.5,-0.5]
    >>> square[1,0,:] = [0.5,-0.5]
    >>> square[1,1,:] = [0.5,0.5]
    >>> square[2,0,:] = [0.5,0.5]
    >>> square[2,1,:] = [-0.5,0.5]
    >>> square[3,0,:] = [-0.5,0.5]
    >>> square[3,1,:] = [-0.5,-0.5]
    >>> rotate_lines2D(square, np.pi/4.)
    [[[-0.70710678  0.        ]
      [ 0.         -0.70710678]]

     [[ 0.         -0.70710678]
      [ 0.70710678  0.        ]]

     [[ 0.70710678  0.        ]
      [ 0.          0.70710678]]

     [[ 0.          0.70710678]
      [-0.70710678  0.        ]]]
    """
    fname = 'rotate_lines2D'

    rotlines = np.zeros(lines.shape, dtype=np.float)

    Nlines = lines.shape[0]
    for il in range(Nlines):
        line = np.zeros((2,2), dtype=np.float)
        line[0,:] = lines[il,0,:]
        line[1,:] = lines[il,1,:]

        rotlines[il,:,:] = rotate_line2D(line, angle)

    return rotlines

####### ###### ##### #### ### ## #
# Shapes/objects

def surface_sphere(radii,Npts):
    """ Function to provide an sphere as matrix of x,y,z coordinates
      radii: radii of the sphere
      Npts: number of points to discretisize longitues (half for latitudes)
    """
    fname = 'surface_sphere'

    sphereup = np.zeros((3,Npts/2,Npts), dtype=np.float)
    spheredown = np.zeros((3,Npts/2,Npts), dtype=np.float)
    for ia in range(Npts):
        alpha = ia*2*np.pi/(Npts-1)
        for ib in range(Npts/2):
            beta = ib*np.pi/(2.*(Npts/2-1))
            sphereup[:,ib,ia] = position_sphere(radii, alpha, beta)
        for ib in range(Npts/2):
            beta = -ib*np.pi/(2.*(Npts/2-1))
            spheredown[:,ib,ia] = position_sphere(radii, alpha, beta)

    return sphereup, spheredown

def ellipse_polar(c, a, b, Nang=100):
    """ Function to determine an ellipse from its center and polar coordinates
        FROM: https://en.wikipedia.org/wiki/Ellipse
      c= coordinates of the center
      a= distance major axis
      b= distance minor axis
      Nang= number of angles to use
    """
    fname = 'ellipse_polar'

    if np.mod(Nang,2) == 0: Nang=Nang+1
  
    dtheta = 2*np.pi/(Nang-1)

    ellipse = np.zeros((Nang,2), dtype=np.float)
    for ia in range(Nang):
        theta = dtheta*ia
        rad = a*b/np.sqrt( (b*np.cos(theta))**2 + (a*np.sin(theta))**2 )
        x = rad*np.cos(theta)
        y = rad*np.sin(theta)
        ellipse[ia,:] = [y+c[0],x+c[1]]

    return ellipse

def hyperbola_polar(a, b, Nang=100):
    """ Fcuntion to determine an hyperbola in polar coordinates
        FROM: https://en.wikipedia.org/wiki/Hyperbola#Polar_coordinates
          x^2/a^2 - y^2/b^2 = 1
      a= x-parameter
      y= y-parameter
      Nang= number of angles to use
      DOES NOT WORK!!!!
    """
    fname = 'hyperbola_polar'

    dtheta = 2.*np.pi/(Nang-1)

    # Positive branch
    hyperbola_p = np.zeros((Nang,2), dtype=np.float)
    for ia in range(Nang):
        theta = dtheta*ia
        x = a*np.cosh(theta)
        y = b*np.sinh(theta)
        hyperbola_p[ia,:] = [y,x]

    # Negative branch
    hyperbola_n = np.zeros((Nang,2), dtype=np.float)
    for ia in range(Nang):
        theta = dtheta*ia
        x = -a*np.cosh(theta)
        y = b*np.sinh(theta)
        hyperbola_n[ia,:] = [y,x]

    return hyperbola_p, hyperbola_n

def dist_points(ptA, ptB):
    """ Function to provide the distance between two points
      ptA: coordinates of the point A [yA, xA]
      ptB: coordinates of the point B [yB, xB]
    >>> dist_points([1.,1.], [-1.,-1.])
    2.82842712475
    """
    fname = 'dist_points'

    dist = np.sqrt( (ptA[0]-ptB[0])**2 + (ptA[1]-ptB[1])**2)

    return dist

def circ_sec(ptA, ptB, radii, Nang=100):
    """ Function union of point A and B by a section of a circle
      ptA= coordinates od the point A [yA, xA]
      ptB= coordinates od the point B [yB, xB]
      radii= radi of the circle to use to unite the points
      Nang= amount of angles to use
    """
    fname = 'circ_sec'

    distAB = dist_points(ptA,ptB)

    if distAB > radii:
        print errormsg
        print '  ' + fname + ': radii=', radii, " too small for the distance " +     \
          "between points !!"
        print '    distance between points:', distAB
        quit(-1)

    if ptA[0] > ptB[0]:
        tottheta = np.atan2(ptA[0]-ptB[0],ptA[1]-ptB[1])
        yc=ptA[0]
        if ptA[1] < ptB[0]: 
            sign = -1
            xc=ptA[1]+radii
        else: 
            sign = 1
            xc=ptA[1]-radii
    else:
        tottheta = np.atan2(ptB[0]-ptA[0],ptB[1]-ptA[1])
        yc=ptA[0]
        if ptA[1] < ptB[0]: 
            sign = 1
            xc=ptA[1]+radii
        else: 
            sign = -1
            xc=ptA[1]-radii

    dtheta = np.abs(tottheta)/(Nang-1)
    if sign == 1:
        dtheta = dtheta*(-1.)

    circ_sec = np.zeros((Nang,2), dtype=np.float)
    for ia in range(Nang):
        theta = dtheta*ia
        x = radii*np.cos(theta)
        y = radii*np.sin(theta)

        circ_sec[ia,:] = [x+xc,y+yc]
    
    return circ_sec

# FROM: http://www.photographers1.com/Sailing/NauticalTerms&Nomenclature.html
def zsailing_boat(length=10., beam=3., sternbp=0.5):
    """ Function to define an schematic boat from the z-plane
      length: length of the boat
      beam: beam of the boat
      sternbp: beam at stern as percentage of beam
    """
    fname = 'zsailing_boat'

    

    return boat

####### ####### ##### #### ### ## #
# Plotting

def plot_sphere(iazm=-60., iele=30., dist=10., Npts=100, radii=10,                   \
  drwsfc=[True,True], colsfc=['#AAAAAA','#646464'],                                  \
  drwxline = True, linex=[':','b',2.], drwyline = True, liney=[':','r',2.],          \
  drwzline = True, linez=['-.','g',2.], drwxcline=[True,True],                       \
  linexc=[['-','#646400',1.],['--','#646400',1.]],                                   \
  drwequator=[True,True], lineeq=[['-','#AA00AA',1.],['--','#AA00AA',1.]],           \
  drwgreeenwhich=[True,True], linegw=[['-','k',1.],['--','k',1.]]):
    """ Function to plot an sphere and determine which standard lines will be also 
        drawn
      iazm: azimut of the camera form the sphere
      iele: elevation of the camera form the sphere
      dist: distance of the camera form the sphere
      Npts: Resolution for the sphere
      radii: radius of the sphere
      drwsfc: whether 'up' and 'down' portions of the sphere should be drawn
      colsfc: colors of the surface of the sphere portions ['up', 'down']
      drwxline: whether x-axis line should be drawn
      linex: properties of the x-axis line ['type', 'color', 'wdith']
      drwyline: whether y-axis line should be drawn
      liney: properties of the y-axis line ['type', 'color', 'wdith']
      drwzline: whether z-axis line should be drawn
      linez: properties of the z-axis line ['type', 'color', 'wdith']
      drwequator: whether 'front' and 'back' portions of the Equator should be drawn
      lineeq: properties of the lines 'front' and 'back' of the Equator
      drwgreeenwhich: whether 'front', 'back' portions of Greenqhich should be drawn
      linegw: properties of the lines 'front' and 'back' Greenwhich
      drwxcline: whether 'front', 'back' 90 line (lon=90., lon=270.) should be drawn
      linexc: properties of the lines 'front' and 'back' for the 90 line
    """
    fname = 'plot_sphere'

    iazmrad = iazm*np.pi/180.
    ielerad = iele*np.pi/180.

    # 3D surface Sphere
    sfcsphereu, sfcsphered = surface_sphere(radii,Npts)
    
    # greenwhich
    if iazmrad > np.pi/2. and iazmrad < 3.*np.pi/2.:
        ia=np.pi-ielerad
    else:
        ia=0.-ielerad
    ea=ia+np.pi
    da = (ea-ia)/(Npts-1)
    beta = np.arange(ia,ea+da,da)[0:Npts]
    alpha = np.zeros((Npts), dtype=np.float)
    greenwhichc = spheric_line(radii,alpha,beta)
    ia=ea+0.
    ea=ia+np.pi
    da = (ea-ia)/(Npts-1)
    beta = np.arange(ia,ea+da,da)[0:Npts]
    greenwhichd = spheric_line(radii,alpha,beta)

    # Equator
    ia=np.pi-iazmrad/2.
    ea=ia+np.pi
    da = (ea-ia)/(Npts-1)
    alpha = np.arange(ia,ea+da,da)[0:Npts]
    beta = np.zeros((Npts), dtype=np.float)
    equatorc = spheric_line(radii,alpha,beta)
    ia=ea+0.
    ea=ia+np.pi
    da = (ea-ia)/(Npts-1)
    alpha = np.arange(ia,ea+da,da)[0:Npts]
    equatord = spheric_line(radii,alpha,beta)

    # 90 line
    if iazmrad > np.pi and iazmrad < 2.*np.pi:
        ia=3.*np.pi/2. + ielerad
    else:
        ia=np.pi/2. - ielerad
    if ielerad < 0.:
        ia = ia + np.pi
    ea=ia+np.pi
    da = (ea-ia)/(Npts-1)
    beta = np.arange(ia,ea+da,da)[0:Npts]
    alpha = np.ones((Npts), dtype=np.float)*np.pi/2.
    xclinec = spheric_line(radii,alpha,beta)
    ia=ea+0.
    ea=ia+np.pi
    da = (ea-ia)/(Npts-1)
    beta = np.arange(ia,ea+da,da)[0:Npts]
    xclined = spheric_line(radii,alpha,beta)

    # x line
    xline = np.zeros((2,3), dtype=np.float)
    xline[0,:] = position_sphere(radii, 0., 0.)
    xline[1,:] = position_sphere(radii, np.pi, 0.)

    # y line
    yline = np.zeros((2,3), dtype=np.float)
    yline[0,:] = position_sphere(radii, np.pi/2., 0.)
    yline[1,:] = position_sphere(radii, 3*np.pi/2., 0.)

    # z line
    zline = np.zeros((2,3), dtype=np.float)
    zline[0,:] = position_sphere(radii, 0., np.pi/2.)
    zline[1,:] = position_sphere(radii, 0., -np.pi/2.)

    fig = plt.figure()
    ax = fig.gca(projection='3d')

    # Sphere surface
    if drwsfc[0]:
        ax.plot_surface(sfcsphereu[0,:,:], sfcsphereu[1,:,:], sfcsphereu[2,:,:],     \
          color=colsfc[0])
    if drwsfc[1]:
        ax.plot_surface(sfcsphered[0,:,:], sfcsphered[1,:,:], sfcsphered[2,:,:],     \
          color=colsfc[1])

    # greenwhich
    linev = linegw[0]
    if drwgreeenwhich[0]:
        ax.plot(greenwhichc[:,0], greenwhichc[:,1], greenwhichc[:,2], linev[0],      \
          color=linev[1], linewidth=linev[2],  label='Greenwich')
    linev = linegw[1]
    if drwgreeenwhich[1]:
        ax.plot(greenwhichd[:,0], greenwhichd[:,1], greenwhichd[:,2], linev[0],      \
          color=linev[1], linewidth=linev[2])

    # Equator
    linev = lineeq[0]
    if drwequator[0]:
        ax.plot(equatorc[:,0], equatorc[:,1], equatorc[:,2], linev[0],               \
          color=linev[1], linewidth=linev[2], label='Equator')
    linev = lineeq[1]
    if drwequator[1]:
        ax.plot(equatord[:,0], equatord[:,1], equatord[:,2], linev[0],               \
          color=linev[1], linewidth=linev[2])

    # 90line
    linev = linexc[0]
    if drwxcline[0]:
        ax.plot(xclinec[:,0], xclinec[:,1], xclinec[:,2], linev[0], color=linev[1],  \
          linewidth=linev[2], label='90-line')
    linev = linexc[1]
    if drwxcline[1]:
        ax.plot(xclined[:,0], xclined[:,1], xclined[:,2], linev[0], color=linev[1],  \
          linewidth=linev[2])

    # x line
    linev = linex
    if drwxline:
        ax.plot([xline[0,0],xline[1,0]], [xline[0,1],xline[1,1]],                    \
          [xline[0,2],xline[1,2]], linev[0], color=linev[1], linewidth=linev[2],  label='xline')

    # y line
    linev = liney
    if drwyline:
        ax.plot([yline[0,0],yline[1,0]], [yline[0,1],yline[1,1]],                    \
          [yline[0,2],yline[1,2]], linev[0], color=linev[1], linewidth=linev[2],  label='yline')

    # z line
    linev = linez
    if drwzline:
        ax.plot([zline[0,0],zline[1,0]], [zline[0,1],zline[1,1]],                    \
          [zline[0,2],zline[1,2]], linev[0], color=linev[1], linewidth=linev[2],  label='zline')

    plt.legend()

    return fig, ax