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

# country name and lon/lat ranges
cntry_name = 'Madagascar'
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
ifiles = '/tigress/wenchang/data/engey/erai/daily/0.75x0.75/tas/tas.daily.*.nc'
data_name = "t2m"
ofile = f'erai.tas.p75.{cntry_name}.nc'
xname = 'lon'
xlim = (xmin, xmax)
yname = 'lat'
ylim = (ymin, ymax)
yincrease = False

x = xr.open_mfdataset(ifiles)[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_mfdataset(ifiles)[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 -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].salem.roi(shape=shdf)

# save the masked data to the ofile
da.to_dataset().to_netcdf(ofile.replace(cntry_name, cntry_name+'.masked'), unlimited_dims='time')