#!/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 = 'Ts'
daname_new = 'GMST'
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_region = lambda x: x
years = range(1871,2017)
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_Ts_FLOR_nudgedHist_6ens_1871-2016.nc'
    print('loading...')
    da = xr.open_dataarray(ifile).pipe(sel_region).load()
    print('calculating...')
    da = da.geo.fldmean().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)

#5-year rolling mean
ofile_r5_nc = ofile_nc.replace('.nc', '.r5.nc')
ofile_r5_csv = ofile_csv.replace('.csv', '.r5.csv')
if os.path.exists(ofile_r5_nc):
    da_r5 = xr.open_dataarray(ofile_r5_nc)
    print('[loaded]:', ofile_r5_nc)
else:
    running5 = lambda x: x.rolling(year=5, center=True, min_periods=1).mean()
    da_r5 = da.pipe(running5)
    da_r5.to_dataset(name=daname_new).to_netcdf(ofile_r5_nc)
    print('[saved]:', ofile_r5_nc)
    da_r5.rename(daname_new).transpose('year', 'en').to_pandas().to_csv(ofile_r5_csv)
    print('[saved]:', ofile_r5_csv)

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    da.plot(hue='en', color='gray', alpha=1/4)
    da_r5.plot(hue='en', ls='--')
    
    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()