#!/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.phase3.nc'
time = ['1979-06-15T00:00:00',
 '1979-06-16T00:00:00',
 '1979-06-17T00:00:00',
 '1979-06-18T00:00:00',
 '1979-06-19T00:00:00',
 '1979-09-19T00:00:00',
 '1979-09-20T00:00:00',
 '1979-09-21T00:00:00',
 '1979-11-05T00:00:00',
 '1979-11-10T00:00:00',
 '1979-11-17T00:00:00',
 '1979-11-18T00:00:00',
 '1979-11-19T00:00:00',
 '1979-11-20T00:00:00',
 '1979-11-21T00:00:00',
 '1979-11-22T00:00:00',
 '1979-11-23T00:00:00',
 '1979-11-24T00:00:00',
 '1979-11-25T00:00:00',
 '1980-08-19T00:00:00',
 '1980-08-20T00:00:00',
 '1980-10-13T00:00:00',
 '1980-10-14T00:00:00',
 '1980-10-15T00:00:00',
 '1981-07-05T00:00:00',
 '1981-07-07T00:00:00',
 '1981-07-08T00:00:00',
 '1981-07-09T00:00:00',
 '1981-07-10T00:00:00',
 '1981-07-11T00:00:00',
 '1981-07-12T00:00:00',
 '1981-07-13T00:00:00',
 '1981-07-22T00:00:00',
 '1981-07-23T00:00:00',
 '1981-07-24T00:00:00',
 '1981-07-25T00:00:00',
 '1981-07-26T00:00:00',
 '1981-09-14T00:00:00',
 '1981-09-15T00:00:00',
 '1981-10-30T00:00:00',
 '1981-11-02T00:00:00',
 '1981-11-09T00:00:00',
 '1982-08-05T00:00:00',
 '1982-08-06T00:00:00',
 '1982-08-09T00:00:00',
 '1982-10-18T00:00:00',
 '1982-11-05T00:00:00',
 '1982-11-27T00:00:00',
 '1982-11-28T00:00:00',
 '1982-11-29T00:00:00',
 '1982-11-30T00:00:00',
 '1983-06-14T00:00:00',
 '1983-06-15T00:00:00',
 '1983-06-16T00:00:00',
 '1983-07-12T00:00:00',
 '1983-07-13T00:00:00',
 '1983-07-14T00:00:00',
 '1983-07-15T00:00:00',
 '1983-07-16T00:00:00',
 '1983-07-17T00:00:00',
 '1983-07-18T00:00:00',
 '1983-07-19T00:00:00',
 '1983-07-20T00:00:00',
 '1983-08-04T00:00:00',
 '1983-08-05T00:00:00',
 '1983-08-06T00:00:00',
 '1983-08-07T00:00:00',
 '1983-08-08T00:00:00',
 '1983-08-09T00:00:00',
 '1983-08-10T00:00:00',
 '1983-09-15T00:00:00',
 '1983-09-16T00:00:00',
 '1983-09-17T00:00:00',
 '1984-06-12T00:00:00',
 '1984-11-17T00:00:00',
 '1984-11-18T00:00:00',
 '1984-11-19T00:00:00',
 '1984-11-20T00:00:00',
 '1984-11-21T00:00:00',
 '1984-11-22T00:00:00',
 '1984-11-23T00:00:00',
 '1984-11-29T00:00:00',
 '1984-11-30T00:00:00',
 '1985-11-06T00:00:00',
 '1985-11-07T00:00:00',
 '1985-11-08T00:00:00',
 '1985-11-09T00:00:00',
 '1986-06-10T00:00:00',
 '1986-06-11T00:00:00',
 '1986-06-12T00:00:00',
 '1986-09-05T00:00:00',
 '1986-09-06T00:00:00',
 '1986-09-07T00:00:00',
 '1986-09-08T00:00:00',
 '1986-10-27T00:00:00',
 '1986-10-28T00:00:00',
 '1986-10-29T00:00:00',
 '1986-10-30T00:00:00',
 '1986-10-31T00:00:00',
 '1986-11-01T00:00:00',
 '1986-11-02T00:00:00',
 '1986-11-03T00:00:00',
 '1987-06-01T00:00:00',
 '1987-06-02T00:00:00',
 '1987-06-03T00:00:00',
 '1987-06-04T00:00:00',
 '1987-08-12T00:00:00',
 '1987-08-13T00:00:00',
 '1987-08-14T00:00:00',
 '1987-08-15T00:00:00',
 '1987-08-16T00:00:00',
 '1987-08-17T00:00:00',
 '1987-08-18T00:00:00',
 '1987-09-23T00:00:00',
 '1987-09-24T00:00:00',
 '1987-09-26T00:00:00',
 '1987-09-27T00:00:00',
 '1987-09-28T00:00:00',
 '1987-11-02T00:00:00',
 '1987-11-03T00:00:00',
 '1987-11-04T00:00:00',
 '1988-09-17T00:00:00',
 '1988-09-18T00:00:00',
 '1990-09-19T00:00:00',
 '1990-09-22T00:00:00',
 '1990-09-25T00:00:00',
 '1990-09-26T00:00:00',
 '1990-10-31T00:00:00',
 '1990-11-01T00:00:00',
 '1990-11-02T00:00:00',
 '1990-11-03T00:00:00',
 '1990-11-04T00:00:00',
 '1991-06-01T00:00:00',
 '1991-06-02T00:00:00',
 '1991-10-12T00:00:00',
 '1991-10-13T00:00:00',
 '1991-10-14T00:00:00',
 '1991-10-15T00:00:00',
 '1991-10-16T00:00:00',
 '1991-10-17T00:00:00',
 '1991-10-18T00:00:00',
 '1991-10-19T00:00:00',
 '1991-10-20T00:00:00',
 '1992-06-06T00:00:00',
 '1992-06-07T00:00:00',
 '1992-06-08T00:00:00',
 '1992-06-09T00:00:00',
 '1992-06-10T00:00:00',
 '1992-06-11T00:00:00',
 '1992-07-15T00:00:00',
 '1992-07-16T00:00:00',
 '1992-10-05T00:00:00',
 '1992-10-08T00:00:00',
 '1993-06-01T00:00:00',
 '1993-06-02T00:00:00',
 '1993-06-03T00:00:00',
 '1993-11-08T00:00:00',
 '1993-11-09T00:00:00',
 '1993-11-26T00:00:00',
 '1993-11-29T00:00:00',
 '1993-11-30T00:00:00',
 '1994-06-06T00:00:00',
 '1994-06-07T00:00:00',
 '1994-08-27T00:00:00',
 '1994-10-24T00:00:00',
 '1994-10-25T00:00:00',
 '1994-10-26T00:00:00',
 '1994-10-27T00:00:00',
 '1994-11-25T00:00:00',
 '1994-11-26T00:00:00',
 '1994-11-27T00:00:00',
 '1994-11-28T00:00:00',
 '1994-11-29T00:00:00',
 '1995-07-05T00:00:00',
 '1995-07-06T00:00:00',
 '1995-07-07T00:00:00',
 '1995-08-26T00:00:00',
 '1995-08-27T00:00:00',
 '1995-08-28T00:00:00',
 '1996-06-06T00:00:00',
 '1996-06-07T00:00:00',
 '1996-06-08T00:00:00',
 '1996-06-09T00:00:00',
 '1996-06-10T00:00:00',
 '1996-06-11T00:00:00',
 '1996-06-12T00:00:00',
 '1996-08-18T00:00:00',
 '1996-08-19T00:00:00',
 '1996-08-20T00:00:00',
 '1996-08-21T00:00:00',
 '1996-08-24T00:00:00',
 '1996-08-25T00:00:00',
 '1996-08-26T00:00:00',
 '1996-08-27T00:00:00',
 '1997-06-29T00:00:00',
 '1997-06-30T00:00:00',
 '1997-07-01T00:00:00',
 '1997-07-02T00:00:00',
 '1997-07-03T00:00:00',
 '1997-07-04T00:00:00',
 '1997-07-05T00:00:00',
 '1997-07-06T00:00:00',
 '1997-07-07T00:00:00',
 '1997-07-08T00:00:00',
 '1997-07-13T00:00:00',
 '1997-07-14T00:00:00',
 '1997-09-27T00:00:00',
 '1997-09-28T00:00:00',
 '1997-11-16T00:00:00',
 '1997-11-17T00:00:00',
 '1997-11-18T00:00:00',
 '1997-11-19T00:00:00',
 '1997-11-20T00:00:00',
 '1998-06-22T00:00:00',
 '1998-06-23T00:00:00',
 '1998-06-24T00:00:00',
 '1998-06-25T00:00:00',
 '1998-06-26T00:00:00',
 '1998-06-27T00:00:00',
 '1998-06-28T00:00:00',
 '1998-09-18T00:00:00',
 '1998-09-19T00:00:00',
 '1998-11-05T00:00:00',
 '1999-08-27T00:00:00',
 '1999-10-17T00:00:00',
 '1999-10-18T00:00:00',
 '1999-10-19T00:00:00',
 '2000-10-03T00:00:00',
 '2000-10-04T00:00:00',
 '2000-10-05T00:00:00',
 '2000-11-17T00:00:00',
 '2000-11-18T00:00:00',
 '2000-11-19T00:00:00',
 '2000-11-20T00:00:00',
 '2000-11-21T00:00:00',
 '2000-11-22T00:00:00',
 '2001-08-03T00:00:00',
 '2001-08-04T00:00:00',
 '2001-08-05T00:00:00',
 '2001-09-29T00:00:00',
 '2001-09-30T00:00:00',
 '2001-10-01T00:00:00',
 '2001-10-02T00:00:00',
 '2001-11-01T00:00:00',
 '2001-11-02T00:00:00',
 '2002-06-13T00:00:00',
 '2002-06-14T00:00:00',
 '2002-06-15T00:00:00',
 '2002-08-03T00:00:00',
 '2002-08-04T00:00:00',
 '2002-08-05T00:00:00',
 '2002-08-06T00:00:00',
 '2002-11-07T00:00:00',
 '2002-11-08T00:00:00',
 '2002-11-09T00:00:00',
 '2002-11-10T00:00:00',
 '2002-11-11T00:00:00',
 '2002-11-12T00:00:00',
 '2002-11-13T00:00:00',
 '2002-11-14T00:00:00',
 '2002-11-15T00:00:00',
 '2002-11-16T00:00:00',
 '2002-11-17T00:00:00',
 '2002-11-18T00:00:00',
 '2003-10-06T00:00:00',
 '2003-11-17T00:00:00',
 '2004-09-24T00:00:00',
 '2004-09-25T00:00:00',
 '2004-09-26T00:00:00',
 '2004-09-27T00:00:00',
 '2004-09-28T00:00:00',
 '2004-09-29T00:00:00',
 '2004-10-30T00:00:00',
 '2004-10-31T00:00:00',
 '2004-11-01T00:00:00',
 '2004-11-02T00:00:00',
 '2005-06-20T00:00:00',
 '2005-09-02T00:00:00',
 '2005-10-14T00:00:00',
 '2006-09-10T00:00:00',
 '2006-09-11T00:00:00',
 '2006-09-12T00:00:00',
 '2006-09-13T00:00:00',
 '2006-09-14T00:00:00',
 '2006-09-15T00:00:00',
 '2006-09-16T00:00:00',
 '2006-09-17T00:00:00',
 '2006-09-18T00:00:00',
 '2006-09-19T00:00:00',
 '2006-09-20T00:00:00',
 '2006-11-22T00:00:00',
 '2006-11-23T00:00:00',
 '2006-11-24T00:00:00',
 '2006-11-25T00:00:00',
 '2007-06-15T00:00:00',
 '2007-06-16T00:00:00',
 '2008-06-05T00:00:00',
 '2008-06-06T00:00:00',
 '2008-06-07T00:00:00',
 '2008-06-08T00:00:00',
 '2008-06-09T00:00:00',
 '2008-07-19T00:00:00',
 '2008-07-20T00:00:00',
 '2008-07-21T00:00:00',
 '2008-09-03T00:00:00',
 '2008-09-04T00:00:00',
 '2008-09-05T00:00:00',
 '2008-10-15T00:00:00',
 '2008-10-16T00:00:00',
 '2008-10-17T00:00:00',
 '2008-10-18T00:00:00',
 '2008-10-19T00:00:00',
 '2009-09-13T00:00:00',
 '2009-09-14T00:00:00',
 '2009-09-15T00:00:00',
 '2009-09-16T00:00:00',
 '2009-11-07T00:00:00',
 '2009-11-08T00:00:00',
 '2009-11-09T00:00:00',
 '2009-11-10T00:00:00',
 '2009-11-11T00:00:00',
 '2009-11-12T00:00:00',
 '2009-11-13T00:00:00',
 '2009-11-14T00:00:00',
 '2009-11-15T00:00:00',
 '2010-07-22T00:00:00',
 '2010-07-23T00:00:00',
 '2010-07-24T00:00:00',
 '2010-07-25T00:00:00',
 '2010-07-28T00:00:00',
 '2010-07-29T00:00:00',
 '2010-09-07T00:00:00',
 '2010-10-03T00:00:00',
 '2010-10-04T00:00:00',
 '2011-07-17T00:00:00',
 '2011-07-19T00:00:00',
 '2011-08-27T00:00:00',
 '2011-08-28T00:00:00',
 '2011-08-29T00:00:00',
 '2011-08-30T00:00:00',
 '2011-09-10T00:00:00',
 '2011-09-12T00:00:00',
 '2011-09-19T00:00:00',
 '2011-11-01T00:00:00',
 '2011-11-26T00:00:00',
 '2011-11-27T00:00:00',
 '2011-11-28T00:00:00',
 '2011-11-29T00:00:00',
 '2011-11-30T00:00:00',
 '2012-07-06T00:00:00',
 '2012-07-07T00:00:00',
 '2012-08-22T00:00:00',
 '2012-08-23T00:00:00',
 '2012-08-24T00:00:00',
 '2012-08-25T00:00:00',
 '2012-08-26T00:00:00',
 '2012-08-27T00:00:00',
 '2012-08-28T00:00:00',
 '2012-08-29T00:00:00',
 '2012-08-30T00:00:00',
 '2012-09-04T00:00:00',
 '2012-11-02T00:00:00',
 '2012-11-03T00:00:00',
 '2012-11-04T00:00:00',
 '2012-11-05T00:00:00',
 '2013-11-08T00:00:00',
 '2014-08-28T00:00:00',
 '2014-11-21T00:00:00',
 '2014-11-22T00:00:00',
 '2014-11-23T00:00:00',
 '2014-11-24T00:00:00',
 '2014-11-25T00:00:00',
 '2014-11-26T00:00:00',
 '2014-11-27T00:00:00',
 '2015-06-11T00:00:00',
 '2015-06-12T00:00:00',
 '2015-06-13T00:00:00',
 '2015-06-14T00:00:00',
 '2015-06-15T00:00:00',
 '2015-11-03T00:00:00',
 '2015-11-04T00:00:00',
 '2015-11-05T00:00:00',
 '2015-11-06T00:00:00',
 '2015-11-07T00:00:00',
 '2015-11-08T00:00:00',
 '2015-11-09T00:00:00',
 '2015-11-10T00:00:00',
 '2015-11-11T00:00:00',
 '2015-11-12T00:00:00',
 '2015-11-13T00:00:00',
 '2015-11-14T00:00:00',
 '2015-11-15T00:00:00',
 '2015-11-25T00:00:00',
 '2015-11-26T00:00:00',
 '2015-11-27T00:00:00',
 '2015-11-28T00:00:00',
 '2015-11-29T00:00:00',
 '2016-06-15T00:00:00',
 '2016-06-16T00:00:00',
 '2016-06-17T00:00:00',
 '2016-06-18T00:00:00',
 '2016-06-19T00:00:00',
 '2016-07-23T00:00:00',
 '2016-07-24T00:00:00',
 '2016-07-25T00:00:00',
 '2016-11-22T00:00:00',
 '2016-11-23T00:00:00',
 '2016-11-25T00:00:00',
 '2016-11-26T00:00:00',
 '2016-11-27T00:00:00',
 '2017-07-05T00:00:00',
 '2017-07-06T00:00:00',
 '2017-07-07T00:00:00',
 '2017-07-08T00:00:00',
 '2017-07-14T00:00:00',
 '2017-07-15T00:00:00',
 '2017-07-16T00:00:00',
 '2017-07-17T00:00:00',
 '2017-07-18T00:00:00',
 '2017-07-19T00:00:00',
 '2017-08-28T00:00:00',
 '2017-09-15T00:00:00',
 '2017-09-16T00:00:00',
 '2017-11-28T00:00:00',
 '2018-11-07T00:00:00',
 '2018-11-08T00:00:00',
 '2018-11-09T00:00:00',
 '2019-06-03T00:00:00',
 '2019-06-04T00:00:00',
 '2019-06-05T00:00:00',
 '2019-06-06T00:00:00',
 '2019-06-07T00:00:00',
 '2019-06-08T00:00:00',
 '2019-08-24T00:00:00',
 '2019-10-25T00:00:00',
 '2019-10-27T00:00:00',
 '2020-06-13T00:00:00',
 '2020-06-14T00:00:00',
 '2020-07-27T00:00:00',
 '2020-07-28T00:00:00',
 '2020-07-29T00:00:00',
 '2020-08-31T00:00:00',
 '2020-09-01T00:00:00',
 '2020-09-04T00:00:00',
 '2020-09-14T00:00:00',
 '2020-11-21T00:00:00',
 '2020-11-22T00: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()