#!/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_precip.nc')
    ds.load()
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    xlim = (-200,200)
    ylim = (-200,200)
    figname = __file__.replace('.py', f'_{tt.today()}.png')
    fig, axes = plt.subplots(3, 2, figsize=(7,7.5))
    ax = axes[0,0]
    ds.precip_xy.mean(['storm', 'stage']).plot.contourf(levels=11, ax=ax)
    ax.set_title('All')
    ax.grid(True)
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)

    ax = axes[0,1]
    ds.precip_xy.where(ds.windmax>32.5).mean(['storm', 'stage']).plot.contourf(levels=11, ax=ax)
    ax.set_title('windmax > 32.5')
    ax.grid(True)
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)

    ax = axes[1,0]
    ds.precip_xy.where(ds.tm>0).mean(['storm', 'stage']).plot.contourf(levels=11, ax=ax)
    ax.set_title('WC')
    ax.grid(True)
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)

    ax = axes[1,1]
    ds.precip_xy.where(ds.tm<0).mean(['storm', 'stage']).plot.contourf(levels=11, ax=ax)
    ax.set_title('non-WC')
    ax.grid(True)
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)

    ax = axes[2,0]
    ds.precip_xy.where(ds.lat>0).mean(['storm', 'stage']).plot.contourf(levels=11, ax=ax)
    ax.set_title('NH')
    ax.grid(True)
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)

    ax = axes[2,1]
    ds.precip_xy.where(ds.lat<0).mean(['storm', 'stage']).plot.contourf(levels=11, ax=ax)
    ax.set_title('SH')
    ax.grid(True)
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)

    
    plt.savefig(figname)
    print('[saved]:', figname)
    tt.check(f'**Done**')
    plt.show()