#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Sep 30 11:49:41 EDT 2024
if __name__ == '__main__':
    import sys,os
    try:
        from misc.timer import Timer
        tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
    except:
        pass
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
#from misc import get_kws_from_argv
#
if __name__ == '__main__':
    try:
        tt.check('end import')
    except:
        pass
#
#start from here
danames = ['total_ocean_heat', 
    'total_ocean_hflux_coupler',
    'pe_tot',
    'total_net_sfc_heating',
    'total_ocean_salt',
    ]
das = []
labels = []
for daname in danames:
    ifile = f'{daname}_FLOR_CTL1860_newdiag_tigercpu_intelmpi_18_576PE_0001-3860_ocean_scalar.annualmean.nc'
    label = ifile.split('_tigercpu_')[0]
    da = xr.open_dataarray(ifile)
    labels.append(label)
    das.append(da)
units = da.attrs['units']
long_name = da.attrs['long_name']
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    import xfilter
    lp = lambda da: da.filter.lowpass(1/100, dim='year', padtype='even')#.assign_attrs(units=da.attrs['units'], long_name=da.attrs['long_name'])
    for da,label,daname in zip(das, labels, danames):
        fig,ax = plt.subplots()
        da.groupby('time.year').mean('time', keep_attrs=True).plot(label=label, alpha=0.5)
        da.groupby('time.year').mean('time', keep_attrs=True).pipe(lp).plot(label='lowpass')

        ax.legend()
        ax.set_title('')
        #ax.set_ylabel(f'{long_name} [{units}]')
    
        #savefig
        if 'savefig' in sys.argv or 's' in sys.argv:
            figname = __file__.replace('.py', f'__{daname}.png')
            if 'overwritefig' in sys.argv or 'o' in sys.argv:
                wysavefig(figname, overwritefig=True)
            else:
                wysavefig(figname)
 
    try:
        tt.check(f'**Done**')
    except:
        pass
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()