;*****************************************************************************************
; This script plots vertical mass flux in Moist Isentropic Coordinates. 
; This script also contours the mass-moist-isentropic streamfunction.
;
; Created by: Josh Alland
; Date: 07 October 2014
;*****************************************************************************************

load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_code.ncl"   
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_csm.ncl"   
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/contributed.ncl" 
load "/nfs/kc11/jalland/Downdrafts_Project/load.ncl"  

begin
omega = 7.292115*10^(-5)   ;---rad/sec
a = 6371000.    ;---m

ncid = addfile("/erai/1979/u.1979.nc","r")
lat = ncid->lat(:)

pi = 4.*atan(1)
lat_rad = 2.*pi/360.*lat(:)

lon = ncid->lon(:)
lev = ncid->lev(:)
lev_to_plot = (/200,300,500,600,700,800,850,900,925,1000/)

d1 = dimsizes(lev_to_plot) 
d2 = dimsizes(lat)
d3 = dimsizes(lon)

shear_av = new((/d2,d3/),"float")

do year=1979,2013,1   ;---Remember leap years
  count = 0
    do j=0,d2-1,1
      do k=0,d3-1,1
        shear_av(j,k) = 0.   
      end do
    end do
    if(year.eq.1980 .OR. year.eq.1984 .OR. year.eq.1988 .OR. year.eq.1992 .OR. year.eq.1996 .OR. year.eq.2000 .OR. year.eq.2004 .OR. year.eq.2008 .OR. year.eq.2012)
        start_time = 607
        end_time   = 975
    else
        start_time = 603
        end_time   = 971
    end if
    ncid1 = addfile("/erai/"+year+"/u."+year+".nc","r")
    ncid2 = addfile("/erai/"+year+"/v."+year+".nc","r")
    do time = start_time,end_time,1    ;---Do loop to average in a single year
         print("Year = "+year+", Time = "+time) 
         u_200 = ncid1->u(time,{200},:,:) ;200,300,500,600,700,800,850,900,925,1000
         u_850 = ncid1->u(time,{850},:,:) ;200,300,500,600,700,800,850,900,925,1000
         v_200 = ncid2->v(time,{200},:,:) ;200,300,500,600,700,800,850,900,925,1000
         v_850 = ncid2->v(time,{850},:,:) ;200,300,500,600,700,800,850,900,925,1000
         
         shear = ((u_200-u_850)^2 + (v_200-v_850)^2)^(1./2.)
         shear_av = shear_av + shear
         count = count+1
    end do   ;---End of do loop to average in a single year
       delete(ncid1)

shear_av = shear_av/count

  shear_av!0 = "lat"
  shear_av&lat = lat
  shear_av!1 = "lon"
  shear_av&lon = lon
  shear_av@units = "m/s"
  shear_av@long_name = "time-averaged deep layer (200-850 hPa) vertical wind shear"
  

;nlev  = dimsizes(lev_to_plot)                                               
nlat  = dimsizes(lat)  
nlon  = dimsizes(lon)   
diro = "shear_av_SUMMER/"                     ; Output directory
filo = "shear_av_SUMMER_"+year+".nc"             ; Output file
system("/bin/rm -f " + diro + filo)    ; remove if exists
fout  = addfile (diro + filo, "c")  ; open output file   

setfileoption(fout,"DefineMode",True)

fAtt               = True            ; assign file attributes
fAtt@title         = "NCL Efficient Approach to netCDF Creation"  
fAtt@source_file   =  "original-file.nc"
fAtt@Conventions   = "None"   
fAtt@creation_date = systemfunc ("date")        
fileattdef( fout, fAtt )            ; copy file attributes  

dimNames = (/"lat", "lon"/)  
dimSizes = (/nlat,nlon/) 
dimUnlim = (/False, False/)   
filedimdef(fout,dimNames,dimSizes,dimUnlim)

;filevardef(fout, "lev"  ,typeof(lev),getvardims(lev) )                           
filevardef(fout, "lat"  ,typeof(lat),getvardims(lat))                          
filevardef(fout, "lon"  ,typeof(lon),getvardims(lon))                          
filevardef(fout, "shear_av"    ,typeof(shear_av),(/"lat","lon"/))    

;filevarattdef(fout,"lev",lev)                           ; copy T attributes
filevarattdef(fout,"lat",lat)                    ; copy time attributes
filevarattdef(fout,"lon",lon)                     ; copy lev attributes
filevarattdef(fout,"shear_av",shear_av)                     ; copy lat attributes

setfileoption(fout,"DefineMode",False)

;fout->lev    = (/lev_to_plot/)
fout->lat    = (/lat/)
fout->lon    = (/lon/) 
fout->shear_av      = (/shear_av/)

end do
end