#
# transform.py
#
# routines to transform polar coordinates and velocities to cartesian equivalents
# specifically implemented to be used by read.ogvar(), and therefore uses the 
# ordering f(nz,ny,nx), not f(nx,ny,nz)
#
# Author:
# J. Aarnes (jorgenaarnes@gmail.com)
import numpy as np

def pospolar2cart(r,th):
    """
    Transform from polar to cartesian coordinates
    """
    rsize=r.size
    thsize=th.size
    
    x = np.empty([thsize+1,rsize])
    y = np.empty([thsize+1,rsize])
    for i in range(rsize):
        for j in range(thsize):
            x[j,i] = r[i]*np.cos(th[j])
            y[j,i] = r[i]*np.sin(th[j])
        x[thsize,i] = r[i]*np.cos(th[0])
        y[thsize,i] = r[i]*np.sin(th[0])

    return x,y

def velpolar2cart(vr,vth,r,th,zcoord=0):
    """
    Transform from polar velocities to cartesian velocities

    call signature:

    velpolar2cart(vr,vth,r,th,zcoord=0)

    Keyword arguments:

    *vr*:
      3D array of velocities in radial direction

    *vth*:
      3D array of velocities in theta direction

    *r*
      coordinates in radial direction

    *th*
      coordinates in theta direction

    *zcoord*
      plane in z-dir
    """

    rsize=r.size
    thsize=th.size
    vrsize=vr.size
    vthsize=vth.size
    if(vrsize!=vthsize):
        print('ERROR: vr and vth not equal in size!')
        return [],[]
    vx = np.empty([thsize+1,rsize])
    vy = np.empty([thsize+1,rsize])
    for i in range(rsize):
        for j in range(thsize):
            vx[j,i] =  vr[zcoord,j,i]*np.cos(th[j])-vth[zcoord,j,i]*np.sin(th[j])
            vy[j,i] =  vr[zcoord,j,i]*np.sin(th[j])+vth[zcoord,j,i]*np.cos(th[j])
        vx[thsize,i] = vr[zcoord,0,i]*np.cos(th[0])-vth[zcoord,0,i]*np.sin(th[0])
        vy[thsize,i] = vr[zcoord,0,i]*np.sin(th[0])+vth[zcoord,0,i]*np.cos(th[0])

    return vx,vy