import sys
import timeit

import numpy as np
from tqdm import tqdm

from turbx import gradient
from turbx.grid_metric import get_metric_tensor_2d

from .utils import even_print, format_time_string

'''
========================================================================
Mass & Momentum Balance Tools
========================================================================
'''

# ======================================================================
# Mass
# ======================================================================

@staticmethod
def mass_balance_kernel_cartesian(x,y, u_Fv,v_Fv,rho, transformation_matrix=None, acc=6, edge_stencil='full'):
    '''
    Get Cartesian [x,y] mass balance terms in 2D plane
    '''
    
    nx,ny = x.shape
    
    ## Rotate & translate [x,y] and rotate vector quantities (global Cartesian)
    if transformation_matrix is not None:
        
        ## Rotate & translate [x,y]
        xy2d_tmp = np.stack((x, y, np.ones_like(x)), axis=-1)
        xy2d_tmp = np.einsum( 'ij,xyj->xyi', transformation_matrix , xy2d_tmp )
        x        = xy2d_tmp[:,:,0]
        y        = xy2d_tmp[:,:,1]
        xy2d_tmp = None; del xy2d_tmp
        
        ## Make a copy of the [2x2] matrix
        rotation_matrix_2x2 = np.copy(transformation_matrix[:2,:2])
        
        ## Rotate [u_Fv,v_Fv]
        uv2d_Fv = np.stack((u_Fv,v_Fv), axis=-1)
        uv2d_Fv = np.einsum('ij,xyj->xyi', rotation_matrix_2x2, uv2d_Fv)
        u_Fv    = uv2d_Fv[:,:,0]
        v_Fv    = uv2d_Fv[:,:,1]
        uv2d_Fv = None; del uv2d_Fv
    
    # ==================================================================
    # Get metric tensor
    # M = δ[q1,q2]/δ[x1,x2] : δ[q1,q2]/δ[x,y] or δ[q1,q2]/δ[r,θ]
    # ==================================================================
    
    M = get_metric_tensor_2d(x, y, acc=acc, edge_stencil=edge_stencil, verbose=False)
    ddx1_q1 = M[:,:,0,0] ## ∂q1/∂x1 --> ∂q1/∂x or ∂q1/∂r
    ddx1_q2 = M[:,:,1,0] ## ∂q2/∂x1 --> ∂q2/∂x or ∂q2/∂r
    ddx2_q1 = M[:,:,0,1] ## ∂q1/∂x2 --> ∂q1/∂y or ∂q1/∂θ
    ddx2_q2 = M[:,:,1,1] ## ∂q2/∂x2 --> ∂q2/∂y or ∂q2/∂θ
    M = None ; del M
    
    # ==================================================================
    # Construct mass balance : Cartesian [x,y]
    # ==================================================================
    
    # === ∂x[ρ·u]
    
    A      = rho * u_Fv
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx1_A = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    #ddx2_A = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t1 = ddx1_A
    
    # === ∂y[ρ·v]
    
    A      = rho * v_Fv
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    #ddx1_A = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddx2_A = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t2 = ddx2_A
    
    # ===
    
    data = {
        't1':t1,
        't2':t2,
        }
    
    return data

@staticmethod
def mass_balance_kernel_cylindrical(x,y, u_Fv,v_Fv,rho, transformation_matrix=None, acc=6, edge_stencil='full'):
    '''
    Get cylindrical [r,θ] mass balance terms in 2D plane
    '''
    
    nx,ny = x.shape
    
    ## Rotate & translate [x,y] and rotate vector quantities (global Cartesian)
    if transformation_matrix is not None:
        
        ## Rotate & translate [x,y]
        xy2d_tmp = np.stack((x, y, np.ones_like(x)), axis=-1)
        xy2d_tmp = np.einsum( 'ij,xyj->xyi', transformation_matrix , xy2d_tmp )
        x        = xy2d_tmp[:,:,0]
        y        = xy2d_tmp[:,:,1]
        xy2d_tmp = None; del xy2d_tmp
        
        ## Make a copy of the [2x2] matrix
        rotation_matrix_2x2 = np.copy(transformation_matrix[:2,:2])
        
        ## Rotate [u_Fv,v_Fv]
        uv2d_Fv = np.stack((u_Fv,v_Fv), axis=-1)
        uv2d_Fv = np.einsum('ij,xyj->xyi', rotation_matrix_2x2, uv2d_Fv)
        u_Fv    = uv2d_Fv[:,:,0]
        v_Fv    = uv2d_Fv[:,:,1]
        uv2d_Fv = None; del uv2d_Fv
    
    ## Calculate coordinate arrays [r,θ] from [x,y]
    r        = np.sqrt(x**2 + y**2)
    theta    = np.arctan2(y,x)
    costheta = np.cos(theta)
    sintheta = np.sin(theta)
    
    eps = np.finfo(np.float64).eps
    if np.any(r<eps):
        raise RuntimeError(f'r has values < {eps:0.3e}')
    
    ## Calculate [ur_Fv,uθ_Fv] from [u_Fv,v_Fv]
    ur_Fv =  u_Fv*costheta + v_Fv*sintheta
    ut_Fv = -u_Fv*sintheta + v_Fv*costheta
    
    # ==================================================================
    # Get metric tensor
    # M = δ[q1,q2]/δ[x1,x2] : δ[q1,q2]/δ[x,y] or δ[q1,q2]/δ[r,θ]
    # ==================================================================
    
    M = get_metric_tensor_2d(r, theta, acc=acc, edge_stencil=edge_stencil, verbose=False)
    ddx1_q1 = M[:,:,0,0] ## ∂q1/∂x1 --> ∂q1/∂x or ∂q1/∂r
    ddx1_q2 = M[:,:,1,0] ## ∂q2/∂x1 --> ∂q2/∂x or ∂q2/∂r
    ddx2_q1 = M[:,:,0,1] ## ∂q1/∂x2 --> ∂q1/∂y or ∂q1/∂θ
    ddx2_q2 = M[:,:,1,1] ## ∂q2/∂x2 --> ∂q2/∂y or ∂q2/∂θ
    M = None ; del M
    
    # ==================================================================
    # Construct mass balance : cylindrical [r,θ]
    # ==================================================================
    
    ## ∇·(ρu) = ∂r[ρ·ur] + (1/r)∂θ[ρ·uθ] + (1/r)[ρ·ur] = 0
    
    # === ∂r[ρ·ur]
    
    A = rho * ur_Fv
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx1_A = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    #ddx2_A = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t1 = ddx1_A
    
    # === (1/r)·∂θ[ρ·uθ]
    
    A = rho * ut_Fv
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    #ddx1_A = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddx2_A = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t2 = (1/r) * ddx2_A
    
    # === (1/r)·[ρ·ur]
    
    t3 = (1/r) * rho * ur_Fv
    
    # ===
    
    data = {
        't1':t1,
        't2':t2,
        't3':t3,
        }
    
    return data

def _calc_mass_balance_terms(self, **kwargs):
    '''
    Calculate terms of time-averaged mass conservation equation
    '''
    
    verbose      = kwargs.get('verbose',True)
    local_csys   = kwargs.get('local_csys',True)
    acc          = kwargs.get('acc',6)
    edge_stencil = kwargs.get('edge_stencil','full')
    xis          = kwargs.get('xis',None) ## which [x] indices to compute (None=all)
    csys_type    = kwargs.get('csys_type','cartesian')
    
    if verbose: print('\n'+'ztmd.calc_mass_balance_terms()'+'\n'+72*'-')
    t_start_func = timeit.default_timer()
    
    if self.open_mode not in ['a','w','r+']:
        raise ValueError('Not able to write to hdf5 file')
    
    if verbose: even_print('local_csys',str(local_csys))
    if verbose: even_print('csys_type',csys_type)
    
    if local_csys:
        if self.rectilinear:
            print('>>> WARNING: this grid is rectilinear. Using local csys is not needed.')
        if 'csys/vtang' not in self:
            raise ValueError('csys/vtang not in hdf5')
        if 'csys/vnorm' not in self:
            raise ValueError('csys/vnorm not in hdf5')
        
        ## Read unit vectors (wall tangent, wall norm) from HDF5
        vtang = np.copy( self['csys/vtang'][()] )
        vnorm = np.copy( self['csys/vnorm'][()] )
        if verbose: even_print('vtang',str(vtang.shape))
        if verbose: even_print('vnorm',str(vnorm.shape))
    
    if (csys_type=='cartesian'):
        gn = 'data_mass_xy'
        gns = [gn,]
    elif (csys_type=='cylindrical'):
        gnr = 'data_mass_rtheta'  ## keep single group for mass in cyl
        gns = [gnr,]
    else:
        raise NotImplementedError
    
    if (csys_type=='cartesian'):
        terms = {
            't1':(r'∂x[ρ·u]', r'$\partial_{x}[\bar{\rho} \tilde{u}]$'),
            't2':(r'∂y[ρ·v]', r'$\partial_{y}[\bar{\rho} \tilde{v}]$'),
            }
    elif (csys_type=='cylindrical'):
        terms = {
            't1':(r'∂r[ρ·ur]'       , r'$\partial_{r}[ \bar{\rho} \tilde{u}_r ]$'),
            't2':(r'(1/r)·∂θ[ρ·uθ]' , r'$\frac{1}{r}\partial_{\theta}[ \bar{\rho} \tilde{u}_{\theta} ]$'),
            't3':(r'(1/r)·[ρ·ur]'   , r'$\frac{1}{r}[ \bar{\rho} \tilde{u}_r ]$'),
            }
    else:
        raise ValueError
    
    ## Create groups in ZTMD
    for gn in gns:
        if verbose: print(72*'-')
        if gn in self:
            del self[gn]
            if verbose: even_print('deleted group',f'/{gn}/')
        
        grp = self.create_group(gn)
        if verbose: even_print('created group',f'{grp.name}/')
        
        ## Set group attributes
        grp.attrs['local_csys'] = local_csys
        if verbose: even_print(f'{gn}.local_csys',str(grp.attrs['local_csys']))
    
    gn=None; del gn
    
    if verbose: print(72*'-')
    
    ## Initialize dsets (Cartesian)
    if (csys_type=='cartesian'):
        for dsn, tags in terms.items():
            unicode,latex = tags
            shape = (self.ny,self.nx)
            if (f'{gns[0]}/{dsn}' in self):
                del self[f'{gns[0]}/{dsn}']
            dset = self.create_dataset(
                                    f'{gns[0]}/{dsn}', 
                                    shape=shape, 
                                    dtype=np.float32,
                                    chunks=None,
                                    fillvalue=np.float32(np.nan),
                                    )
            if verbose: even_print(f'{gns[0]}/{dsn}',str(shape))
            dset.attrs['latex'] = latex
            dset.attrs['unicode'] = unicode
        if verbose: print(72*'-')
    
    ## Initialize dsets (cylindrical)
    if (csys_type=='cylindrical'):
        for dsn, tags in terms.items():
            unicode,latex = tags
            shape = (self.ny,self.nx)
            if (f'{gnr}/{dsn}' in self):
                del self[f'{gnr}/{dsn}']
            dset = self.create_dataset(
                                    f'{gnr}/{dsn}', 
                                    shape=shape, 
                                    dtype=np.float32,
                                    chunks=None,
                                    fillvalue=np.float32(np.nan),
                                    )
            if verbose: even_print(f'{gnr}/{dsn}',str(shape))
            dset.attrs['latex'] = latex
            dset.attrs['unicode'] = unicode
        if verbose: print(72*'-')
    
    ## Read data
    if csys_type=='cylindrical' and hasattr(self,'crv_R'):
        crv_R = np.copy(self.crv_R)
        if verbose:
            even_print('crv_R',str(crv_R.shape))
            print(72*'-')
    
    if verbose:
        print('Reading data'+'\n'+'------------')
    #u   = np.copy( self['data/u'][()].T   )
    #v   = np.copy( self['data/v'][()].T   )
    rho = np.copy( self['data/rho'][()].T )
    
    #if verbose: even_print('u',str(u.shape))
    #if verbose: even_print('v',str(v.shape))
    if verbose: even_print('ρ',str(rho.shape))
    
    u_Fv = np.copy( self['data/u_Fv'][()].T )
    v_Fv = np.copy( self['data/v_Fv'][()].T )
    if verbose: even_print('u_Fv',str(u_Fv.shape))
    if verbose: even_print('v_Fv',str(v_Fv.shape))
    
    x = np.copy(self.x)
    y = np.copy(self.y)
    
    if x.ndim==1 and y.ndim==1:
        x,y = np.meshgrid(x,y, indexing='ij')
    if verbose: even_print('x',str(x.shape))
    if verbose: even_print('y',str(y.shape))
    if verbose: print(72*'-')
    
    nx,ny = x.shape
    
    # ==================================================================
    
    def _get_transformation_matrix(xi):
        '''
        Get transformation matrix : translate,rotate,[translate]
        '''
        rotation_angle_z = np.arcsin(vnorm[xi,0,0])
        
        rotation_matrix = np.array( [[ np.cos(rotation_angle_z) , -np.sin(rotation_angle_z) , 0 ],
                                     [ np.sin(rotation_angle_z) ,  np.cos(rotation_angle_z) , 0 ],
                                     [ 0                        ,  0                        , 1 ]], dtype=np.float64 )
        
        ## Translate to [cx,cy]
        cx = x[xi,0]
        cy = y[xi,0]
        translation_matrix_A = np.array( [[ 1 , 0 , -cx ],
                                          [ 0 , 1 , -cy ],
                                          [ 0 , 0 ,  1  ]], dtype=np.float64 )
        
        ## For cylindrical csys, translate in [y] by -R (the local curvature radius)
        ##   such that origin is distance 'R' from the wall
        if (csys_type=='cylindrical'):
            cx = 0.
            try:
                cy = crv_R[xi]
            except NameError:
                cy = 0.
            translation_matrix_B = np.array( [[ 1 , 0 ,  0  ],
                                              [ 0 , 1 , -cy ],
                                              [ 0 , 0 ,  1  ]], dtype=np.float64 )
        else:
            translation_matrix_B = np.identity(3, dtype=np.float64) ## No second translation
        
        ## Form the full [3x3] matrix by multiplying together
        ##  --> ORDER MATTERS!
        transformation_matrix = np.einsum( 'ij,jk,kl->il', translation_matrix_B , rotation_matrix, translation_matrix_A )
        
        return transformation_matrix
    
    # ==================================================================
    
    if local_csys: ## Compute mass balance in local csys
        
        if xis is None: ## Do all
            xis = np.arange(self.nx)
            n_prog = self.nx
        else: ## Do only for selected x-indices
            n_prog = len(xis)
        
        ## Assert passed 'xis' is OK
        if isinstance(xis,np.ndarray): ## numpy array
            if not xis.ndim==1:
                raise ValueError('xis must have .ndim==1')
            if not np.issubdtype(xis.dtype,np.integer):
                raise RuntimeError("xis must have subtype np.integer")
        else: ## some other iterable... are all elements ints?
            try:
                xis = [int(xi) for xi in xis]
            except Exception as e:
                raise RuntimeError("elements of xis are not interpretable as ints") from e
        
        # ==============================================================
        # Main loop
        # ==============================================================
        
        if csys_type=='cartesian':
            desc='Mass balance [x,y]'
        elif csys_type=='cylindrical':
            desc='Mass balance [r,θ]'
        else:
            raise ValueError
        
        if verbose:
            progress_bar = tqdm(
                total=n_prog,
                ncols=100,
                desc=desc,
                leave=True,
                file=sys.stdout,
                mininterval=0.1,
                smoothing=0.,
                bar_format="{l_bar}{bar}| {n}/{total} [{percentage:.1f}%] {elapsed}/{remaining}",
                ascii="░█",
                colour='#00ffff',
                )
        
        for xi in xis:
            
            ## Get transformation matrix (global-->local) for this [x] index
            transformation_matrix = _get_transformation_matrix(xi)
            
            ## Take N points in [x] on either side
            ## This allows for full FD stencil at central point
            xiA = max( 0  , xi-12   )
            xiB = min( nx , xi+12+1 )
            xii = xi - xiA ## Index local to slice
            
            if (csys_type=='cartesian'):
                
                data = mass_balance_kernel_cartesian(
                                                x=x[xiA:xiB,:],
                                                y=y[xiA:xiB,:],
                                                u_Fv=u_Fv[xiA:xiB,:],
                                                v_Fv=v_Fv[xiA:xiB,:],
                                                rho=rho[xiA:xiB,:],
                                                transformation_matrix=transformation_matrix,
                                                acc=acc,
                                                edge_stencil=edge_stencil,
                                                )
                
                ## Write to HDF5
                for dsn, arr in data.items():
                    dset = self[f'{gns[0]}/{dsn}']
                    dset[:,xi] = arr[xii,:] ## Write only 1 [x] position
            
            elif (csys_type=='cylindrical'):
                
                data = mass_balance_kernel_cylindrical(
                                                x=x[xiA:xiB,:],
                                                y=y[xiA:xiB,:],
                                                u_Fv=u_Fv[xiA:xiB,:],
                                                v_Fv=v_Fv[xiA:xiB,:],
                                                rho=rho[xiA:xiB,:],
                                                transformation_matrix=transformation_matrix,
                                                acc=acc,
                                                edge_stencil=edge_stencil,
                                                )
                
                ## Write to HDF5
                for dsn, arr in data.items():
                    dset = self[f'{gnr}/{dsn}']
                    dset[:,xi] = arr[xii,:] ## Write only 1 [x] position
            
            else:
                raise ValueError
            
            if verbose: progress_bar.update()
        if verbose: progress_bar.close()
    
    else: ## Compute mass balance in global csys
        
        ## Send to kernel
        if csys_type=='cartesian':
            
            data = mass_balance_kernel_cartesian(
                                            x=x,
                                            y=y,
                                            u_Fv=u_Fv,
                                            v_Fv=v_Fv,
                                            rho=rho,
                                            #transformation_matrix=transformation_matrix,
                                            acc=acc,
                                            edge_stencil=edge_stencil,
                                            )
            
            ## Write to HDF5
            for dsn, arr in data.items():
                dset = self[f'{gns[0]}/{dsn}']
                if verbose:
                    even_print(dset.name,str(arr.T.shape))
                dset[:,:] = arr.T
        
        elif csys_type=='cylindrical':
            
            data = mass_balance_kernel_cylindrical(
                                            x=x,
                                            y=y,
                                            u_Fv=u_Fv,
                                            v_Fv=v_Fv,
                                            rho=rho,
                                            #transformation_matrix=transformation_matrix,
                                            acc=acc,
                                            edge_stencil=edge_stencil,
                                            )
            
            ## Write to HDF5
            for dsn, arr in data.items():
                dset = self[f'{gnr}/{dsn}']
                if verbose:
                    even_print(dset.name,str(arr.T.shape))
                dset[:,:] = arr.T
        
        else:
            raise ValueError
    
    self.get_header(verbose=False)
    if verbose: print(72*'-')
    if verbose: print('total time : ztmd.calc_mass_balance_terms() : %s'%format_time_string((timeit.default_timer() - t_start_func)))
    if verbose: print(72*'-')
    return

# ======================================================================
# Momentum
# ======================================================================

@staticmethod
def momentum_balance_kernel_cartesian(x,y, u,v, u_Fv,v_Fv, p,rho,mu, r_uII_uII,r_uII_vII,r_vII_vII, transformation_matrix=None, acc=6, edge_stencil='full'):
    '''
    Get Cartesian [x,y] momentum balance terms in 2D plane
    '''
    
    nx,ny = x.shape
    
    ## Rotate & translate [x,y] and rotate vector/tensor quantities (global Cartesian)
    if transformation_matrix is not None:
        
        ## Rotate & translate [x,y]
        xy2d_tmp = np.stack((x, y, np.ones_like(x)), axis=-1)
        xy2d_tmp = np.einsum( 'ij,xyj->xyi', transformation_matrix , xy2d_tmp )
        x        = xy2d_tmp[:,:,0]
        y        = xy2d_tmp[:,:,1]
        xy2d_tmp = None; del xy2d_tmp
        
        ## Make a copy of the [2x2] matrix
        rotation_matrix_2x2 = np.copy(transformation_matrix[:2,:2])
        
        ## Rotate [u,v]
        uv2d = np.stack((u,v), axis=-1)
        uv2d = np.einsum('ij,xyj->xyi' , rotation_matrix_2x2 , uv2d)
        u    = uv2d[:,:,0]
        v    = uv2d[:,:,1]
        uv2d = None; del uv2d
        
        ## Rotate [u_Fv,v_Fv]
        uv2d_Fv = np.stack((u_Fv,v_Fv), axis=-1)
        uv2d_Fv = np.einsum('ij,xyj->xyi', rotation_matrix_2x2, uv2d_Fv)
        u_Fv    = uv2d_Fv[:,:,0]
        v_Fv    = uv2d_Fv[:,:,1]
        uv2d_Fv = None; del uv2d_Fv
        
        # === Turbulent stress tensor: global Cartesian rotate
        
        r_uiIIujII = np.zeros((nx,ny,2,2), dtype=np.float64)
        r_uiIIujII[:,:,0,0] = r_uII_uII
        r_uiIIujII[:,:,0,1] = r_uII_vII
        r_uiIIujII[:,:,1,0] = r_uII_vII
        r_uiIIujII[:,:,1,1] = r_vII_vII
        
        ## Assert symmetry
        np.testing.assert_allclose(r_uiIIujII, np.transpose(r_uiIIujII, (0,1,3,2)), atol=1e-8, rtol=1e-8)
        
        ## Tensor rotation : [A'] = [R][A][R]^T
        r_uiIIujII = np.copy( np.einsum( 'ij,xyjk,kl->xyil' , rotation_matrix_2x2 , r_uiIIujII , rotation_matrix_2x2.T ) )
        
        r_uII_uII = r_uiIIujII[:,:,0,0]
        r_uII_vII = r_uiIIujII[:,:,0,1]
        r_vII_uII = r_uiIIujII[:,:,1,0]
        r_vII_vII = r_uiIIujII[:,:,1,1]
        
        r_uiIIujII = None ; del r_uiIIujII
        
        np.testing.assert_allclose(r_uII_vII, r_vII_uII, atol=1e-7, rtol=1e-7)
        r_vII_uII = None ; del r_vII_uII
    
    # ==================================================================
    # Get metric tensor
    # M = δ[q1,q2]/δ[x1,x2] : δ[q1,q2]/δ[x,y] or δ[q1,q2]/δ[r,θ]
    # ==================================================================
    
    M = get_metric_tensor_2d(x, y, acc=acc, edge_stencil=edge_stencil, verbose=False)
    ddx1_q1 = np.copy( M[:,:,0,0] ) ## ∂q1/∂x1 --> ∂q1/∂x or ∂q1/∂r
    ddx1_q2 = np.copy( M[:,:,1,0] ) ## ∂q2/∂x1 --> ∂q2/∂x or ∂q2/∂r
    ddx2_q1 = np.copy( M[:,:,0,1] ) ## ∂q1/∂x2 --> ∂q1/∂y or ∂q1/∂θ
    ddx2_q2 = np.copy( M[:,:,1,1] ) ## ∂q2/∂x2 --> ∂q2/∂y or ∂q2/∂θ
    M = None ; del M
    
    # ==================================================================
    # Calculate gradients : Cartesian [x,y]
    # ==================================================================
    
    ## [x1,x2] = [x,y]
    ## ∂(A)/∂x1 = ∂(A)/∂(q1)·∂(q1)/∂(x1) + ∂(A)/∂(q2)·∂(q2)/∂(x1)
    ## ∂(A)/∂x2 = ∂(A)/∂(q1)·∂(q1)/∂(x2) + ∂(A)/∂(q2)·∂(q2)/∂(x2)
    
    A = u
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx_u  = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddy_u  = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    A = v
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx_v  = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddy_v  = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    A = p
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx_p  = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddy_p  = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    # ==================================================================
    # Construct momentum balance : Cartesian [x,y]
    # ==================================================================
    
    ## τ_xx, τ_yy, τ_xy
    tau_xx = +(4/3)*mu*ddx_u - (2/3)*mu*ddy_v
    tau_yy = -(2/3)*mu*ddx_u + (4/3)*mu*ddy_v
    tau_xy = mu*(ddx_v + ddy_u)
    
    # === ∂x[ρ·u·u]
    
    A      = rho * u_Fv * u_Fv
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx_A  = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    #ddy_A  = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t1x = ddx_A
    
    # === ∂x[ρ·u·v] & ∂y[ρ·u·v]
    
    A      = rho * u_Fv * v_Fv
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx_A  = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddy_A  = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t1y = ddx_A
    t2x = ddy_A
    
    # === ∂y[ρ·v·v]
    
    A      = rho * v_Fv * v_Fv
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    #ddx_A  = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddy_A  = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t2y = ddy_A
    
    # === ∂x[ρu″u″]
    
    A      = r_uII_uII
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx_A  = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    #ddy_A  = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t5x = ddx_A
    
    # === ∂x[ρu″v″] & ∂y[ρu″v″]
    
    A      = r_uII_vII
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx_A  = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddy_A  = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t5y = ddx_A
    t7x = ddy_A
    
    # === ∂y[ρv″v″]
    
    A      = r_vII_vII
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    #ddx_A  = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddy_A  = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t7y = ddy_A
    
    # === ∂x[p] & ∂y[p]
    
    t3x = ddx_p
    t3y = ddy_p
    
    # === -∂x[τxx]
    
    ddq1_A = gradient(tau_xx, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(tau_xx, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx_A  = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    #ddy_A  = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t4x = -1 * ddx_A
    
    # === -∂x[τxy] & -∂y[τxy]
    
    ddq1_A = gradient(tau_xy, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(tau_xy, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx_A  = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddy_A  = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t4y = -1 * ddx_A
    t6x = -1 * ddy_A
    
    # === -∂y[τyy]
    
    ddq1_A = gradient(tau_yy, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(tau_yy, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    #ddx_A  = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddy_A  = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t6y = -1 * ddy_A
    
    # ===
    
    data_x = {
        't1':t1x,
        't2':t2x,
        't3':t3x,
        't4':t4x,
        't5':t5x,
        't6':t6x,
        't7':t7x,
        }
    
    data_y = {
        't1':t1y,
        't2':t2y,
        't3':t3y,
        't4':t4y,
        't5':t5y,
        't6':t6y,
        't7':t7y,
        }
    
    return data_x, data_y

@staticmethod
def momentum_balance_kernel_cylindrical(x,y, u,v, u_Fv,v_Fv, p,rho,mu, r_uII_uII,r_uII_vII,r_vII_vII, transformation_matrix=None, acc=6, edge_stencil='full'):
    '''
    Get cylindrical [r,θ] momentum balance terms in 2D plane
    '''
    
    nx,ny = x.shape
    
    ## Rotate & translate [x,y] and rotate vector/tensor quantities (global Cartesian)
    if transformation_matrix is not None:
        
        ## Rotate & translate [x,y]
        xy2d_tmp = np.stack((x, y, np.ones_like(x)), axis=-1)
        xy2d_tmp = np.einsum( 'ij,xyj->xyi', transformation_matrix , xy2d_tmp )
        x        = xy2d_tmp[:,:,0]
        y        = xy2d_tmp[:,:,1]
        xy2d_tmp = None; del xy2d_tmp
        
        ## Make a copy of the [2x2] matrix
        rotation_matrix_2x2 = np.copy(transformation_matrix[:2,:2])
        
        ## Rotate [u,v]
        uv2d = np.stack((u,v), axis=-1)
        uv2d = np.einsum('ij,xyj->xyi' , rotation_matrix_2x2 , uv2d)
        u    = uv2d[:,:,0]
        v    = uv2d[:,:,1]
        uv2d = None; del uv2d
        
        ## Rotate [u_Fv,v_Fv]
        uv2d_Fv = np.stack((u_Fv,v_Fv), axis=-1)
        uv2d_Fv = np.einsum('ij,xyj->xyi', rotation_matrix_2x2, uv2d_Fv)
        u_Fv    = uv2d_Fv[:,:,0]
        v_Fv    = uv2d_Fv[:,:,1]
        uv2d_Fv = None; del uv2d_Fv
        
        # === Turbulent stress tensor: global Cartesian rotate
        
        # a = rotation_matrix_2x2[0,0]
        # b = rotation_matrix_2x2[0,1]
        # c = rotation_matrix_2x2[1,0]
        # d = rotation_matrix_2x2[1,1]
        # 
        # r_uII_uII_new = a*a*r_uII_uII + 2*a*b*r_uII_vII + b*b*r_vII_vII
        # r_uII_vII_new = a*c*r_uII_uII + (a*d + b*c)*r_uII_vII + b*d*r_vII_vII
        # r_vII_vII_new = c*c*r_uII_uII + 2*c*d*r_uII_vII + d*d*r_vII_vII
        # 
        # r_uII_uII = r_uII_uII_new
        # r_uII_vII = r_uII_vII_new
        # r_vII_vII = r_vII_vII_new
        # 
        # r_uII_uII_new = None; del r_uII_uII_new
        # r_uII_vII_new = None; del r_uII_vII_new
        # r_vII_vII_new = None; del r_vII_vII_new
        
        # === Turbulent stress tensor: global Cartesian rotate
        
        r_uiIIujII = np.zeros((nx,ny,2,2), dtype=np.float64)
        r_uiIIujII[:,:,0,0] = np.copy( r_uII_uII )
        r_uiIIujII[:,:,0,1] = np.copy( r_uII_vII )
        r_uiIIujII[:,:,1,0] = np.copy( r_uII_vII )
        r_uiIIujII[:,:,1,1] = np.copy( r_vII_vII )
        
        ## Assert symmetry
        np.testing.assert_allclose(r_uiIIujII, np.transpose(r_uiIIujII, (0,1,3,2)), atol=1e-8, rtol=1e-8)
        
        ## Tensor rotation : [A'] = [R][A][R]^T
        r_uiIIujII = np.copy( np.einsum( 'ij,xyjk,kl->xyil' , rotation_matrix_2x2 , r_uiIIujII , rotation_matrix_2x2.T ) )
        
        r_uII_uII = np.copy(r_uiIIujII[:,:,0,0])
        r_uII_vII = np.copy(r_uiIIujII[:,:,0,1])
        r_vII_uII = np.copy(r_uiIIujII[:,:,1,0])
        r_vII_vII = np.copy(r_uiIIujII[:,:,1,1])
        
        r_uiIIujII = None ; del r_uiIIujII
        
        np.testing.assert_allclose(r_uII_vII, r_vII_uII, atol=1e-7, rtol=1e-7)
        r_vII_uII = None ; del r_vII_uII
    
    ## Calculate coordinate arrays [r,θ] from [x,y]
    r        = np.sqrt(x**2 + y**2)
    theta    = np.arctan2(y,x)
    costheta = np.cos(theta)
    sintheta = np.sin(theta)
    
    eps = np.finfo(np.float64).eps
    if np.any(r<eps):
        raise RuntimeError(f'r has values < {eps:0.3e}')
    
    ## Calculate [ur,uθ] from [u,v]
    ur =  u*costheta + v*sintheta
    ut = -u*sintheta + v*costheta
    
    ## Calculate [ur_Fv,uθ_Fv] from [u_Fv,v_Fv]
    ur_Fv =  u_Fv*costheta + v_Fv*sintheta
    ut_Fv = -u_Fv*sintheta + v_Fv*costheta
    
    ## Re-Stress tensor : convert from Cartesian (rotated) to cylindrical components
    ## Tensor transform for symmetric 2x2: [R][T][R]^T, with R = [[c,s],[-s,c]]
    ## Components:
    ## rr = c^2*uu + 2cs*uv + s^2*vv
    ## rt = -cs*uu + (c^2 - s^2)*uv + cs*vv
    ## tt = s^2*uu - 2cs*uv + c^2*vv
    
    c2  = costheta*costheta
    s2  = sintheta*sintheta
    cs  = costheta*sintheta
    c2ms2 = c2 - s2
    
    r_urII_urII =  c2*r_uII_uII +  2*cs*r_uII_vII + s2*r_vII_vII
    r_urII_utII = -cs*r_uII_uII + c2ms2*r_uII_vII + cs*r_vII_vII
    r_utII_urII = r_urII_utII
    r_utII_utII = s2*r_uII_uII - 2*cs*r_uII_vII + c2*r_vII_vII
    
    ## Assert symmetry
    np.testing.assert_allclose(r_urII_utII, r_utII_urII, atol=1e-7, rtol=1e-7)
    r_utII_urII = None ; del r_utII_urII
    
    # ==================================================================
    # Get metric tensor
    # M = δ[q1,q2]/δ[x1,x2] : δ[q1,q2]/δ[x,y] or δ[q1,q2]/δ[r,θ]
    # ==================================================================
    
    M = get_metric_tensor_2d(r, theta, acc=acc, edge_stencil=edge_stencil, verbose=False)
    ddx1_q1 = M[:,:,0,0] ## ∂q1/∂x1 --> ∂q1/∂x or ∂q1/∂r
    ddx1_q2 = M[:,:,1,0] ## ∂q2/∂x1 --> ∂q2/∂x or ∂q2/∂r
    ddx2_q1 = M[:,:,0,1] ## ∂q1/∂x2 --> ∂q1/∂y or ∂q1/∂θ
    ddx2_q2 = M[:,:,1,1] ## ∂q2/∂x2 --> ∂q2/∂y or ∂q2/∂θ
    M = None ; del M
    
    # ==================================================================
    # Calculate gradients : cylindrical [r,θ]
    # ==================================================================
    
    ## [x1,x2] = [r,θ]
    ## ∂(A)/∂x1 = ∂(A)/∂(q1)·∂(q1)/∂(x1) + ∂(A)/∂(q2)·∂(q2)/∂(x1)
    ## ∂(A)/∂x2 = ∂(A)/∂(q1)·∂(q1)/∂(x2) + ∂(A)/∂(q2)·∂(q2)/∂(x2)
    
    A = ur
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddr_ur = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddt_ur = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    A = ut
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddr_ut = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddt_ut = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    A = p
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddr_p = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddt_p = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    # ==================================================================
    # Construct momentum balance : cylindrical [r,θ]
    # ==================================================================
    
    ## τ_rr, τ_θθ, τ_rθ
    tau_rr = +(4/3)*mu*ddr_ur - (2/3)*(1/r)*mu*ddt_ut - (2/3)*mu*(ur/r)
    tau_tt = -(2/3)*mu*ddr_ur + (4/3)*(1/r)*mu*ddt_ut + (4/3)*mu*(ur/r)
    tau_rt = mu*( (1/r)*ddt_ur + ddr_ut - (ut/r) )
    
    # === ∂r[ρ·ur·ur]
    
    A = rho * ur_Fv * ur_Fv
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx1_A = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    #ddx2_A = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t1r = ddx1_A
    
    # === (1/r)·∂θ[ρ·ur·uθ]
    
    A = rho * ur_Fv * ut_Fv
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    #ddx1_A = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddx2_A = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t2r = (1/r)*ddx2_A
    
    # === ∂r[ρ·ur·uθ]
    
    A = rho * ur_Fv * ut_Fv
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx1_A = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    #ddx2_A = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t1t = ddx1_A
    
    # === (1/r)·∂θ[ρ·uθ·uθ]
    
    A = rho * ut_Fv * ut_Fv
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    #ddx1_A = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddx2_A = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t2t = (1/r)*ddx2_A
    
    # === ∂r[p] & (1/r)·∂θ[p]
    
    t3r = ddr_p
    t3t = (1/r)*ddt_p
    
    # === -∂r[τ_rr]
    
    A = tau_rr
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx1_A = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    #ddx2_A = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t4r = -1 * ddx1_A
    
    # === -(1/r)·∂θ[τ_θθ]
    
    A = tau_tt
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    #ddx1_A = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddx2_A = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t6t = -(1/r)*ddx2_A
    
    # === ∂r[ρur″ur″]
    
    A = r_urII_urII
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx1_A = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    #ddx2_A = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t5r = ddx1_A
    
    # === -(1/r)·∂θ[τ_rθ]  &  -∂r[τ_rθ]
    
    A = tau_rt
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx1_A = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddx2_A = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t4t = -ddx1_A
    t6r = -(1/r)*ddx2_A
    
    # === (1/r)·∂θ[ρur″uθ″] & ∂r[ρur″uθ″]
    
    A = r_urII_utII
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    ddx1_A = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddx2_A = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t7r = (1/r)*ddx2_A
    t5t = ddx1_A
    
    # === (1/r)·∂θ[ρuθ″uθ″]
    
    A = r_utII_utII
    ddq1_A = gradient(A, 1., axis=0, acc=acc, edge_stencil=edge_stencil, d=1)
    ddq2_A = gradient(A, 1., axis=1, acc=acc, edge_stencil=edge_stencil, d=1)
    #ddx1_A = ddq1_A*ddx1_q1 + ddq2_A*ddx1_q2
    ddx2_A = ddq1_A*ddx2_q1 + ddq2_A*ddx2_q2
    
    t7t = (1/r)*ddx2_A
    
    # ===
    
    ## +(1/r)·ρ·ur·ur
    t8r = (1/r) * rho * ur_Fv * ur_Fv
    
    ## +(1/r)·ρur″ur″
    t9r = (1/r) * r_urII_urII
    
    ## -(1/r)·ρ·uθ·uθ
    t10r = -1 * (1/r) * rho * ut_Fv * ut_Fv
    
    ## -(1/r)·ρuθ″uθ″
    t11r = -1 * (1/r) * r_utII_utII
    
    ## -(1/r)·τ_rr
    t12r = -1 * (1/r) * tau_rr
    
    ## +(1/r)·τ_θθ
    t13r = (1/r) * tau_tt
    
    # ===
    
    ## +(2/r)·ρ·ur·uθ
    t8t = (2/r) * rho * ur_Fv * ut_Fv
    
    ## +(2/r)·ρur″uθ″
    t9t = (2/r) * r_urII_utII
    
    ## -(2/r)·τ_rθ
    t10t = -(2/r) * tau_rt
    
    # ===
    
    data_r = {
        't1':t1r,
        't2':t2r,
        't3':t3r,
        't4':t4r,
        't5':t5r,
        't6':t6r,
        't7':t7r,
        't8':t8r,
        't9':t9r,
        't10':t10r,
        't11':t11r,
        't12':t12r,
        't13':t13r,
        }
    
    data_t = {
        't1':t1t,
        't2':t2t,
        't3':t3t,
        't4':t4t,
        't5':t5t,
        't6':t6t,
        't7':t7t,
        't8':t8t,
        't9':t9t,
        't10':t10t,
        }
    
    return data_r, data_t

def _calc_momentum_balance_terms(self, **kwargs):
    '''
    Calculate terms of the 2D time-averaged momentum conservation equation
    '''
    
    verbose      = kwargs.get('verbose',True)
    local_csys   = kwargs.get('local_csys',True)
    acc          = kwargs.get('acc',6)
    edge_stencil = kwargs.get('edge_stencil','full')
    xis          = kwargs.get('xis',None) ## which [x] indices to compute (None=all)
    csys_type    = kwargs.get('csys_type','cartesian')
    
    if verbose: print('\n'+'ztmd.calc_momentum_balance_terms()'+'\n'+72*'-')
    t_start_func = timeit.default_timer()
    
    if self.open_mode not in ['a','w','r+']:
        raise ValueError('Not able to write to hdf5 file')
    
    if verbose: even_print('local_csys',str(local_csys))
    if verbose: even_print('csys_type',csys_type)
    
    if local_csys:
        if self.rectilinear:
            print('>>> WARNING: this grid is rectilinear. Using local csys is not needed.')
        if 'csys/vtang' not in self:
            raise ValueError('csys/vtang not in hdf5')
        if 'csys/vnorm' not in self:
            raise ValueError('csys/vnorm not in hdf5')
        
        ## Read unit vectors (wall tangent, wall norm) from HDF5
        vtang = np.copy( self['csys/vtang'][()] )
        vnorm = np.copy( self['csys/vnorm'][()] )
        if verbose: even_print('vtang',str(vtang.shape))
        if verbose: even_print('vnorm',str(vnorm.shape))
    
    if (csys_type=='cartesian'):
        gnx = 'data_momentum_x'
        gny = 'data_momentum_y'
        gns = [gnx,gny]
    elif (csys_type=='cylindrical'):
        gnr = 'data_momentum_r'
        gnt = 'data_momentum_theta'
        gns = [gnr,gnt]
    else:
        raise NotImplementedError
    
    terms_x = {
        't1':( r'∂x[ρ·u·u]' , r'$\partial_{x}[\bar{\rho} \tilde{u} \tilde{u}]$'),
        't2':( r'∂y[ρ·u·v]' , r'$\partial_{y}[\bar{\rho} \tilde{u} \tilde{v}]$'),
        't3':( r'∂x[p]'     , r'$\partial_{x}[\bar{p}]$'),
        't4':( r'-∂x[τxx]'  , r'$-\partial_{x}[\overline{\tau}_{xx}]$'),
        't5':( r'∂x[ρu″u″]' , r'$\partial_{x}[\overline{\rho u^{\prime\prime} u^{\prime\prime}}]$'),
        't6':( r'-∂y[τxy]'  , r'$-\partial_{y}[\overline{\tau}_{xy}]$'),
        't7':( r'∂y[ρu″v″]' , r'$\partial_{y}[\overline{\rho u^{\prime\prime} v^{\prime\prime}}]$'),
        }
    
    terms_y = {
        't1':( r'∂x[ρ·u·v]' , r'$\partial_{x}[\bar{\rho} \tilde{u} \tilde{v}]$'),
        't2':( r'∂y[ρ·v·v]' , r'$\partial_{y}[\bar{\rho} \tilde{v} \tilde{v}]$'),
        't3':( r'∂y[p]'     , r'$\partial_{y}[\bar{p}]$'),
        't4':( r'-∂x[τxy]'  , r'$-\partial_{x}[\overline{\tau}_{xy}]$'),
        't5':( r'∂x[ρu″v″]' , r'$\partial_{x}[\overline{\rho u^{\prime\prime} v^{\prime\prime}}]$'),
        't6':( r'-∂y[τyy]'  , r'$-\partial_{y}[\overline{\tau}_{yy}]$'),
        't7':( r'∂y[ρv″v″]' , r'$\partial_{y}[\overline{\rho v^{\prime\prime} v^{\prime\prime}}]$'),
        }
    
    terms_r = {
        't1':( r'∂r[ρ·ur·ur]'       , r'$\partial_{r}[ \bar{\rho} \tilde{u}_r \tilde{u}_r ]$'),
        't2':( r'(1/r)·∂θ[ρ·ur·uθ]' , r'$\frac{1}{r} \partial_{\theta}[ \bar{\rho} \tilde{u}_r \tilde{u}_\theta ]$'),
        't3':( r'∂r[p]'             , r'$\partial_{r}[\bar{p}]$'),
        't4':( r'-∂r[τ_rr]'         , r'$-\partial_{r}[ \overline{\tau}_{r r} ]$'),
        't5':( r'∂r[ρur″ur″]'       , r'$\partial_{r}[ \overline{ \rho u^{\prime \prime}_r  u^{\prime \prime}_r } ]$'),
        't6':( r'-(1/r)·∂θ[τ_rθ]'   , r'$-\frac{1}{r} \partial_{\theta}[ \overline{\tau}_{r \theta} ]$'),
        't7':( r'(1/r)·∂θ[ρur″uθ″]' , r'$\frac{1}{r} \partial_{\theta}[ \overline{\rho u^{\prime \prime}_r  u^{\prime \prime}_\theta } ]$'),
        't8':( r'(1/r)·ρ·ur·ur'     , r'$\frac{1}{r} \bar{\rho} \tilde{u}_{r} \tilde{u}_{r}$'),
        't9':( r'(1/r)·ρur″ur″'     , r'$\frac{1}{r} \overline{ \rho u^{\prime \prime}_r u^{\prime \prime}_r }$'),
        't10':( r'-(1/r)·ρ·uθ·uθ'   , r'$-\frac{1}{r} \bar{\rho} \tilde{u}_{\theta} \tilde{u}_{\theta}$'),
        't11':( r'-(1/r)·ρuθ″uθ″'   , r'$-\frac{1}{r} \overline{ \rho u^{\prime \prime}_\theta u^{\prime \prime}_\theta }$'),
        't12':( r'-(1/r)·τ_rr'      , r'$-\frac{1}{r} \overline{\tau}_{r r}$'),
        't13':( r'(1/r)·τ_θθ'       , r'$\frac{1}{r} \overline{\tau}_{\theta \theta}$'),
        }
    
    terms_theta = {
        't1':( r'∂r[ρ·ur·uθ]'       , r'$\partial_{r}[ \bar{\rho} \tilde{u}_r \tilde{u}_\theta ]$'),
        't2':( r'(1/r)·∂θ[ρ·uθ·uθ]' , r'$\frac{1}{r} \partial_{\theta}[ \bar{\rho} \tilde{u}_\theta \tilde{u}_\theta ]$'),
        't3':( r'(1/r)·∂θ[p]'       , r'$\frac{1}{r} \partial_{\theta}[\bar{p}]$'),
        't4':( r'-∂r[τ_rθ]'         , r'$-\partial_{r}[ \overline{\tau}_{r \theta} ]$'),
        't5':( r'∂r[ρur″uθ″]'       , r'$\partial_{r}[ \overline{\rho u^{\prime \prime}_r  u^{\prime \prime}_\theta } ]$'),
        't6':( r'-(1/r)·∂θ[τ_θθ]'   , r'$-\frac{1}{r} \partial_{\theta}[ \overline{\tau}_{\theta \theta} ]$'),
        't7':( r'(1/r)·∂θ[ρuθ″uθ″]' , r'$\frac{1}{r} \partial_{\theta}[ \overline{\rho u^{\prime \prime}_\theta  u^{\prime \prime}_\theta } ]$'),
        't8':( r'(2/r)·ρ·ur·uθ'     , r'$\frac{2}{r} \bar{\rho} \tilde{u}_{r} \tilde{u}_{\theta}$'),
        't9':( r'(2/r)·ρur″uθ″'     , r'$\frac{2}{r} \overline{ \rho u^{\prime \prime}_r u^{\prime \prime}_\theta }$'),
        't10':( r'-(2/r)·τ_rθ'      , r'$-\frac{2}{r} \overline{\tau}_{r \theta}$'),
        }
    
    ## Create groups in ZTMD
    for gn in gns:
        if verbose: print(72*'-')
        if gn in self:
            del self[gn]
            if verbose: even_print('deleted group',f'/{gn}/')
        
        grp = self.create_group(gn)
        if verbose: even_print('created group',f'{grp.name}/')
        
        ## Set group attributes
        grp.attrs['local_csys'] = local_csys
        if verbose: even_print(f'{gn}.local_csys',str(grp.attrs['local_csys']))
    
    gn=None; del gn
    
    if verbose: print(72*'-')
    
    ## Initialize dsets for [x] terms
    if csys_type=='cartesian':
        for dsn, tags in terms_x.items():
            unicode,latex = tags
            shape = (self.ny,self.nx)
            if (f'{gnx}/{dsn}' in self):
                del self[f'{gnx}/{dsn}']
            dset = self.create_dataset(
                                    f'{gnx}/{dsn}', 
                                    shape=shape, 
                                    dtype=np.float32,
                                    chunks=None,
                                    fillvalue=np.float32(np.nan),
                                    )
            if verbose: even_print(f'{gnx}/{dsn}',str(shape))
            dset.attrs['latex'] = latex
            dset.attrs['unicode'] = unicode
            #if verbose: even_print(f'{dsn}.latex',str(dset.attrs['latex']))
        if verbose: print(72*'-')
    
    ## Initialize dsets for [y] terms
    if csys_type=='cartesian':
        for dsn, tags in terms_y.items():
            unicode,latex = tags
            shape = (self.ny,self.nx)
            if (f'{gny}/{dsn}' in self):
                del self[f'{gny}/{dsn}']
            dset = self.create_dataset(
                                    f'{gny}/{dsn}', 
                                    shape=shape, 
                                    dtype=np.float32,
                                    chunks=None,
                                    fillvalue=np.float32(np.nan),
                                    )
            if verbose: even_print(f'{gny}/{dsn}',str(shape))
            dset.attrs['latex'] = latex
            dset.attrs['unicode'] = unicode
            #if verbose: even_print(f'{dsn}.latex',str(dset.attrs['latex']))
        if verbose: print(72*'-')
    
    ## Initialize dsets for [r] terms
    if csys_type=='cylindrical':
        for dsn, tags in terms_r.items():
            unicode,latex = tags
            shape = (self.ny,self.nx)
            if (f'{gnr}/{dsn}' in self):
                del self[f'{gnr}/{dsn}']
            dset = self.create_dataset(
                                    f'{gnr}/{dsn}', 
                                    shape=shape, 
                                    dtype=np.float32,
                                    chunks=None,
                                    fillvalue=np.float32(np.nan),
                                    )
            if verbose: even_print(f'{gnr}/{dsn}',str(shape))
            dset.attrs['latex'] = latex
            dset.attrs['unicode'] = unicode
            #if verbose: even_print(f'{dsn}.latex',str(dset.attrs['latex']))
        if verbose: print(72*'-')
    
    ## Initialize dsets for [θ] terms
    if csys_type=='cylindrical':
        for dsn, tags in terms_theta.items():
            unicode,latex = tags
            shape = (self.ny,self.nx)
            if (f'{gnt}/{dsn}' in self):
                del self[f'{gnt}/{dsn}']
            dset = self.create_dataset(
                                    f'{gnt}/{dsn}', 
                                    shape=shape, 
                                    dtype=np.float32,
                                    chunks=None,
                                    fillvalue=np.float32(np.nan),
                                    )
            if verbose: even_print(f'{gnt}/{dsn}',str(shape))
            dset.attrs['latex'] = latex
            dset.attrs['unicode'] = unicode
            #if verbose: even_print(f'{dsn}.latex',str(dset.attrs['latex']))
        if verbose: print(72*'-')
    
    ## Read data
    if csys_type=='cylindrical' and hasattr(self,'crv_R'):
        crv_R = np.copy(self.crv_R)
        if verbose:
            even_print('crv_R',str(crv_R.shape))
            print(72*'-')
    
    if verbose:
        print('Reading data'+'\n'+'------------')
    u   = np.copy( self['data/u'][()].T   )
    v   = np.copy( self['data/v'][()].T   )
    rho = np.copy( self['data/rho'][()].T )
    p   = np.copy( self['data/p'][()].T   )
    mu  = np.copy( self['data/mu'][()].T  )
    
    if verbose: even_print('u',str(u.shape))
    if verbose: even_print('v',str(v.shape))
    if verbose: even_print('ρ',str(rho.shape))
    if verbose: even_print('p',str(p.shape))
    if verbose: even_print('μ',str(mu.shape))
    
    u_Fv = np.copy( self['data/u_Fv'][()].T )
    v_Fv = np.copy( self['data/v_Fv'][()].T )
    if verbose: even_print('u_Fv',str(u_Fv.shape))
    if verbose: even_print('v_Fv',str(v_Fv.shape))
    
    r_uII_uII = np.copy( self['data/r_uII_uII'][()].T )
    r_vII_vII = np.copy( self['data/r_vII_vII'][()].T )
    r_uII_vII = np.copy( self['data/r_uII_vII'][()].T )
    if verbose: even_print('ρu″u″',str(r_uII_uII.shape))
    if verbose: even_print('ρu″v″',str(r_uII_vII.shape))
    if verbose: even_print('ρv″v″',str(r_vII_vII.shape))
    
    x = np.copy(self.x)
    y = np.copy(self.y)
    
    if x.ndim==1 and y.ndim==1:
        x,y = np.meshgrid(x,y, indexing='ij')
    if verbose: even_print('x',str(x.shape))
    if verbose: even_print('y',str(y.shape))
    if verbose: print(72*'-')
    
    nx,ny = x.shape
    
    # ==================================================================
    
    def _get_transformation_matrix(xi):
        '''
        Get transformation matrix : translate,rotate,[translate]
        '''
        rotation_angle_z = np.arcsin(vnorm[xi,0,0])
        
        rotation_matrix = np.array( [[ np.cos(rotation_angle_z) , -np.sin(rotation_angle_z) , 0 ],
                                     [ np.sin(rotation_angle_z) ,  np.cos(rotation_angle_z) , 0 ],
                                     [ 0                        ,  0                        , 1 ]], dtype=np.float64 )
        
        ## Translate to [cx,cy]
        cx = x[xi,0]
        cy = y[xi,0]
        translation_matrix_A = np.array( [[ 1 , 0 , -cx ],
                                          [ 0 , 1 , -cy ],
                                          [ 0 , 0 ,  1  ]], dtype=np.float64 )
        
        ## For cylindrical csys, translate in [y] by -R (the local curvature radius)
        ##   such that origin is distance 'R' from the wall
        if (csys_type=='cylindrical'):
            cx = 0.
            try:
                cy = crv_R[xi]
            except NameError:
                cy = 0.
            
            translation_matrix_B = np.array( [[ 1 , 0 ,  0  ],
                                              [ 0 , 1 , -cy ],
                                              [ 0 , 0 ,  1  ]], dtype=np.float64 )
        else:
            translation_matrix_B = np.identity(3, dtype=np.float64) ## No second translation
        
        ## Form the full [3x3] matrix by multiplying together
        ##  --> ORDER MATTERS!
        transformation_matrix = np.einsum( 'ij,jk,kl->il', translation_matrix_B , rotation_matrix, translation_matrix_A )
        
        return transformation_matrix
    
    # ==================================================================
    
    if local_csys: ## Compute momentum balance in local csys
        
        if xis is None: ## Do all
            xis = np.arange(self.nx)
            n_prog = self.nx
        else: ## Do only for selected x-indices
            n_prog = len(xis)
        
        ## Assert passed 'xis' is OK
        if isinstance(xis,np.ndarray): ## numpy array
            if not xis.ndim==1:
                raise ValueError('xis must have .ndim==1')
            if not np.issubdtype(xis.dtype,np.integer):
                raise RuntimeError("xis must have subtype np.integer")
        else: ## some other iterable... are all elements ints?
            try:
                xis = [int(xi) for xi in xis]
            except Exception as e:
                raise RuntimeError("elements of xis are not interpretable as ints") from e
        
        # ==============================================================
        # Main loop
        # ==============================================================
        
        if csys_type=='cartesian':
            desc='Momentum balance [x,y]'
        elif csys_type=='cylindrical':
            desc='Momentum balance [r,θ]'
        else:
            raise ValueError
        
        if verbose:
            progress_bar = tqdm(
                total=n_prog,
                ncols=100,
                desc=desc,
                leave=True,
                file=sys.stdout,
                mininterval=0.1,
                smoothing=0.,
                #bar_format="\033[B{l_bar}{bar}| {n}/{total} [{percentage:.1f}%] {elapsed}/{remaining}\033[A\n\b",
                bar_format="{l_bar}{bar}| {n}/{total} [{percentage:.1f}%] {elapsed}/{remaining}",
                ascii="░█",
                colour='#00ffff',
                )
        
        for xi in xis:
            
            ## Get transformation matrix (global-->local) for this [x] index
            transformation_matrix = _get_transformation_matrix(xi)
            
            ## Take N points in [x] on either side
            ## This allows for full FD stencil at central point
            xiA = max( 0  , xi-12   )
            xiB = min( nx , xi+12+1 )
            xii = xi - xiA ## Index local to slice
            #tqdm.write(f'xi={xi:d},xiA={xiA:d},xiB={xiB:d},xii={xii:d}') ## DEBUG
            
            if (csys_type=='cartesian'):
                
                data_x, data_y = momentum_balance_kernel_cartesian(
                                                        x=x[xiA:xiB,:],
                                                        y=y[xiA:xiB,:],
                                                        u=u[xiA:xiB,:],
                                                        v=v[xiA:xiB,:],
                                                        u_Fv=u_Fv[xiA:xiB,:],
                                                        v_Fv=v_Fv[xiA:xiB,:],
                                                        p=p[xiA:xiB,:],
                                                        rho=rho[xiA:xiB,:],
                                                        mu=mu[xiA:xiB,:],
                                                        r_uII_uII=r_uII_uII[xiA:xiB,:],
                                                        r_uII_vII=r_uII_vII[xiA:xiB,:],
                                                        r_vII_vII=r_vII_vII[xiA:xiB,:],
                                                        transformation_matrix=transformation_matrix,
                                                        acc=acc,
                                                        edge_stencil=edge_stencil,
                                                        )
                
                ## Write to HDF5
                for dsn, arr in data_x.items():
                    dset = self[f'{gnx}/{dsn}']
                    dset[:,xi] = arr[xii,:] ## Write only 1 [x] position
                for dsn, arr in data_y.items():
                    dset = self[f'{gny}/{dsn}']
                    dset[:,xi] = arr[xii,:] ## Write only 1 [x] position
            
            elif (csys_type=='cylindrical'):
                
                data_r, data_t = momentum_balance_kernel_cylindrical(
                                                        x=x[xiA:xiB,:],
                                                        y=y[xiA:xiB,:],
                                                        u=u[xiA:xiB,:],
                                                        v=v[xiA:xiB,:],
                                                        u_Fv=u_Fv[xiA:xiB,:],
                                                        v_Fv=v_Fv[xiA:xiB,:],
                                                        p=p[xiA:xiB,:],
                                                        rho=rho[xiA:xiB,:],
                                                        mu=mu[xiA:xiB,:],
                                                        r_uII_uII=r_uII_uII[xiA:xiB,:],
                                                        r_uII_vII=r_uII_vII[xiA:xiB,:],
                                                        r_vII_vII=r_vII_vII[xiA:xiB,:],
                                                        transformation_matrix=transformation_matrix,
                                                        acc=acc,
                                                        edge_stencil=edge_stencil,
                                                        )
                
                ## Write to HDF5
                for dsn, arr in data_r.items():
                    dset = self[f'{gnr}/{dsn}']
                    dset[:,xi] = arr[xii,:] ## Write only 1 [x] position
                for dsn, arr in data_t.items():
                    dset = self[f'{gnt}/{dsn}']
                    dset[:,xi] = arr[xii,:] ## Write only 1 [x] position
            
            else:
                raise ValueError
            
            if verbose: progress_bar.update()
        if verbose: progress_bar.close()
    
    else: ## Compute momentum balance in global csys
        
        ## Send to kernel
        if csys_type=='cartesian': ## Cartesian
            
            data_x, data_y = momentum_balance_kernel_cartesian(
                                                        x=x,
                                                        y=y,
                                                        u=u,
                                                        v=v,
                                                        u_Fv=u_Fv,
                                                        v_Fv=v_Fv,
                                                        p=p,
                                                        rho=rho,
                                                        mu=mu,
                                                        r_uII_uII=r_uII_uII,
                                                        r_uII_vII=r_uII_vII,
                                                        r_vII_vII=r_vII_vII,
                                                        #transformation_matrix=transformation_matrix,
                                                        acc=acc,
                                                        edge_stencil=edge_stencil,
                                                        )
            
            ## Write to HDF5
            for dsn,arr in data_x.items():
                dset = self[f'{gnx}/{dsn}']
                if verbose:
                    even_print(dset.name,str(arr.T.shape))
                dset[:,:] = arr.T
            for dsn,arr in data_y.items():
                dset = self[f'{gny}/{dsn}']
                if verbose:
                    even_print(dset.name,str(arr.T.shape))
                dset[:,:] = arr.T
        
        elif csys_type=='cylindrical': ## Cylindrical
            
            data_r, data_t = momentum_balance_kernel_cylindrical(
                                                            x=x,
                                                            y=y,
                                                            u=u,
                                                            v=v,
                                                            u_Fv=u_Fv,
                                                            v_Fv=v_Fv,
                                                            p=p,
                                                            rho=rho,
                                                            mu=mu,
                                                            r_uII_uII=r_uII_uII,
                                                            r_uII_vII=r_uII_vII,
                                                            r_vII_vII=r_vII_vII,
                                                            #transformation_matrix=transformation_matrix,
                                                            acc=acc,
                                                            edge_stencil=edge_stencil,
                                                            )
            
            ## Write to HDF5
            for dsn,arr in data_r.items():
                dset = self[f'{gnr}/{dsn}']
                if verbose:
                    even_print(dset.name,str(arr.T.shape))
                dset[:,:] = arr.T
            for dsn,arr in data_t.items():
                dset = self[f'{gnt}/{dsn}']
                if verbose:
                    even_print(dset.name,str(arr.T.shape))
                dset[:,:] = arr.T
        
        else:
            raise ValueError
    
    self.get_header(verbose=False)
    if verbose: print(72*'-')
    if verbose: print('total time : ztmd.calc_momentum_balance_terms() : %s'%format_time_string((timeit.default_timer() - t_start_func)))
    if verbose: print(72*'-')
    return
