#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Feb 18 04:26:56 PM EST 2025
if __name__ == '__main__':
    import sys,os
    try:
        from misc.timer import Timer
        tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
    except:
        pass
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
#from misc import get_kws_from_argv
#
if __name__ == '__main__':
    try:
        tt.check('end import')
    except:
        pass
#
#start from here
print('u200...')
ifiles = '/projects/w/wenchang/data/era5/analysis_wy/plevels/u200/daily/era5.u200.daily.*.nc'
u200 = xr.open_mfdataset(ifiles)['u']

print('u850...')
ifiles = '/projects/w/wenchang/data/era5/analysis_wy/plevels/u850/daily/era5.u850.daily.*.nc'
u850 = xr.open_mfdataset(ifiles)['u']
 
print('v200...')
ifiles = '/projects/w/wenchang/data/era5/analysis_wy/plevels/v200/daily/era5.v200.daily.*.nc'
v200 = xr.open_mfdataset(ifiles)['v']

print('v850...')
ifiles = '/projects/w/wenchang/data/era5/analysis_wy/plevels/v850/daily/era5.v850.daily.*.nc'
v850 = xr.open_mfdataset(ifiles)['v']

#time-mean
ofile = f'era5.wind_shear.daily.1979-2020.phase8.nc'
time = ['1979-07-07T00:00:00',
 '1979-07-08T00:00:00',
 '1979-07-09T00:00:00',
 '1979-07-10T00:00:00',
 '1979-07-11T00:00:00',
 '1979-07-15T00:00:00',
 '1979-07-17T00:00:00',
 '1979-07-18T00:00:00',
 '1979-08-18T00:00:00',
 '1979-08-19T00:00:00',
 '1979-08-20T00:00:00',
 '1979-08-21T00:00:00',
 '1979-10-18T00:00:00',
 '1979-10-19T00:00:00',
 '1979-10-20T00:00:00',
 '1979-10-21T00:00:00',
 '1979-10-22T00:00:00',
 '1981-10-04T00:00:00',
 '1981-10-05T00:00:00',
 '1981-10-06T00:00:00',
 '1982-07-24T00:00:00',
 '1982-07-25T00:00:00',
 '1982-11-17T00:00:00',
 '1982-11-18T00:00:00',
 '1982-11-19T00:00:00',
 '1982-11-20T00:00:00',
 '1982-11-21T00:00:00',
 '1984-07-06T00:00:00',
 '1984-07-07T00:00:00',
 '1984-09-08T00:00:00',
 '1984-09-09T00:00:00',
 '1984-09-10T00:00:00',
 '1984-09-11T00:00:00',
 '1984-09-12T00:00:00',
 '1984-09-13T00:00:00',
 '1984-09-17T00:00:00',
 '1984-09-18T00:00:00',
 '1984-09-19T00:00:00',
 '1984-09-20T00:00:00',
 '1984-09-21T00:00:00',
 '1984-10-24T00:00:00',
 '1984-10-27T00:00:00',
 '1984-10-28T00:00:00',
 '1984-10-29T00:00:00',
 '1984-10-30T00:00:00',
 '1984-10-31T00:00:00',
 '1985-09-21T00:00:00',
 '1985-09-22T00:00:00',
 '1985-09-23T00:00:00',
 '1985-09-24T00:00:00',
 '1985-10-23T00:00:00',
 '1985-10-24T00:00:00',
 '1985-10-25T00:00:00',
 '1985-10-26T00:00:00',
 '1985-10-27T00:00:00',
 '1985-10-28T00:00:00',
 '1985-10-29T00:00:00',
 '1986-07-05T00:00:00',
 '1986-09-23T00:00:00',
 '1986-09-24T00:00:00',
 '1986-11-21T00:00:00',
 '1986-11-22T00:00:00',
 '1986-11-23T00:00:00',
 '1986-11-24T00:00:00',
 '1988-06-17T00:00:00',
 '1988-06-18T00:00:00',
 '1988-06-23T00:00:00',
 '1988-06-24T00:00:00',
 '1989-06-22T00:00:00',
 '1989-06-23T00:00:00',
 '1989-07-26T00:00:00',
 '1989-07-27T00:00:00',
 '1989-07-28T00:00:00',
 '1989-07-29T00:00:00',
 '1989-11-08T00:00:00',
 '1989-11-09T00:00:00',
 '1989-11-30T00:00:00',
 '1990-06-28T00:00:00',
 '1990-06-29T00:00:00',
 '1990-06-30T00:00:00',
 '1990-07-01T00:00:00',
 '1990-11-20T00:00:00',
 '1990-11-21T00:00:00',
 '1990-11-22T00:00:00',
 '1990-11-23T00:00:00',
 '1990-11-24T00:00:00',
 '1990-11-25T00:00:00',
 '1991-06-21T00:00:00',
 '1991-06-22T00:00:00',
 '1991-06-23T00:00:00',
 '1991-06-24T00:00:00',
 '1991-09-17T00:00:00',
 '1991-09-18T00:00:00',
 '1991-09-19T00:00:00',
 '1991-09-20T00:00:00',
 '1991-09-21T00:00:00',
 '1991-10-01T00:00:00',
 '1991-10-02T00:00:00',
 '1991-11-13T00:00:00',
 '1991-11-14T00:00:00',
 '1991-11-15T00:00:00',
 '1991-11-16T00:00:00',
 '1991-11-17T00:00:00',
 '1991-11-18T00:00:00',
 '1992-09-14T00:00:00',
 '1992-09-15T00:00:00',
 '1992-09-17T00:00:00',
 '1992-09-18T00:00:00',
 '1992-09-19T00:00:00',
 '1992-11-04T00:00:00',
 '1992-11-05T00:00:00',
 '1992-11-06T00:00:00',
 '1992-11-07T00:00:00',
 '1992-11-08T00:00:00',
 '1992-11-24T00:00:00',
 '1993-06-15T00:00:00',
 '1993-06-16T00:00:00',
 '1993-06-17T00:00:00',
 '1993-06-18T00:00:00',
 '1993-06-19T00:00:00',
 '1993-06-20T00:00:00',
 '1993-06-21T00:00:00',
 '1993-06-22T00:00:00',
 '1993-06-23T00:00:00',
 '1993-10-25T00:00:00',
 '1993-10-26T00:00:00',
 '1993-10-27T00:00:00',
 '1993-10-28T00:00:00',
 '1993-10-29T00:00:00',
 '1994-06-22T00:00:00',
 '1994-09-23T00:00:00',
 '1994-09-24T00:00:00',
 '1994-09-25T00:00:00',
 '1994-09-26T00:00:00',
 '1994-09-27T00:00:00',
 '1994-09-28T00:00:00',
 '1994-09-29T00:00:00',
 '1994-09-30T00:00:00',
 '1994-10-01T00:00:00',
 '1994-10-02T00:00:00',
 '1996-06-23T00:00:00',
 '1996-06-24T00:00:00',
 '1996-06-25T00:00:00',
 '1996-06-26T00:00:00',
 '1996-06-27T00:00:00',
 '1996-06-28T00:00:00',
 '1996-10-09T00:00:00',
 '1996-11-13T00:00:00',
 '1996-11-14T00:00:00',
 '1996-11-15T00:00:00',
 '1996-11-16T00:00:00',
 '1996-11-17T00:00:00',
 '1996-11-18T00:00:00',
 '1997-06-02T00:00:00',
 '1997-06-03T00:00:00',
 '1997-06-04T00:00:00',
 '1997-06-05T00:00:00',
 '1997-06-06T00:00:00',
 '1997-06-07T00:00:00',
 '1997-06-08T00:00:00',
 '1997-06-09T00:00:00',
 '1997-06-10T00:00:00',
 '1997-06-11T00:00:00',
 '1997-10-10T00:00:00',
 '1997-10-11T00:00:00',
 '1997-10-12T00:00:00',
 '1998-08-28T00:00:00',
 '1998-08-29T00:00:00',
 '1998-08-30T00:00:00',
 '1998-08-31T00:00:00',
 '1998-09-01T00:00:00',
 '1998-09-02T00:00:00',
 '1998-09-03T00:00:00',
 '1998-09-04T00:00:00',
 '1998-09-05T00:00:00',
 '1998-09-06T00:00:00',
 '1998-09-07T00:00:00',
 '1998-09-08T00:00:00',
 '1998-09-09T00:00:00',
 '1998-09-10T00:00:00',
 '1999-06-16T00:00:00',
 '1999-06-17T00:00:00',
 '1999-06-18T00:00:00',
 '1999-06-19T00:00:00',
 '1999-06-20T00:00:00',
 '1999-06-21T00:00:00',
 '1999-06-22T00:00:00',
 '1999-06-23T00:00:00',
 '1999-06-24T00:00:00',
 '1999-06-25T00:00:00',
 '1999-06-26T00:00:00',
 '1999-06-27T00:00:00',
 '1999-09-08T00:00:00',
 '1999-11-15T00:00:00',
 '1999-11-16T00:00:00',
 '2000-06-08T00:00:00',
 '2000-06-14T00:00:00',
 '2000-06-15T00:00:00',
 '2000-06-16T00:00:00',
 '2000-06-17T00:00:00',
 '2000-07-17T00:00:00',
 '2000-07-18T00:00:00',
 '2000-07-19T00:00:00',
 '2000-07-20T00:00:00',
 '2000-07-21T00:00:00',
 '2000-07-22T00:00:00',
 '2000-07-23T00:00:00',
 '2000-07-24T00:00:00',
 '2000-07-25T00:00:00',
 '2000-09-02T00:00:00',
 '2000-09-03T00:00:00',
 '2000-09-04T00:00:00',
 '2000-09-05T00:00:00',
 '2000-09-06T00:00:00',
 '2000-09-07T00:00:00',
 '2000-09-08T00:00:00',
 '2000-09-09T00:00:00',
 '2000-09-10T00:00:00',
 '2000-09-12T00:00:00',
 '2000-11-07T00:00:00',
 '2000-11-08T00:00:00',
 '2000-11-09T00:00:00',
 '2001-06-13T00:00:00',
 '2001-07-11T00:00:00',
 '2001-07-12T00:00:00',
 '2001-07-13T00:00:00',
 '2001-07-14T00:00:00',
 '2001-07-15T00:00:00',
 '2001-09-07T00:00:00',
 '2001-10-21T00:00:00',
 '2001-10-22T00:00:00',
 '2001-10-23T00:00:00',
 '2001-10-24T00:00:00',
 '2001-10-25T00:00:00',
 '2001-10-26T00:00:00',
 '2002-07-14T00:00:00',
 '2002-07-15T00:00:00',
 '2002-07-16T00:00:00',
 '2002-07-17T00:00:00',
 '2002-07-18T00:00:00',
 '2002-07-19T00:00:00',
 '2002-07-20T00:00:00',
 '2002-08-22T00:00:00',
 '2002-08-23T00:00:00',
 '2002-08-24T00:00:00',
 '2002-08-25T00:00:00',
 '2002-08-26T00:00:00',
 '2002-08-27T00:00:00',
 '2002-08-28T00:00:00',
 '2002-08-29T00:00:00',
 '2002-08-30T00:00:00',
 '2003-06-01T00:00:00',
 '2003-06-02T00:00:00',
 '2003-06-03T00:00:00',
 '2003-06-04T00:00:00',
 '2003-06-05T00:00:00',
 '2003-06-06T00:00:00',
 '2003-06-07T00:00:00',
 '2003-06-08T00:00:00',
 '2003-06-09T00:00:00',
 '2003-10-27T00:00:00',
 '2003-10-28T00:00:00',
 '2003-10-29T00:00:00',
 '2003-10-30T00:00:00',
 '2003-11-18T00:00:00',
 '2004-07-01T00:00:00',
 '2004-07-02T00:00:00',
 '2004-07-03T00:00:00',
 '2004-07-04T00:00:00',
 '2004-07-05T00:00:00',
 '2004-08-21T00:00:00',
 '2004-08-22T00:00:00',
 '2004-08-23T00:00:00',
 '2004-08-24T00:00:00',
 '2004-08-25T00:00:00',
 '2004-08-26T00:00:00',
 '2004-08-28T00:00:00',
 '2004-10-18T00:00:00',
 '2004-10-19T00:00:00',
 '2004-10-20T00:00:00',
 '2005-09-21T00:00:00',
 '2005-09-22T00:00:00',
 '2005-09-23T00:00:00',
 '2005-09-24T00:00:00',
 '2005-11-18T00:00:00',
 '2005-11-19T00:00:00',
 '2005-11-20T00:00:00',
 '2005-11-21T00:00:00',
 '2005-11-22T00:00:00',
 '2006-06-01T00:00:00',
 '2006-06-02T00:00:00',
 '2006-07-13T00:00:00',
 '2006-07-14T00:00:00',
 '2006-07-15T00:00:00',
 '2006-07-17T00:00:00',
 '2006-07-18T00:00:00',
 '2006-07-19T00:00:00',
 '2006-07-20T00:00:00',
 '2006-08-20T00:00:00',
 '2006-10-12T00:00:00',
 '2006-10-13T00:00:00',
 '2006-10-14T00:00:00',
 '2006-10-15T00:00:00',
 '2006-10-16T00:00:00',
 '2006-10-17T00:00:00',
 '2006-11-09T00:00:00',
 '2006-11-10T00:00:00',
 '2006-11-11T00:00:00',
 '2007-07-13T00:00:00',
 '2007-07-14T00:00:00',
 '2007-10-08T00:00:00',
 '2007-10-09T00:00:00',
 '2007-10-10T00:00:00',
 '2007-10-11T00:00:00',
 '2007-10-12T00:00:00',
 '2007-10-13T00:00:00',
 '2007-10-14T00:00:00',
 '2007-10-15T00:00:00',
 '2007-10-16T00:00:00',
 '2007-10-17T00:00:00',
 '2007-11-30T00:00:00',
 '2009-06-11T00:00:00',
 '2009-06-12T00:00:00',
 '2009-06-13T00:00:00',
 '2009-06-14T00:00:00',
 '2009-06-15T00:00:00',
 '2009-06-16T00:00:00',
 '2009-06-17T00:00:00',
 '2009-06-18T00:00:00',
 '2009-06-19T00:00:00',
 '2009-06-20T00:00:00',
 '2009-06-21T00:00:00',
 '2009-06-22T00:00:00',
 '2009-06-23T00:00:00',
 '2009-10-26T00:00:00',
 '2009-10-27T00:00:00',
 '2009-10-28T00:00:00',
 '2010-07-10T00:00:00',
 '2010-07-11T00:00:00',
 '2010-11-13T00:00:00',
 '2010-11-14T00:00:00',
 '2010-11-15T00:00:00',
 '2011-07-09T00:00:00',
 '2011-10-10T00:00:00',
 '2011-10-11T00:00:00',
 '2011-10-12T00:00:00',
 '2011-10-13T00:00:00',
 '2011-10-14T00:00:00',
 '2011-11-14T00:00:00',
 '2011-11-15T00:00:00',
 '2012-06-15T00:00:00',
 '2012-06-16T00:00:00',
 '2012-06-22T00:00:00',
 '2012-08-07T00:00:00',
 '2012-08-08T00:00:00',
 '2013-06-24T00:00:00',
 '2013-06-25T00:00:00',
 '2013-06-26T00:00:00',
 '2013-06-27T00:00:00',
 '2013-06-28T00:00:00',
 '2013-08-18T00:00:00',
 '2013-08-19T00:00:00',
 '2013-08-24T00:00:00',
 '2013-08-25T00:00:00',
 '2013-10-30T00:00:00',
 '2013-10-31T00:00:00',
 '2014-10-13T00:00:00',
 '2014-10-14T00:00:00',
 '2014-11-08T00:00:00',
 '2014-11-09T00:00:00',
 '2014-11-10T00:00:00',
 '2015-07-12T00:00:00',
 '2015-07-13T00:00:00',
 '2015-07-14T00:00:00',
 '2015-07-15T00:00:00',
 '2015-10-01T00:00:00',
 '2015-10-02T00:00:00',
 '2015-10-03T00:00:00',
 '2015-10-04T00:00:00',
 '2016-11-12T00:00:00',
 '2016-11-13T00:00:00',
 '2016-11-14T00:00:00',
 '2016-11-15T00:00:00',
 '2017-10-30T00:00:00',
 '2018-09-05T00:00:00',
 '2018-09-06T00:00:00',
 '2018-09-07T00:00:00',
 '2018-09-08T00:00:00',
 '2018-09-09T00:00:00',
 '2018-09-10T00:00:00',
 '2018-09-11T00:00:00',
 '2018-09-22T00:00:00',
 '2018-09-23T00:00:00',
 '2018-09-24T00:00:00',
 '2018-09-25T00:00:00',
 '2018-09-26T00:00:00',
 '2018-09-27T00:00:00',
 '2018-09-28T00:00:00',
 '2018-09-29T00:00:00',
 '2018-09-30T00:00:00',
 '2018-10-01T00:00:00',
 '2018-10-29T00:00:00',
 '2018-10-30T00:00:00',
 '2018-11-28T00:00:00',
 '2018-11-29T00:00:00',
 '2018-11-30T00:00:00',
 '2019-09-15T00:00:00',
 '2019-09-16T00:00:00',
 '2019-09-17T00:00:00',
 '2019-09-18T00:00:00',
 '2019-09-19T00:00:00',
 '2019-09-20T00:00:00',
 '2019-10-09T00:00:00',
 '2019-10-10T00:00:00',
 '2019-10-11T00:00:00',
 '2019-10-12T00:00:00',
 '2019-11-13T00:00:00',
 '2019-11-14T00:00:00',
 '2019-11-15T00:00:00',
 '2019-11-16T00:00:00',
 '2019-11-17T00:00:00',
 '2019-11-18T00:00:00',
 '2019-11-19T00:00:00',
 '2020-08-15T00:00:00',
 '2020-08-16T00:00:00',
 '2020-08-17T00:00:00',
 '2020-08-18T00:00:00',
 '2020-08-19T00:00:00',
 '2020-08-20T00:00:00',
 '2020-08-21T00:00:00',
 '2020-11-07T00:00:00',
 '2020-11-10T00:00:00',
 '2020-11-11T00:00:00']
if os.path.exists(ofile): 
    print('[exists]:', ofile)
else:
    print('wind_shear...')
    wind_shear = ((u200-u850)**2 + (v200 - v850)**2)**0.5
    print('phase...')
    wind_shear = wind_shear.sel(time=time).load()
    print('wind_shear_mean...')
    wind_shear_mean = wind_shear.mean('time')
    print('wind_shear_mean save...')
    wind_shear_mean.to_dataset(name='wind_shear_mean').to_netcdf(ofile)
    print('[saved]:', ofile)

 
if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
 
    try:
        tt.check(f'**Done**')
    except:
        pass
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()