#!/usr/bin/env python
import numpy as np
import xarray as xr
import os
import salem

# country name and lon/lat ranges
cntry_name = 'Mozambique'
cntry_sname = cntry_name.replace(' ', '')
shdf = salem.read_shapefile(salem.get_demo_file('world_borders.shp'))
shdf = shdf[ shdf['CNTRY_NAME']==cntry_name ]
xmin = shdf.min_x.min() - 1
xmax = shdf.max_x.max() + 1
ymin = shdf.min_y.min() - 1
ymax = shdf.max_y.max() + 1

# IO files
expname = 'CTL1860_newdiag_tigercpu_intelmpi_18_576PE'
ifiles = f'/tigress/wenchang/MODEL_OUT/{expname}/POSTP/????0101.atmos_daily.nc'
ifile1 = ifiles.replace('????', '0001')
data_name = "precip"
ofile = f'flor.{expname}.{data_name}.{cntry_sname}.nc'
xname = 'grid_xt'
xlim = (xmin, xmax)
yname = 'grid_yt'
ylim = (ymin, ymax)
yincrease = True

x = xr.open_dataset(ifile1)[xname].values
ix = x.argsort()
L = (x>=xlim[0]) & (x<=xlim[1])
ixlim = ix[L][0], ix[L][-1]
print(f'xlim: {xlim}; ixlim: {ixlim}')

y = xr.open_dataset(ifile1)[yname].values
iy = y.argsort()
if yincrease is False:
    iy = iy[-1::-1]# reverse the indices when y is decreasing along axis
L = (y>=ylim[0]) & (y<=ylim[1])
iylim = iy[L][0], iy[L][-1]
print(f'ylim: {ylim}; iylim: {iylim}')

# extract the data
cmd = f'ncrcat -h -v {data_name} -d {xname},{ixlim[0]},{ixlim[1]} -d {yname},{iylim[0]},{iylim[1]} {ifiles} {ofile}'
print(cmd, '...')
s = os.system(cmd)
if s==0:
    print('[OK]:', cmd)

# mask the grids outside of the country
da = xr.open_dataset(ofile)[data_name].rename({xname:'lon', yname:'lat'}).salem.roi(shape=shdf)

# save the masked data to the ofile
ofile_mask = ofile.replace(cntry_sname, cntry_sname+'.masked')
da.to_dataset().to_netcdf(ofile_mask, unlimited_dims='time')
print('[saved]:', ofile_mask)
print('*done*')