#!/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 True:#'ds' not in globals():
    ds = xr.open_dataset('test_windspeed_nn.nc')
    ds.load()
da = ds.windspeed_xy
units = 'm/s'
long_name = 'wind speed'
da = da.assign_attrs(units=units, long_name=long_name)
 
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]
    da.where(ds.tm>0).mean(['storm', 'stage']).assign_attrs(units=units, long_name=long_name).plot.contourf(levels=11, ax=ax)
    ax.set_title('WC')
    ax.grid(True)
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)
    ax.annotate('', xy=(-100,0), xytext=(0,0), arrowprops=dict(arrowstyle='->'))
    ax.text(-100, 0, 'TC moving left', ha='left', va='bottom')

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

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

    ax = axes[1,1]
    da.where(ds.tm>0).where(ds.lat<0).mean(['storm', 'stage']).assign_attrs(units=units, long_name=long_name).plot.contourf(levels=11, ax=ax)
    ax.set_title('WC SH')
    ax.grid(True)
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)

    ax = axes[2,0]
    da.where(ds.tm>0).where(ds.lat>0).where(ds.windmax>32.5).mean(['storm', 'stage']).assign_attrs(units=units, long_name=long_name).plot.contourf(levels=11, ax=ax)
    ax.set_title('WC NH HU')
    ax.grid(True)
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)

    ax = axes[2,1]
    da.where(ds.tm>0).where(ds.lat<0).where(ds.windmax>32.5).mean(['storm', 'stage']).assign_attrs(units=units, long_name=long_name).plot.contourf(levels=11, ax=ax)
    ax.set_title('WC SH HU')
    ax.grid(True)
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)
    
    plt.savefig(figname)
    print('[saved]:', figname)
    tt.check(f'**Done**')
    plt.show()