#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Dec 15 22:09:06 EST 2021
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    s = ' '
    tt = Timer(f'start {s.join(sys.argv)}')
import sys, os.path, os, glob, datetime, re
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
#from misc.landmask import flagland
import xlinregress
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
#expname, hem, season = sys.argv[1:4] 
daname = 'ts'
expname = 'historical' #from sys.argv[1]
#timespan = slice('1958', '2014')
timespan = slice('1980', '2014')
idir = f'/tigress/wenchang/data/cmip6/variables/{daname}/{expname}/wy_regrid_all_members'
func_years = lambda ncfile: ncfile.split('.')[-2].split('-') # e.g. ts.historical.UKESM1-0-LL.r9i1p1f2.850hPa.1850-2014.nc -> ['1850', '2014']
ncfiles = [f for f in os.listdir(idir)
    if f.endswith('.nc')
    and func_years(f)[0]<=timespan.start
    and func_years(f)[1]>=timespan.stop ]
#sort by model/member
func_key = lambda ncfile: ncfile.split('.')[2:3] + [int(s) for s in re.split('[ipf]', ncfile.split('.')[3][1:])] #used in model/member sorting
ncfiles.sort(key=func_key)
print('**ncfiles**:')
for ncfile in ncfiles:
    print(ncfile)
print('**ncfiles done**')
mask_land = lambda da: da.where(flagland(da)<0.1)
print(idir)
n_files = len(ncfiles)

ofile = f'{daname}_{expname}_{n_files}ens_{timespan.start}-{timespan.stop}.nc'
if os.path.exists(ofile):
    print('[exists]:', ofile)
    sys.exit()
long_name = f'{expname} {daname} linear trend over {timespan.start}-{timespan.stop}'
units = 'K/year'

dss = []
models = []
members = []
model_members = []
for ii,ncfile in enumerate(ncfiles, start=1):
    model,member = ncfile.split('.')[2:4]
    ifile = os.path.join(idir, ncfile)
    print(f'{ii:3d} of {n_files:3d}: {ifile}')
    #process
    da = xr.open_dataarray(ifile)
    da = da.sel(time=timespan).load()
    if ii==1: #check data selection
        print('**check data selection**:')
        print(da) 
    #print(da)
    #da = da.pipe(mask_land)
    #if ii==1:  #check land mask
    #    da.isel(time=0).plot()
    #    plt.title(f'land mask check: {long_name}')
    #    plt.savefig(ofile.replace('.nc', '.landmask.png'))
    #da.isel(time=0).plot()
    da = da.groupby('time.year').mean('time')
    rg = da.linregress.on(da.year, dim='year') #linear trend
    #da.plot(ls='--')
    dss.append(rg)
    models.append(model)
    members.append(member)
    model_members.append(f'{model}_{member}')
print('concat...')
ds = xr.concat(dss, dim=pd.Index(model_members, name='model_member')) \
    .assign_coords(model=('model_member', models)) \
    .assign_coords(members=('model_member', members)) 
ds.slope.attrs['units'] = units
ds.slope.attrs['long_name'] = long_name
print('saving...')
ds.to_netcdf(ofile)
print('[saved]:', ofile)
 
if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    #plt.show()