#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Oct  2 17:49:43 EDT 2020
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#more imports
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
if 'ds' not in globals():
    ds = xr.open_dataset('test_slp_nn.nc')
    ds.load()
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    figname = __file__.replace('.py', f'_{tt.today()}.png')
    fig, axes = plt.subplots(3, 2, figsize=(7,7.5))
    ax = axes[0,0]
    ds.slp_xy.pipe(lambda da: da-da.sel(x=0,y=0)).mean(['storm', 'stage']).plot.contourf(levels=11, ax=ax)
    ax.set_title('All')

    ax = axes[0,1]
    ds.slp_xy.where(ds.windmax>32.5).pipe(lambda da: da-da.sel(x=0,y=0)).mean(['storm', 'stage']).plot.contourf(levels=11, ax=ax)
    ax.set_title('windmax > 32.5')

    ax = axes[1,0]
    ds.slp_xy.where(ds.tm>0).pipe(lambda da: da-da.sel(x=0,y=0)).mean(['storm', 'stage']).plot.contourf(levels=11, ax=ax)
    ax.set_title('WC')

    ax = axes[1,1]
    ds.slp_xy.where(ds.tm<0).pipe(lambda da: da-da.sel(x=0,y=0)).mean(['storm', 'stage']).plot.contourf(levels=11, ax=ax)
    ax.set_title('non-WC')

    ax = axes[2,0]
    ds.slp_xy.where(ds.lat>0).pipe(lambda da: da-da.sel(x=0,y=0)).mean(['storm', 'stage']).plot.contourf(levels=11, ax=ax)
    ax.set_title('NH')

    ax = axes[2,1]
    ds.slp_xy.where(ds.lat<0).pipe(lambda da: da-da.sel(x=0,y=0)).mean(['storm', 'stage']).plot.contourf(levels=11, ax=ax)
    ax.set_title('SH')

    
    plt.savefig(figname)
    tt.check(f'**Done**')
    #plt.show()