#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Jul 25 14:43:30 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 geoxarray
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
daname = 'ts'
amip = True if os.path.basename(os.path.dirname(os.getcwd())) == 'highresSST' else False
model = os.path.basename(os.path.abspath(os.getcwd())) #model is in current dir, e.g. CNRM-CM6-1-HR 
members = [m for m in os.listdir(os.getcwd()) if os.path.isdir(m) and m.startswith('r') and not m.endswith('_drop')] 
members.sort()
for ii,member in enumerate(members, start=1):
    #print(f'{amip = }'); sys.exit()
    #highresSST
    print(f'[{ii:2d} of {len(members):2d}]:', daname, 'HighResMIP', 'amip' if amip else 'cmip', model, member)
    expname_hist = 'highresSST-present' if amip else 'hist-1950'
    idir = f'/tigress/wenchang/data/cmip6/variables/{daname}/{expname_hist}/{model}/{member}'
    ncfiles = sorted( [ncfile for ncfile in os.listdir(idir) if ncfile.endswith('.nc')] )
    ifiles = [os.path.join(idir, ncfile) for ncfile in ncfiles]
    #highresSST-future
    expname_future = 'highresSST-future' if amip else 'highres-future'
    idir = f'/tigress/wenchang/data/cmip6/variables/{daname}/{expname_future}/{model}/{member}'
    ncfiles = sorted( [ncfile for ncfile in os.listdir(idir) if ncfile.endswith('.nc')] )
    ifiles = ifiles + [os.path.join(idir, ncfile) for ncfile in ncfiles]

    
    ofile = os.path.basename(ifiles[0]) \
        .replace(ifiles[0].split('.')[-2].split('-')[-1], ifiles[-1].split('.')[-2].split('-')[-1])
    ofile = ofile.replace(f'{daname}_Amon', 'wwa_gmst_smoothed')
    ofile = os.path.join(member, ofile) #put data to the sub dir, e.g. r1i1p1f1/
    print(f'{ofile = }')
    #print(ofile); sys.exit()
    if os.path.exists(ofile):
        print('[exists]:', ofile)
        print()
        continue #sys.exit()
    
    #load data
    print('loading...')
    da = xr.open_mfdataset(ifiles)[daname].load()
    
    #
    print('calculating...')
    da = da.geo.fldmean().groupby('time.year').mean('time')
    #check if missing data
    year_start = da.year.values[0]
    year_stop = da.year.values[-1]
    if (year_stop - year_start + 1) != da.year.size:
        print('**missing data error**')
        with open(f'wy_{daname}_{member}.log', 'w') as f:
            f.write('**missing data of some years. please check: **\n')
            f.write(f'/tigress/wenchang/data/cmip6/variables/{daname}/{expname_hist}/{model}/{member}\n')
            f.write(f'/tigress/wenchang/data/cmip6/variables/{daname}/{expname_future}/{model}/{member}\n')
        continue #sys.exit()
    time = [datetime.datetime(year, 6, 30) for year in da.year.values]
    da = da.rename(year='time').assign_coords(time=time) \
        .pipe(lambda x: x - x.sel(time=slice('1951', '1980')).mean('time')) \
        .rolling(time=5, min_periods=1).mean()
    
    print('saving...')
    da.to_dataset(name=daname).to_netcdf(ofile, encoding={daname: {'_FillValue': None}})
    print('[saved]:', ofile)
    print()
 
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')
        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()