#!/usr/bin/env python

### A. Spiga LMD 29/05/2011

def usage():
    print 'USAGE:  winds.py nvert file (target)'
    print 'file  : name of input netcdf file.'
    print 'nvert : vertical level.'
    print '(target) : a directory with write rights.'
    print 'Example: winds.py 0 /d5/aslmd/LMD_MM_MARS_SIMUS/OM/OM6_TI85/wrfout_d01_2024-06-43_06:00:00_zabg ~/'

def winds (namefile,nvert,proj="cyl",back=None,target=None): 
    from netCDF4 import Dataset
    from myplot import getcoord2d,define_proj,makeplotpng,simplinterv,vectorfield,ptitle
    from matplotlib.pyplot import contourf, subplot, figure
    nc  = Dataset(namefile)
    [lon2d,lat2d] = getcoord2d(nc)
    [wlon,wlat] = simplinterv(lon2d,lat2d)
    [u,v] = getwinds(nc)
    nt = len(u[:,0,0,0])
    fig = figure()
    fig.subplots_adjust(wspace = 0.0, hspace = 0.3)
    sub = 221
    for i in range(0,nt-1,int(nt/4.)):
       subplot(sub)
       ptitle("Winds time"+str(i)+" level"+str(nvert))
       m = define_proj(proj,wlon,wlat,back=back)
       x, y = m(lon2d, lat2d)
       contourf(x, y, nc.variables['HGT'][0,:,:] / 1000., 30)
       vectorfield(u[i,nvert,:,:], v[i,nvert,:,:],\
                      x, y, stride=5, csmooth=3,\
                      scale=20., factor=300., color='k')
       sub += 1
    if not target:   zeplot = namefile+".winds"+str(nvert)
    else:            zeplot = target+"winds"+str(nvert)
    makeplotpng(zeplot,pad_inches_value=0.35)   

def getwinds (nc,charu='Um',charv='Vm'):
    u = nc.variables[charu]
    v = nc.variables[charv]
    return u,v

if __name__ == "__main__":
    import sys
    if (len(sys.argv)) == 3:     winds( str(sys.argv[2]) , int(sys.argv[1]) )
    elif (len(sys.argv)) == 4:   winds( str(sys.argv[2]) , int(sys.argv[1]) , target=str(sys.argv[3]))