#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Oct 26 16:48:43 EDT 2021
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
import  salem
import geoxarray
from misc.shell import run_shell

cntry_name = 'China'
cntry_sname = cntry_name.replace(' ', '')
daname = 'tas'
xname = 'lon'
yname = 'lat'
shdf = salem.read_shapefile('chn_admbnda_adm1_ocha/chn_admbnda_adm1_ocha.shp')
name_cols = ['ADM1_EN', 'ADM1_PCODE']
print(daname)
if len(sys.argv)>1:
    idir = sys.argv[1]
else:
    idir = 'cmip6_tas_historical'
odir = f'{idir}_China_subregion'
if not os.path.exists(odir):
    os.mkdir(odir)
    print('[created]:', odir)
ifiles = [os.path.basename(ifile) for ifile in glob.glob(f'{idir}/*.nc')] 
ifiles.sort()
nfiles = len(ifiles)
#print(ifiles)
#for year,month in zip(t.year, t.month):
for ii,ifile in enumerate(ifiles, start=1):
    ifile_ = os.path.join(idir, ifile) #dir + basname
    print(f'{ii:02d} of {nfiles}:', ifile_)
    model = ifile.split('.')[-2]
    ofile = os.path.join(odir, ifile.replace(model, f'China.subregion.{model}'))
    if os.path.exists(ofile):
        print('[exists]:', ofile)
        continue
    else:
        print('[ofile]:', ofile)

    print('loading...')
    ds = salem.open_xr_dataset(ifile_).load()

    #rename
    #ds = ds.rename({'air': 't2m', 'lon': 'longitude', 'lat': 'latitude'})

    # transform lon [0,360) -> [-180,180)
    ds = ds.roll(lon=ds.lon.size//2)
    lon_new = lon = ds.lon.where(ds.lon<180, other=ds.lon-360).values
    ds = ds.assign_coords(lon=lon_new)

    #interp onto fine grids
    #ds[daname] = ds[daname].interp(longitude=np.arange(0,180,0.05), latitude=np.arange(0, 89, 0.05))
    
    das = []
    #region_IDs = []
    region_names = []
    N = len(shdf)
    for i in range(N):
        shape_sub = shdf.iloc[i:i+1,:]
        r = shdf.iloc[i, :]
        #print(r); sys.exit()
        lon_slice = slice(r.min_x, r.max_x)
        lat_slice = slice(r.min_y, r.max_y)
        g = r.geometry
        #province, pcode = r[name_cols[0]].split(' ')[0], r[name_cols[1]]#name_cols = ['ADM1_EN', 'ADM1_PCODE']
        province, pcode = '_'.join(r[name_cols[0]].split()), r[name_cols[1]]#name_cols = ['ADM1_EN', 'ADM1_PCODE']
        region_name = '_'.join([province,pcode])
     
        try:
            #da = ds[daname].salem.roi(geometry=g).geo.fldmean()
            da = ds[daname].salem.subset(shape=shape_sub, margin=1) \
                .salem.roi(shape=shape_sub) \
                .geo.fldmean()
            print('[OK]:', i+1,'of', N, region_name)
        except ValueError:
            centroid0 = shape_sub.centroid.iloc[0]
            lon0 = centroid0.x
            lat0 = centroid0.y
            da = ds[daname].interp(lon=lon0, lat=lat0).drop(['lon', 'lat'])
            #print('\t[Failed]', i+1,'of', N, region_name)
            print('\t[single grid point used]:', i+1,'of', N, region_name)
        das.append(da)
        region_names.append(region_name)
    
    #i_record = range(len(das))
    da = xr.concat(das, pd.Index(region_names, name='region_name')).transpose()
    da.attrs = ds[daname].attrs
    ds_new = da.to_dataset(name=daname)
    ds_new.attrs = ds.attrs
    ds_new.attrs['region_name'] = '; '.join(region_names)
    
    #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={daname: {'_FillValue': None}})
    print('[saved]:', ofile)

#ncrcat
#ofile = ofiles[0].replace('tmp/', '').replace(f'{years[0]}-{months[0]:02d}', f'{years[0]}-{years[-1]}')
#ofile = ofiles[0].replace('tmp/', '').replace(f'{years[0]}', f'{years[0]}-{years[-1]}')
#cmd = f'ncrcat {" ".join(ofiles)} {ofile}'
#run_shell(cmd)

if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    
    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()