#!/usr/bin/env python
import xarray as xr
import sys, time
import  salem
import pandas as pd
import numpy as np

cntry_name = 'Sri Lanka'
cntry_sname = cntry_name.replace(' ', '')
data_name = 'precip'
xname = 'longitude'
yname = 'latitude'
shdf = salem.read_shapefile('LKA_adm1.shp')
ifile = f'../data/chirps.p05.{cntry_sname}.nc' 
ofile = f'chirps.{cntry_sname}.subregion.nc'
ds = salem.open_xr_dataset(ifile)

das = []
ID_region = []
NAME_region = []
N = len(shdf)
for i in range(N):
    r = shdf.iloc[i, :]
    lon_slice = slice(r.min_x, r.max_x)
    lat_slice = slice(r.min_y, r.max_y)
    g = r.geometry
 
    try:
        da = ds[data_name].salem.roi(geometry=g).mean([xname, yname])
        das.append(da)
        ID_region.append(r.ID_1)
        NAME_region.append(r.NAME_1)
        print('[OK]:', i+1,'of', N, r.ID_1, r.NAME_1)
    except:
        print('\t[Failed]', i+1,'of', N, r.ID_1, r.NAME_1)
        pass

i_record = range(len(das))
da = xr.concat(das, pd.Index(ID_region, name='ID_region')).transpose()
da.attrs = ds[data_name].attrs
ds_new = da.to_dataset(name=data_name)
ds_new.attrs = ds.attrs
ds_new.attrs['NAME_region'] = '; '.join(NAME_region)

#scodes = xr.DataArray(scodes, dims='i_county', coords=[i_county,])
#ccodes = xr.DataArray(ccodes, dims='i_county', coords=[i_county,])
#xr.Dataset(dict(precip=da, scode=scodes, ccode=ccodes)).to_netcdf(f'pr{year}.nc', unlimited_dims='time')
ds_new.to_netcdf(ofile, unlimited_dims='time', 
    encoding={'ID_region': {'dtype': np.int32}, 'time': {'dtype': 'int32'}})