#!/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_u = xr.open_dataset('u10_xy_2000.nc')
    ds_v = xr.open_dataset('v10_xy_2000.nc')
    ds_u.load()
    ds_v.load()
    ds = ds_u
da = (ds_u.u10_xy**2 + ds_v.v10_xy**2)**0.5
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 = (-500,500)#(-200,200)
    ylim = (-500,500)#(-200,200)
    vmin = 0
    levels = 21
    figname = __file__.replace('.py', f'_{tt.today()}.png')
    fig, axes = plt.subplots(3, 2, figsize=(7,7.5))
    ax = axes[0,0]
    da.mean(['storm', 'stage']).assign_attrs(units=units, long_name=long_name).plot.contourf(levels=levels, ax=ax, vmin=vmin)
    ax.set_title('All')
    ax.grid(True)
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)

    ax = axes[0,1]
    da.where(ds.windmax>33).mean(['storm', 'stage']).assign_attrs(units=units, long_name=long_name).plot.contourf(levels=levels, ax=ax, vmin=vmin)
    ax.set_title('HU')
    ax.grid(True)
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)

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

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

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

    ax = axes[2,1]
    da.where(ds.lat<0).where(ds.windmax>33).mean(['storm', 'stage']).assign_attrs(units=units, long_name=long_name).plot.contourf(levels=levels, ax=ax, vmin=vmin)
    ax.set_title('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()