#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Jun 29 16:14:32 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
import geoxarray
from misc.landmask import flagland
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
daname = 'Tas'
daname_new = 'Tas'
rename_grid = lambda x: x.rename(grid_xt='lon', grid_yt='lat')
region = 'central'
#lons = slice(360-123, 360-119)
#lats = slice(45, 52)
lons = slice(360-103, 360-94)
lats = slice(30, 45)
if 'southeast' in sys.argv:
    region = 'southeast'
    lons = slice(360-93, 360-80)
    lats = slice(25, 33)
elif 'northeast' in sys.argv:
    region = 'northeast'
    lons = slice(360-78, 360-70)
    lats = slice(36, 44)
long_name_region = f'{360-lons.start}-{360-lons.stop}W_{lats.start}-{lats.stop}N_land'
long_name_region = f'{360-lons.start}-{360-lons.stop}W_{lats.start}-{lats.stop}N_land'
sel_region = lambda x: x.sel(lon=lons, lat=lats)
sel_season = lambda x: x.isel(time=(x.time.dt.month>=6)&(x.time.dt.month<=8)) #JJA only 
years = range(1870,2021)
ofile_nc = __file__.replace('.py', f'_{region}.nc')
#ofile_csv = __file__.replace('.py', '.csv')
if os.path.exists(ofile_nc):
    da = xr.open_dataarray(ofile_nc)
    print('[loaded]:', ofile_nc)
else:
    ifile = '../data_Tas_AM4urban_wasteCool_amip_5ens_1870-2020.nc'
    print('loading...')
    da = xr.open_dataarray(ifile).pipe(sel_region).load().sel(en=[0,21,22,23,24])
    #more ens
    ifile = './data_Tas_AM4urban_wasteCool_amip_new5ens_1870-2020.nc'
    da_more = xr.open_dataarray(ifile).pipe(sel_region).load()
    da = xr.concat([da, da_more], dim='en')
    print('calculating...')
    #land only
    landflag = flagland(da)
    da = da.where(landflag>0.5)
    #da = da.geo.fldmean().groupby('time.year').max('time')
    da = da.geo.fldmean().pipe(sel_season).groupby('time.year').mean('time')
    if da.max() > 200:
        da = da - 273.15 #units K -> C
        da.attrs['units'] = 'degC'
    da.attrs['region'] = long_name_region
    print('saving...')
    da.to_dataset(name=daname_new).to_netcdf(ofile_nc)
    print('[saved]:', ofile_nc)
    #da.rename(daname_new).transpose('year', 'en').to_pandas().to_csv(ofile_csv)
    #print('[saved]:', ofile_csv)
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    da.plot(hue='en')
    
    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()