#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Apr 21 15:31:22 EDT 2022
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    tt = Timer('start ' + ' '.join(sys.argv))
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
import xlinregress
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
#ifile = 'cru_ts4.05.1901.2020.pre.dat.nc'
ifile = 'precip.mon.total.v2018.nc'
season = 'annual'
sel_season = lambda x: x
if 'JJA' in sys.argv:
    season = 'JJA'
    sel_season = lambda x: x.isel(time=(x.time.dt.month>=6)&(x.time.dt.month<=8))
years = '1901-2016'
sel_years = lambda x: x.sel(time=slice('1901', '2016'))
if '1981-2016' in sys.argv:
    years = '1981-2016'
    sel_years = lambda x: x.sel(time=slice('1981', '2016'))
ofile = __file__.replace('.py', f'_{season}{years}.nc')
if os.path.exists(ofile):
    ds = xr.open_dataset(ofile)
    print('[loaded]:', ofile)
else:
    da = xr.open_dataset(ifile)['precip'].load()
    da = da.pipe(sel_season).pipe(sel_years)
    da = da.groupby('time.year').mean('time')
    da = da.pipe(lambda x: x/30).assign_attrs(units='mm/day') #units convert: mm/month -> mm/day
    ds = da.linregress.on(da.year, dim='year')
    ds.slope.attrs['units'] = 'mm/day/year'
    ds = ds.assign_attrs(season=season, years=years)
    ds.to_netcdf(ofile)
    print('[saved]:', ofile)
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, ax = plt.subplots()
    da = ds.slope.pipe(lambda x: x*100).assign_attrs(units='mm/day/century', long_name='linear trend')
    da = da.roll(lon=da.lon.size//2, roll_coords=True)
    lon = da.lon.where(da.lon<180, other=da.lon-360)
    da = da.assign_coords(lon=lon.values)
    da.plot(ax=ax, vmax=1, levels=11, cmap='BrBG')

    ax.set_title(ifile[:-3] + f' {season} {years}', loc='left')

    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'_{season}{years}.png')
        if 'overwritefig' in sys.argv[1:]:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    print()
    plt.show()