"""
Regridding tool for MICOM output

"""

# import libraries 
import numpy as np 
from netCDF4 import Dataset
import xesmf as xe
import os
import sys
import warnings

# surpress warnings
warnings.filterwarnings("ignore")

def micom2reg(fldIn, gridPath, res, lonBounds=(0,360), latBounds=(-90,90),
              method='conservative', oceanMaskType='best', isMasked=True):
    """
Regrids MICOM output to a regular lon/lat grid. 

Parameters:
-----------
Mandatory 
    
     fldIn : input field on MICOM grid; the innermost/rightmost two 
             dimensions must be y,x and correspond to the global dimensions 
             in MICOM's grid file; additional dimensions (e.g. vertical and 
             time) are allowed 
        
  gridPath : full path to MICOM's grid file 
        
       res : resolution in deg; if skalar is specified the value is used for 
             both longitude and latitude; if tuple is specified the first value 
             is used for longitude and the second for latitude
              
Optional 

 lonBounds : tuple that specifies the western and eastern domain bounds; 
             note that the first longitude value will be shifted by res/2

 latBounds : tuple that specifies the southern and northern domain bounds; 
             note that the first longitude value will be shifted by res/2

    method : interpolation method; must be one of 'conservative', 'bilinear',
             'patch', 'nearest_s2d' and 'nearest_d2s' 
                 
 oceanMaskType : must be one of 'conservative', 'best' and 'maximized';
                 'conservative' sets points to nan if they contain any land,
                 'best' if more than 50% land and 'maximized' if 100% land

Returns:
--------
    fldInterp, lon2d, lat2d : output field, longitude matrix, latitude matrix        
        
Example:
--------
    
# import libraries
import matplotlib.pyplot as plt
from netCDF4 import Dataset
import micom2reg as mr
    
# read gx1 sample data (3d temperature) 
nc = Dataset('datasample_gx1.nc')
T_gx1 = nc.variables['templvl'][:]
depth = nc.variables['depth'][:]
nc.close()    

# interpolate data to 1x1 deg lonlat grid
T, lon, lat = mr.micom2reg(T_gx1,'gridsample_gx1.nc',1)
    
# plot map of surface temperature
plt.pcolormesh(lon,lat,T[0,0,:,:]) # horizontal plot of SST 
plt.show()

# plot meridional temperature section through Atlantic
plt.pcolormesh(lat[:,0],-depth,T[0,:,:,330]) 
plt.show()

History: 
--------
2019.10.30 created (Ingo.Bethke@uib.no)
    """
    global regridder, gridOut, maskOutInv, shapeInOld, gridPathOld, resOld, \
        methodOld, oceanMaskTypeOld
    #
    # make sure that input is a masked array; convert missing to mask
    fldIn = np.ma.masked_invalid(fldIn,copy=False)
    #
    # prepare weights if necessary
    if (not 'regridder' in globals() or fldIn.shape != shapeInOld 
        or gridPath != gridPathOld or res != resOld  or method !=methodOld 
        or oceanMaskType != oceanMaskTypeOld):   
        shapeInOld, gridPathOld, resOld, methodOld, oceanMaskTypeOld = \
            fldIn.shape, gridPath, res, method, oceanMaskType
        #
        # prepare MICOM coordinate information
        nc = Dataset(gridPath)
        lon, lat = nc.variables['plon'][:], nc.variables['plat'][:]
        clon, clat = nc.variables['pclon'][:], nc.variables['pclat'][:]
        if not isMasked: 
            pmask = nc.variables['pmask'][:]
        nc.close()            
        ydm, xdm = lon.shape 
        lonC, latC = np.zeros((ydm+1,xdm+1)), np.zeros((ydm+1,xdm+1))
        lonC[0:ydm,0:xdm], latC[0:ydm,0:xdm] = clon[0,:,:], clat[0,:,:]
        lonC[ydm,0:xdm], latC[ydm,0:xdm] = clon[3,ydm-1,:], clat[3,ydm-1,:] 
        lonC[0:ydm,xdm], latC[0:ydm,xdm] = clon[1,:,xdm-1], clat[1,:,xdm-1]
        lonC[ydm,xdm], latC[ydm,xdm] = clon[2,ydm-1,xdm-1], clat[2,ydm-1,xdm-1] 
        gridIn = {'lon': lon, 'lon_b': lonC, 'lat': lat, 'lat_b': latC}
        # 
        # prepare regular grid coordinate information 
        periodic=True 
        if type(res) == int or type(res) == float:   
            resLon, resLat = res, res 
        else: 
            resLon, resLat = res[0], res[1]
        if lonBounds[0] % 360 != lonBounds[1] % 360: 
            periodic=False 
        gridOut = xe.util.grid_2d(lonBounds[0],lonBounds[1],resLon,
                                  latBounds[0],latBounds[1],resLat)
        #
        # compute weights 
        weightsFile = '{:s}_{:d}x{:d}_{:d}x{:d}.nc'.format(method,ydm,xdm, \
                      gridOut.y.size,gridOut.x.size) 
        if os.path.isfile(weightsFile):
            regridder = xe.Regridder(gridIn, gridOut, method, periodic=periodic,
                                     reuse_weights=True,filename=weightsFile)
        else:
            regridder = xe.Regridder(gridIn, gridOut, method, periodic=periodic,
                                     reuse_weights=False)
        #
        # create ocean mask from data and interpolate to output grid
        if not isMasked: 
            maskIn = np.where(pmask < 0.5, np.zeros(fldIn.shape), 
                              np.ones(fldIn.shape))
        else:
            maskIn = np.where(fldIn.mask, np.zeros(fldIn.shape), 
                              np.ones(fldIn.shape))
        maskOut = regridder(maskIn)        
        if oceanMaskType == 'conservative':
            maskOut = np.where(maskOut>1.-1e-10, maskOut, np.nan)
        elif oceanMaskType == 'maximized':
            maskOut = np.where(maskOut>0.+1e-10, maskOut, np.nan) 
        else:
            if oceanMaskType != 'best':
                sys.exit('Unknown value for oceanMaskType. Will use "best".')
            maskOut = np.where(maskOut>0.5, maskOut, np.nan)
        maskOutInv = 1. / maskOut
    #
    # interpolate, normalize and mask
    # NOTE: To allow interpolation if one or more of contributing input cells 
    #       contain missing values, all missing values are set to zero in input. 
    #       This has the same effect as setting the grid cell area to zero, so 
    #       that cell effectively does not contribute. To ensure conservation, 
    #       the same is done in the interpolation of the ocean mask which is  
    #       used to normalize the output field.
    fldOut = regridder(np.where(fldIn.mask,0.,fldIn)) * maskOutInv
    return fldOut, np.array(gridOut.lon), np.array(gridOut.lat)