#!/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 = 'tmax'
daname_new = 'TXm'
rename_grid = lambda x: x.rename(grid_xt='lon', grid_yt='lat')
lons = slice(360-123, 360-119)
lats = slice(45, 52)
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.where((x.time.dt.month>=6)&(x.time.dt.month<=8)) #JJA only 
years = range(1870,2015)
ofile_nc = __file__.replace('.py', '.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_tmax_AM4urban_wasteCool_0urban_amip_5ens_1870-2014.nc'
    print('loading...')
    da = xr.open_dataarray(ifile).pipe(sel_region).load()
    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'
    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()