#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Jun 10 13:01:07 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
tspan = slice('1958', '2014')
tag = f'{tspan.start}-{tspan.stop}'
ofile = __file__.replace('.py', f'_{tag}.nc')
if os.path.exists(ofile):
    rg = xr.open_dataset(ofile)
    print('[loaded]:', ofile)
else:
    #ifile = '/tigress/wenchang/HIRAM/input_wy/HadISST/HadISST_sst.187001-202202.nc'
    ifile = '/tigress/wenchang/data/HadISST/v2/HadISST.2.1.0.0_realisation_dec2010_1346.nc'
    da = xr.open_dataset(ifile)['sst'].sel(time=tspan).load().rename(dict(longitude='lon', latitude='lat'))
    #roll lon to center 180
    da = da.roll(lon=da.lon.size//2, roll_coords=True)
    lon_new = da.lon.where(da.lon>0, other=da.lon+360).values
    da = da.assign_coords(lon=lon_new)
    #mask non-valid value
    #da = da.where(da>-500)
    #da.isel(time=0).plot()
    da = da.groupby('time.year').mean('time') #to yearly
    
    
    #linear trend
    rg = da.linregress.on(da.year, dim='year')
    rg.to_netcdf(ofile)
    print('[saved]:', ofile)

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    fig, ax = plt.subplots()
    da = rg.slope.pipe(lambda x: (x*100).sel(lat=slice(30,-30)).assign_attrs(units='K/century', long_name='linear trend'))
    da.plot.contourf(vmax=3.2, levels=33, extend='both')
    cs = da.plot.contour(levels=[1,], linewidths=0.5, colors='gray')
    ax.clabel(cs)
    mapplot(xticks=range(0,360,60), fill_continents=True)
    title = f'HadISST2 SST linear trend over {tspan.start}-{tspan.stop}'
    ax.set_title(title)
    ax.set_ylabel('')
    ax.set_xlabel('')
    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'_{tag}.png')
        if 'overwritefig' in sys.argv[1:]:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv:
        pass
    else:
        plt.show()