#!/usr/bin/env python3

#;;DIAG_NORCPM;; OCNGRIDFILE: ../../Data/grid.nc
#;;DIAG_NORCPM;; OCNGRID: p
#;;DIAG_NORCPM;; DSTFILE: VAR_ocngrid.nc
#;;DIAG_NORCPM;; SRCFILE: 
#;;DIAG_NORCPM;; SRCVAR: 

from os.path import exists
from sys import exit

def regrid2ocn(ocngridfile,srcf,srcvn,ocngrid='p'):  # not really convinent
    import xarray as xr
    import xesmf  as xe
    ## output grid
    if ocngrid == 'plat plon': ocngrid = 'p'
    gridf = xr.open_dataset(ocngridfile)
    plat = gridf.variables[ocngrid+'lat']
    plon = gridf.variables[ocngrid+'lon']

    ## input data and grid
    ifile = xr.open_dataset(srcf)
    inlat = ifile.variables['lat']
    inlon = ifile.variables['lon']

    ## set grid
    srcgrid = xr.Dataset({
        'lat': inlat,
        'lon': inlon
        })
    dstgrid = xr.Dataset({
        'lat': plat,
        'lon': plon
        })
    regridder = xe.Regridder(srcgrid,dstgrid,'bilinear',periodic=True)
    return regridder(ifile[srcvn])

if __name__ == '__main__':
    ocngridfile = '/nird/datalake/NS9039K/users/pgchiu/diag_norcpm/grid_tnx1v4_20170622.nc'
    ocngrid = 'p'
    srcfile = 'amip_tos.nc'
    srcvar  = 'tos'
    dstfile = 'amip_tos_ocngrid.nc'

    if exists(dstfile): 
        print(f'{dstfile} exist, skip.')
        exit()
    ## regrid
    import xarray as xr
    ovar = regrid2ocn(ocngridfile,srcfile,srcvar,ocngrid)
    outds = xr.Dataset({srcvar:ovar})
    outds.to_netcdf(dstfile,mode='w')