import xarray as xr
import xesmf as xe
#
def NorESM_to_WOA_regridder(NorESM_grid_path='/cluster/shared/noresm/inputdata/ocn/micom/tnx1v4/20170601/', WOA_grid_path='https://data.nodc.noaa.gov/thredds/dodsC/ncei/woa/temperature/decav81B0/1.00/woa18_decav81B0_t00_01.nc', grid_weight_path = '/cluster/work/users/cgu025/OMIP_diag/grid_weights/',regrid_mode='conservative',reuse_weights=True):
    '''
    '''
    # open WOA data to get the grid 
    WOA =xr.open_dataset(WOA_grid_path, decode_times=False)
    mask_WOA = WOA.t_an.isel(depth=0,time=0).notnull().drop('depth').drop('time')
    mask_WOA = mask_WOA.rename('mask_WOA').to_dataset().rename_dims({'lon':'x','lat':'y'}).mask_WOA
    # setup WOA grid
    # 
    lat = WOA.lat.rename({'lat':'y'})
    lon = WOA.lon.rename({'lon':'x'})
    lat_b = xr.concat([WOA.lat_bnds.isel(nbounds=0),WOA.lat_bnds.isel(lat=-1).isel(nbounds=1)],dim='lat').rename('lat_b').rename({'lat':'y_b'})
    lon_b = xr.concat([WOA.lon_bnds.isel(nbounds=0),WOA.lon_bnds.isel(lon=-1).isel(nbounds=1)],dim='lon').rename('lon_b').rename({'lon':'x_b'})
    #
    WOA_grid = xr.merge([lon,lat,lon_b,lat_b])
    #
    # NorESM grid 
    grid = xr.open_dataset(NorESM_grid_path+'grid.nc')
    mask_NorESM = grid.pmask.astype(float).rename('mask_NorESM')
    #
    lat_noresm = grid.plat.rename('lat')
    lon_noresm = grid.plon.rename('lon')
    #
    lon_b_noresm = xr.concat([grid.pclon.isel(nv=0),grid.pclon.isel(nv=1,x=-1)],dim='x')
    lon_b_noresm = xr.concat([lon_b_noresm,xr.concat([grid.pclon.isel(nv=3,y=-1),grid.pclon.isel(nv=2,y=-1,x=-1)],dim='x')],dim='y')
    lat_b_noresm= xr.concat([grid.pclat.isel(nv=0),grid.pclat.isel(nv=1,x=-1)],dim='x')
    lat_b_noresm = xr.concat([lat_b_noresm,xr.concat([grid.pclat.isel(nv=3,y=-1),grid.pclat.isel(nv=2,y=-1,x=-1)],dim='x')],dim='y')
    #
    lon_b_noresm = lon_b_noresm.rename('lon_b').rename({'y':'y_b','x':'x_b'})
    lat_b_noresm = lat_b_noresm.rename('lat_b').rename({'y':'y_b','x':'x_b'})     
    #
    NorESM_grid = xr.merge([lon_noresm,lat_noresm,lon_b_noresm,lat_b_noresm])  
    #  
    # create the regridder
    regridder = xe.Regridder(NorESM_grid, WOA_grid, regrid_mode,filename=grid_weight_path+'NORESM_to_WOA_'+regrid_mode+'.nc',reuse_weights=reuse_weights)
    #   
    return regridder, mask_WOA, mask_NorESM


#
fpath = '/cluster/projects/nn2345k/cgu025/omip_diag_core2/output/'
datestr='core2_20190804'
fnames=['tos_gn_'+datestr+'.nc','sos_gn_'+datestr+'.nc','zos_gn_'+datestr+'.nc','mlotst_gn_'+datestr+'.nc']
regridder, mask_WOA, mask_NorESM = NorESM_to_WOA_regridder(regrid_mode='conservative',reuse_weights=True)
#
mask_out = regridder(mask_NorESM)
#
for j,fname in enumerate(fnames):
    data     = xr.open_dataset(fpath+fname)
    varname  = list(data.variables).pop(-1)
    data_out = regridder(data[varname])
    data_out = data_out/mask_out # make sure land-ocean boundary is taken care of
    data_out = data_out.where(mask_WOA) # assign WOA mask
    data_out.to_dataset(name=varname[:-3]).to_netcdf(fpath+fname[:-20]+datestr+'_WOA.nc')
    data.close()