# -*- Mode: python -*-
### -
### - Computes polynomial fit
### -
## -
## - Olivier Marti (olivier.marti@lsce.ipsl.fr)
## - Institut Pierre Simon Laplace
## - Laboratoire des Sciences du Climate et de L'environnment
## -
print 'Demarrage'
## - Needed Python modules
import cdms2, MV2, numpy as np
import nemo

def clo_lon (lon, lon_ref) : # Find closest longitude
      z0 = lon
      for nn in range (2) :
         z0 = np.where ( z0 > lon_ref+180.0, z0-360.0, z0)
         z0 = np.where ( z0 < lon_ref-180.0, z0+360.0, z0)

      return z0

## - NetCDF parameters
# - If you want to turn that off or set different values of compression use the functions:
cdms2.setNetcdfShuffleFlag     (0)
cdms2.setNetcdfDeflateFlag     (0)
cdms2.setNetcdfDeflateLevelFlag(0)

print 'Ouverture fichier'
f_grid = cdms2.open ( 'tmp_ORCA025_coordinates_mask.nc', 'r+' )

print 'Lecture tmask'
tmask = f_grid ( 'mask_T' )

print 'Lecture des axes'
yy      = tmask.getAxis(0) ; ny = len(yy)
xx      = tmask.getAxis(1) ; nx = len(xx)

print 'Creation axe'
nvertex = cdms2.createAxis ( np.arange (4) )

print 'Lecture lon lat'
glamt =  f_grid ( 'nav_lon_grid_T' ).filled(0)
gphit =  f_grid ( 'nav_lat_grid_T' ).filled(0)
glamu =  f_grid ( 'nav_lon_grid_U' ).filled(0)
gphiu =  f_grid ( 'nav_lat_grid_U' ).filled(0)
glamv =  f_grid ( 'nav_lon_grid_V' ).filled(0)
gphiv =  f_grid ( 'nav_lat_grid_V' ).filled(0)
glamf =  f_grid ( 'nav_lon_grid_F' ).filled(0)
gphif =  f_grid ( 'nav_lon_grid_F' ).filled(0)


# Construction des limites
#   1,2->3      2,2->2
#   1,1->0      2,1->1

# :, :, \([1,2]\), \([1,2]\)
# \1, \2, :, :

print 'Creation bounds_*_T'
bounds_lon_grid_T = np.zeros ( (ny, nx, 4), dtype=np.float32 )
bounds_lat_grid_T = np.zeros ( (ny, nx, 4), dtype=np.float32 )

print 'Calcul bounds_*_T'
bounds_lon_grid_T [ :, :, 0] = np.roll ( np.roll ( glamf [:,:], 1, axis=0), 1, axis=1) # EOSHIFT (CSHIFT (glamf(:,:), -1, dim = 1), -1, dim = 2 )
bounds_lon_grid_T [ :, :, 1] = np.roll ( glamf [:,:], 1, axis=0) # CSHIFT  (glamf(:,:), -1, dim = 2 )
bounds_lon_grid_T [ :, :, 3] = np.roll ( glamf [:,:], 1, axis=1) # CSHIFT  (glamf(:,:), -1, dim = 1 )
bounds_lon_grid_T [ :, :, 2] =           glamf [:,:]

bounds_lat_grid_T [ :, :, 0] = np.roll ( np.roll ( gphif [:,:], 1, axis=0), 1, axis=1)
bounds_lat_grid_T [ :, :, 1] = np.roll ( gphif [:,:], 1, axis=0)
bounds_lat_grid_T [ :, :, 3] = np.roll ( gphif [:,:], 1, axis=1)
bounds_lat_grid_T [ :, :, 2] =           gphif [:,:]

bounds_lon_grid_T = cdms2.createVariable ( bounds_lon_grid_T, axes=(yy, xx, nvertex), id='bounds_lon_grid_T', dtype=np.float32 )
bounds_lat_grid_T = cdms2.createVariable ( bounds_lat_grid_T, axes=(yy, xx, nvertex), id='bounds_lat_grid_T', dtype=np.float32 )

print 'Calcul bounds_*_U'
bounds_lon_grid_U = np.zeros ( (ny, nx, 4), dtype=np.float32 )
bounds_lat_grid_U = np.zeros ( (ny, nx, 4), dtype=np.float32 )

bounds_lon_grid_U [ :, :, 0] = np.roll ( glamv [:,:], 1, axis=0) # CSHIFT  (glamv(:,:), -1, dim = 2 )
bounds_lon_grid_U [ :, :, 1] = np.roll ( np.roll ( glamv [:,:], 1, axis=0), 1, axis=1) # EOSHIFT (CSHIFT (glamv(:,:),  1, dim = 1 ), -1, dim = 2)
bounds_lon_grid_U [ :, :, 3] =           glamv [:,:]
bounds_lon_grid_U [ :, :, 2] = np.roll ( glamv [:,:], 1, axis=1) # CSHIFT  (glamv(:,:),  1, dim = 1 )

bounds_lat_grid_U [ :, :, 0] = np.roll ( gphiv [:,:], 1, axis=0)
bounds_lat_grid_U [ :, :, 1] = np.roll ( np.roll ( gphiv [:,:], 1, axis=0), 1, axis=1)
bounds_lat_grid_U [ :, :, 3] =           gphiv [:,:]
bounds_lat_grid_U [ :, :, 2] = np.roll ( gphiv [:,:], 1, axis=1)

bounds_lon_grid_U = cdms2.createVariable ( bounds_lon_grid_U, axes = ( yy, xx, nvertex ), id='bounds_lon_grid_U' )
bounds_lat_grid_U = cdms2.createVariable ( bounds_lat_grid_U, axes = ( yy, xx, nvertex ), id='bounds_lon_grid_U' )

print 'Calcul bounds_*_V'
bounds_lon_grid_V = np.zeros ( (ny, nx, 4), dtype=np.float32 )
bounds_lat_grid_V = np.zeros ( (ny, nx, 4), dtype=np.float32 )

bounds_lon_grid_V [ :, :, 0] = np.roll (glamu [:,:], 1, axis=1) # CSHIFT  (glamu(:,:), -1, dim = 1 )
bounds_lon_grid_V [ :, :, 1] =          glamu [:,:]
bounds_lon_grid_V [ :, :, 3] = np.roll (np.roll (glamu [:,:], 1, axis=0), 1, axis=1) # EOSHIFT (CSHIFT (glamu(:,:), -1, dim = 1 ),  1, dim = 2)
bounds_lon_grid_V [ :, :, 2] = np.roll (glamu [:,:], 1, axis=0) # EOSHIFT (glamu(:,:), 1, dim = 2 )

bounds_lat_grid_V [ :, :, 0] = np.roll (gphiu [:,:], 1, axis=1)
bounds_lat_grid_V [ :, :, 1] =          gphiu [:,:]
bounds_lat_grid_V [ :, :, 3] = np.roll (np.roll (glamu [:,:], 1, axis=0), 1, axis=1)
bounds_lat_grid_V [ :, :, 2] = np.roll (gphiu [:,:], 1, axis=0) 

bounds_lon_grid_V = cdms2.createVariable ( bounds_lon_grid_V, axes = ( yy, xx, nvertex ), id='bounds_lon_grid_V' )
bounds_lat_grid_V = cdms2.createVariable ( bounds_lat_grid_V, axes = ( yy, xx, nvertex ), id='bounds_lat_grid_V' )

print 'Calcul bounds_*_F'
bounds_lon_grid_F = np.zeros((ny, nx, 4), dtype=np.float32)
bounds_lat_grid_F = np.zeros((ny, nx, 4), dtype=np.float32)

bounds_lon_grid_F [ :, :, 0] =          glamt [:,:]
bounds_lon_grid_F [ :, :, 1] = np.roll (glamt [:,:], 1, axis=1) # CSHIFT  (glamt(:,:), 1, dim = 1 )
bounds_lon_grid_F [ :, :, 3] = np.roll (glamt [:,:], 1, axis=0) # EOSHIFT (glamt(:,:), 1, dim = 2 )
bounds_lon_grid_F [ :, :, 2] = np.roll (np.roll (glamt [:,:], 1, axis=0), 1, axis=1) # EOSHIFT (CSHIFT (glamt(:,:),  1, dim = 1),  1, dim = 2)

bounds_lat_grid_F [ :, :, 0] =          gphit [:,:]
bounds_lat_grid_F [ :, :, 1] = np.roll (gphit [:,:], 1, axis=1) # CSHIFT  (glamt(:,:), 1, dim = 1 )
bounds_lat_grid_F [ :, :, 3] = np.roll (gphit [:,:], 1, axis=0) # EOSHIFT (glamt(:,:), 1, dim = 2 )
bounds_lat_grid_F [ :, :, 2] = np.roll (np.roll (gphit [:,:], 1, axis=0), 1, axis=1) # EOSHIFT (CSHIFT (glamt(:,:),  1, dim = 1),  1, dim = 2)

bounds_lon_grid_F = cdms2.createVariable ( bounds_lon_grid_F, axes = ( yy, xx, nvertex ), id='bounds_lon_grid_F' )
bounds_lat_grid_F = cdms2.createVariable ( bounds_lat_grid_F, axes = ( yy, xx, nvertex ), id='bounds_lat_grid_F' )

print 'Trucage au sud'

#bounds_lat_grid_T [ 0, :, 0:1] = -90.0
#bounds_lat_grid_U [ 0, :, 0:1] = -90.0
#bounds_lat_grid_V [ 0, :, 0:1] = -90.0
#bounds_lat_grid_F [ 0, :, 0:1] = -90.0

#bounds_lon_grid_T [ 0, :, 0:1] = bounds_lon_grid_T  [ 0, :, 2:3 ]
#bounds_lon_grid_U [ 0, :, 0:1] = bounds_lon_grid_U  [ 0, :, 2:3 ]
#bounds_lon_grid_V [ 0, :, 0:1] = bounds_lon_grid_V  [ 0, :, 2:3 ]
#bounds_lon_grid_F [ 0, :, 0:1] = bounds_lon_grid_F  [ 0, :, 2:3 ]

print 'lbc'
for nn in range (4) :
    print 'nn : ', nn
    bounds_lon_grid_T [:,:,nn] = nemo.lbc (bounds_lon_grid_T [:,:,nn], npolj=6, cd_type='T', psgn=1.0)
    bounds_lat_grid_T [:,:,nn] = nemo.lbc (bounds_lat_grid_T [:,:,nn], npolj=6, cd_type='T', psgn=1.0)
    bounds_lon_grid_U [:,:,nn] = nemo.lbc (bounds_lon_grid_U [:,:,nn], npolj=6, cd_type='U', psgn=1.0)
    bounds_lat_grid_U [:,:,nn] = nemo.lbc (bounds_lat_grid_U [:,:,nn], npolj=6, cd_type='U', psgn=1.0)
    bounds_lon_grid_V [:,:,nn] = nemo.lbc (bounds_lon_grid_V [:,:,nn], npolj=6, cd_type='V', psgn=1.0)
    bounds_lat_grid_V [:,:,nn] = nemo.lbc (bounds_lat_grid_V [:,:,nn], npolj=6, cd_type='V', psgn=1.0)
    bounds_lon_grid_F [:,:,nn] = nemo.lbc (bounds_lon_grid_F [:,:,nn], npolj=6, cd_type='F', psgn=1.0)
    bounds_lat_grid_F [:,:,nn] = nemo.lbc (bounds_lat_grid_F [:,:,nn], npolj=6, cd_type='F', psgn=1.0)

    bounds_lon_grid_T[:,:,nn] = clo_lon ( bounds_lon_grid_T[:,:,nn], lon_ref=glamt )
    bounds_lon_grid_U[:,:,nn] = clo_lon ( bounds_lon_grid_U[:,:,nn], lon_ref=glamu )
    bounds_lon_grid_V[:,:,nn] = clo_lon ( bounds_lon_grid_V[:,:,nn], lon_ref=glamv )
    bounds_lon_grid_F[:,:,nn] = clo_lon ( bounds_lon_grid_F[:,:,nn], lon_ref=glamf )
    

f_grid.write ( bounds_lon_grid_T )
f_grid.write ( bounds_lat_grid_T )
f_grid.write ( bounds_lon_grid_U )
f_grid.write ( bounds_lat_grid_U )
f_grid.write ( bounds_lon_grid_V )
f_grid.write ( bounds_lat_grid_V )
f_grid.write ( bounds_lon_grid_F )
f_grid.write ( bounds_lat_grid_F )