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