#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Dec 15 22:09:06 EST 2021
#noIO: exclude Indian Ocean (60-90E)
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
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
expname, hem, season = sys.argv[1:4] 
dsname = 'PI'
daname = 'vmax'
#expname = 'hist-nat' #from sys.argv[1]
units = 'm/s'
time_span = slice('1981', '2014')
years = slice(1982, 2014)
idir = f'/tigress/wenchang/data/cmip6/variables/{dsname}/{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]<=time_span.start
    and func_years(f)[1]>=time_span.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)
#for ncfile in ifiles:
#    print(ncfile)
#hem = 'NH' #from sys.argv[2]
if hem == 'NH':
    latmin, latmax = 10, 30
    lonmin, lonmax = 40, 360-20
    #lonmin, lonmax = 90, 360-20 #noIO
elif hem == 'SH':
    latmin, latmax = -30, -10
    lonmin, lonmax = 30, 360-150
    #lonmin, lonmax = 90, 360-150 #noIO
mask_land = lambda da: da.where(flagland(da)<0.1)
excludeIO = lambda da: da.where((da.lon>90)|(da.lon<60)) #noIO
#season = 'ASO' #from sys.argv[3]
if season == 'ASO':
    isSeason = lambda month: (month>=8)&(month<=10)
elif season == 'JJASO':
    isSeason = lambda month: (month>=6)&(month<=10)
elif season == 'FMA':
    isSeason = lambda month: (month>=2)&(month<=4)
elif season == 'DJFMA':
    isSeason = lambda month: (month==12)|((month>=1)&(month<=4))
print(idir)
print(hem)
print(season)
n_files = len(ncfiles)

ofile = f'{dsname}_{expname}_{hem}_{season}_ens{n_files}_{years.start}-{years.stop}_noIO.nc'
if os.path.exists(ofile):
    print('[exists]:', ofile)
    sys.exit()
long_name = f'{expname} {hem} {season} PI vmax'

tss = []
models = []
members = []
model_members = []
note = ''
for ii,ncfile in enumerate(ncfiles, start=1):
    model,member = ncfile.split('.')[2:4]
    year_end = ncfile.split('.')[4].split('-')[-1]
    print(f'{ii:3d} of {n_files:3d}: {daname} {expname} {model}, {member}, {ncfile}')
    if int(year_end) < 2017:
        note += f'{model}_{member} ends with {year_end};\n'
    ifile = os.path.join(idir, ncfile)
    da = xr.open_dataset(ifile)[daname]
    da = da.sel(time=isSeason(da.time.dt.month)) \
        .sel(time=time_span) \
        .sel(lon=slice(lonmin, lonmax)).sel(lat=slice(latmin, latmax))
    da = da.pipe(excludeIO) #noIO
    if ii==1: #check data selection
        print('**check data selection**:')
        print(da) 
    da.load()
    #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()
    ##spatial and seasonal mean
    #import geoxarray
    #_da = da.geo.fldmean() \
    #    .groupby('time.year').mean('time').sel(year=years)
    #_da.plot()
    da = da.weighted(np.cos(np.deg2rad(da.lat))).mean(['lat', 'lon']) \
    #seasonal mean
    if season == 'DJFMA':
        da['time'] = da.indexes['time'].shift(1, 'MS').values #shift time to move year0 Dec to year1 Jan
        if ii==1: print(f'**shifted time coords (1MS)**: {da.time.values}')
    da = da.groupby('time.year').mean('time').sel(year=years)
    #da.plot(ls='--')
    tss.append(da)
    models.append(model)
    members.append(member)
    model_members.append(f'{model}_{member}')
print('concat...')
ts = xr.concat(tss, dim=pd.Index(model_members, name='model_member')) \
    .assign_coords(model=('model_member', models)) \
    .assign_coords(members=('model_member', members)) \
    .assign_attrs(note=note) \
    .assign_attrs(long_name=long_name) \
    .assign_attrs(units=units)
ts.to_dataset().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()