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

# country name and lon/lat ranges
cntry_name = 'Peru'
cntry_sname = cntry_name.replace(' ', '')
shdf = salem.read_shapefile(salem.get_demo_file('world_borders.shp'))
shdf = shdf[ shdf['CNTRY_NAME']==cntry_name ]


# IO files
data_name = "mx2t"
ifile = f'erai.tmax.p05.{cntry_sname}.nc'
xname = 'lon'
yname = 'lat'

# mask the grids outside of the country
da = xr.open_dataset(ifile)[data_name]
x = da[xname].values
x[x>180] = x[x>180] - 360
da[xname] = x
da = da.salem.roi(shape=shdf)

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