#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Mar 17 23:58:39 EDT 2021
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifiles = 'data/*.ocean.nc'
daname = 'wt'
encoding = {daname: {'zlib':True, 'complevel':1}}
print(f'{daname} from {ifiles}...')
ofile = f'ocean.{daname}.nc'
if os.path.exists(ofile):
    print('[loaded]:', ofile)
    ds = xr.open_dataset(ofile)
else:    
    with xr.open_mfdataset(ifiles) as ds_:
        print('loading...')
        ds = ds_[[daname]].load()
        print('saving...')
        ds.to_netcdf(ofile, encoding=encoding) 
        print('[saved]:', ofile)

#time mean
print('time mean...')
ofile_tmean = ofile.replace('.nc', '.tmean.nc')
if os.path.exists(ofile_tmean):
    print('[loaded]:', ofile_tmean)
    ds_tmean = xr.open_dataset(ofile_tmean)
else:
    print('calculating...')
    ds_tmean = ds.mean('time').load()
    print('saving...')
    ds_tmean.to_netcdf(ofile_tmean, encoding=encoding)
    print('[saved]:', ofile_tmean)

#mclim
print('monthly climatology...')
ofile_mclim = ofile.replace('.nc', '.mclim.nc')
if os.path.exists(ofile_mclim):
    print('[loaded]:', ofile_mclim)
    ds_mclim = xr.open_dataset(ofile_mclim)
else:
    print('calculating...')
    ds_mclim = ds.groupby('time.month').mean('time').load()
    print('saving...')
    ds_mclim.to_netcdf(ofile_mclim, encoding=encoding)
    print('[saved]:', ofile_mclim)

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    da = ds_tmean[daname].assign_attrs(units='m/s')
    figsize = (6.4*2, 6.4*9/16*2)
    fig, axes = plt.subplots(2,2,figsize=figsize, dpi=100)
    fc = '0.3'#ax facecolor
    ax = axes[0,0]
    da.sel(sw_ocean=200, method='nearest').plot.contourf(ax=ax, cmap='RdBu_r', levels=21, robust=True)
    ax.axhline(0, color='gray', ls='--')
    ax.axhline(35, color='gray', ls='--')
    ax.axvline(-40, color='gray', ls='--')
    ax.set_facecolor(fc)

    #fig, ax = plt.subplots()
    ax = axes[1,1]
    da.sel(yt_ocean=0, method='nearest').plot.contourf(ax=ax, cmap='RdBu_r', levels=21, robust=True, yincrease=False)
    ax.set_ylim(1000, None)
    ax.set_facecolor(fc)
    
    #fig, ax = plt.subplots()
    ax = axes[0,1]
    da.sel(yt_ocean=35, method='nearest').plot.contourf(ax=ax, cmap='RdBu_r', levels=21, robust=True, yincrease=False)
    ax.set_ylim(1000, None)
    #ax.set_xlim(-100,-50)
    ax.set_facecolor(fc)

    #fig, ax = plt.subplots()
    ax = axes[1,0]
    da.sel(xt_ocean=-40, method='nearest').plot.contourf(ax=ax, cmap='RdBu_r', levels=21, robust=True, yincrease=False)
    ax.set_ylim(1000, None)
    ax.set_facecolor(fc)

    #savefig
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        figtype = 'png'
        figname = __file__.replace('.py', f'.{figtype}')
        #archive existing fig
        if os.path.exists(figname):
            mdate = tt.getmdate(figname)
            figname_archive = figname.replace(f'.{figtype}', f'.{figtype}.{mdate}')
            os.rename(figname, figname_archive)
            print('[renamed]:', figname, '->', figname_archive)
        plt.savefig(figname)
        print('[saved]:', figname)
    tt.check(f'**Done**')
    plt.show()