load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_code.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/wrf/WRFUserARW.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/wrf/WRF_contributed.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_csm.ncl" 
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/contributed.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/shea_util.ncl"
load "./time_avg.ncl"
external SOR "./metlib.so"

begin
    months=(/"JAN","FEB","MAR","APR","MAY","JUN","JUL","AUG","SEP","OCT","NOV","DEC"/)
    ihour       = "00"  ;init hour
    imonth  = "01"  ;init month
    iday        = "08"   ;init day
    iyr     = "2009";init year
    mon_str = months(stringtoint(imonth)-1)
    strat="strat"
    rf="fnl"

    ;inFile1="output/20090108_strat_fnl.nc"
    inFile1="output/20090114_strat_rt.nc"
    ;inFile2="output/20090108_warm_coolarea.nc"

; The WRF ARW input file.  
; This needs to have a ".nc" appended, so just do it.
    a = addfile(inFile1,"r")
    b = addfile(inFile2,"r")
    ; We generate plots, but what kind do we prefer?
    ; type = "x11"
    ; type = "pdf"
    type = "ps"
    ; type = "ncgm"
    wks = gsn_open_wks(type,"plt_xy_int_avg_"+name);ihour+"Z"+iday2+mon_str+iyr)

    pressure_levels = (/900,850,700,600,550,500,450,400,350,300,250,200,150,100,80,60,50,40,20,15/)
    pressure_levels_R = (/900.0,850.0,700.0,600.0,550.0,500.0,450.0,400.0,350.0,300.0,\
                                250.0,200.0,150.0,100.0,80.0,60.0,50.0,40.0,20.0,15.0/)

    nlevels=dimsizes(pressure_levels)  
    if(pres_s.eq.20) then
        pStrat = 12   ;;Index of where stratosphere starts (in fortran, idx=1...nlevs)
        pTop = nlevels-1
    end if
    if(pres_s.eq.500)then
        pStrat=6
        pTop=12
    end if

    ;What times and how many time steps are in the data set?
    times  = wrf_user_list_times(a)  ; get times in the file
    times2 = wrf_user_list_times(b)
    ntimes = dimsizes(times)         ; number of times in the file
    ntimes2 = dimsizes(times2)         ; number of times in the file
    TimeConst=2
    T1 = a->Times
    T2 = b->Times
    T1A = wrf_times_c( T1,1 )
    T2A = wrf_times_c( T2,1 )
    T1_S = chartostring( T1 )
    T2_S = chartostring( T2 )
    start_index = 0
    do k = 0, dimsizes(T1(:,0))-1
        if(T1A(k).eq.T2A(0)) then
            start_index = k
        end if
    end do
    ktimes=(ntimes2/TimeConst)+1
    endTime=ntimes2+start_index-1
    print("ENDTIME:"+endTime)  
    if(endTime .gt. ntimes-1) then
        endTime=ntimes-1
        ktimes=((endTime-start_index)/TimeConst)+1
        print("#### KTIMES RESET TO :"+ktimes+" ####")
        print("#### ENDTIME RESET TO :"+endTime+" ####")
    end if
    ktimesB=ktimes
    
    lat2d = wrf_user_getvar(a, "lat", 0)   ; Lattitude
    lon2d = wrf_user_getvar(a, "lon", 0)   ; Longitude
    lat1d=lat2d(:,0)           ; Lattitude 1-D
    lon1d=lon2d(0,:)           ; Longitude 1-D
    x_size=dimsizes(lon1d)
    y_size=dimsizes(lat1d)
    ;print("latmin="+lat1d(0)+" latmax="+lat1d(y_size-1)+" lonmin="+lon1d(0)+" lonmax="+lon1d(x_size-1))
    s       =new((/y_size,x_size,nlevels,ktimes/),float)
    Z_Array =new((/y_size,x_size,nlevels,ktimes/),float)
    T       =new((/y_size,x_size,nlevels,ktimes/),float)
    TH      =new((/y_size,x_size,nlevels,ktimes/),float)
    W       =new((/y_size,x_size,nlevels,ktimes/),float)
    omega   =new((/y_size,x_size,nlevels,ktimes/),float)
    dT      =new((/y_size,x_size,nlevels,ktimes/),float)
    t_adv   =new((/y_size,x_size,nlevels,ktimes/),float)
    

    int_tot =new((/y_size,x_size,ktimes/),float)
    int_adia=new((/y_size,x_size,ktimes/),float)
    int_adv =new((/y_size,x_size,ktimes/),float)


    
    lat2dB  = wrf_user_getvar(a, "lat", 0) ; Lattitude
    lon2dB  = wrf_user_getvar(a, "lon", 0) ; Longitude
    lat1dB  =lat2dB(:,0)            ; Lattitude 1-D
    lon1dB  =lon2dB(0,:)            ; Longitude 1-D
    x_sizeB =dimsizes(lon1dB)
    y_sizeB =dimsizes(lat1dB)
    sB      =new((/y_sizeB,x_sizeB,nlevels,ktimesB/),float)
    Z_ArrayB=new((/y_sizeB,x_sizeB,nlevels,ktimesB/),float)
    TB      =new((/y_sizeB,x_sizeB,nlevels,ktimesB/),float)
    THB     =new((/y_sizeB,x_sizeB,nlevels,ktimesB/),float)
    WB      =new((/y_sizeB,x_sizeB,nlevels,ktimesB/),float)
    omegaB  =new((/y_sizeB,x_sizeB,nlevels,ktimesB/),float)
    dTB     =new((/y_sizeB,x_sizeB,nlevels,ktimesB/),float)
    t_advB  =new((/y_sizeB,x_sizeB,nlevels,ktimesB/),float)

    int_totB    =new((/y_sizeB,x_sizeB,ktimesB/),float)
    int_adiaB   =new((/y_sizeB,x_sizeB,ktimesB/),float)
    int_advB    =new((/y_sizeB,x_sizeB,ktimesB/),float)

    lat_rad =lat1d*(3.14159/180.0)
    lon_rad =lon1d*(3.14159/180.0)
    lat_radB=lat1dB*(3.14159/180.0)
    lon_radB=lon1dB*(3.14159/180.0)    

    x=new(ktimes,string)
    dT=0.0
    dTB=0.0
    do it=start_index,endTime,TimeConst;
        it2=it-start_index
        idx=(it2/TimeConst)
        x(idx)=times2(it2)  
        print("it1="+it+" idx="+idx)
        print("Working on time: " + times2(it2))
        p  = wrf_user_getvar(a, "pressure",it)    ; pressure as vertical coordinate
        z  = wrf_user_getvar(a, "z",it)           ; grid point height
        t  = wrf_user_getvar(a, "tk", it)     ; temperature in Kelvins
        th = wrf_user_getvar(a, "th", it)     ; potential temperature
        w  = wrf_user_getvar(a, "W",it)           ; grid point height

        pB  = wrf_user_getvar(b, "pressure",it2); pressure as vertical coordinate
        zB  = wrf_user_getvar(b, "z",it2)     ; grid point height
        tB  = wrf_user_getvar(b, "tk", it2)       ; temperature in Kelvins
        thB = wrf_user_getvar(b, "th", it2)       ; potential temperature
        wB  = wrf_user_getvar(b, "W",it2)     ; grid point height

        z_size=dimsizes(w(:,0,0))
        z_sizeB=dimsizes(wB(:,0,0))
        w_new=w(1:z_size-1,:,:)
        w_newB=wB(1:z_sizeB-1,:,:)
    
        dT(:,:,:,0)   = 0.0
        dTB(:,:,:,0)= 0.0
        do level = 0,nlevels-1    ;;;BEGIN LEVEL LOOP
            ;;;;;;;;Then interpolate variables on \pres\ surface;;;;;;;;;;;
            pres=pressure_levels(level)
            Z_Array(:,:,level,idx)  = wrf_user_intrp3d(z,p,"h",pres,0.,False)
            T(:,:,level,idx)        = wrf_user_intrp3d(t,p,"h",pres,0.,False)
            TH(:,:,level,idx)       = wrf_user_intrp3d(th,p,"h",pres,0.,False)
            W(:,:,level,idx)        = wrf_user_intrp3d(w_new,p,"h",pres,0.,False)
            omega(:,:,level,idx)    = -pres*9.81*W(:,:,level,idx)/(287*T(:,:,level,idx))


            Z_ArrayB(:,:,level,idx) = wrf_user_intrp3d(zB,pB,"h",pres,0.,False)
            TB(:,:,level,idx)       = wrf_user_intrp3d(tB,pB,"h",pres,0.,False)
            THB(:,:,level,idx)      = wrf_user_intrp3d(thB,pB,"h",pres,0.,False)
            WB(:,:,level,idx)       = wrf_user_intrp3d(w_newB,pB,"h",pres,0.,False)
            omegaB(:,:,level,idx)   = -pres*9.81*WB(:,:,level,idx)/(287*TB(:,:,level,idx))
            wrf_smooth_2d(Z_Array(:,:,level,idx),3)
            wrf_smooth_2d(Z_ArrayB(:,:,level,idx),3)
            wrf_smooth_2d(T(:,:,level,idx),8)
            wrf_smooth_2d(TB(:,:,level,idx),8)

            if (idx.gt.1)
                dT(:,:,:,idx)   = T(:,:,:,idx)-T(:,:,:,idx-2)
                dTB(:,:,:,idx)  = TB(:,:,:,idx)-TB(:,:,:,idx-2)
            end if
        end do
    end do
    int_adia=0.0
    int_adiaB=0.0
    int_tot=0.0
    int_totB=0.0
    t_adv=0.0
    t_advB=0.0
    int_adv=0.0
    int_advB=0.0
    s=0.0
    sB=0.0

    print("Starting stsb loop")
    SOR::stsb(TH,T,pressure_levels_R,x_size,y_size,nlevels,ktimes,s)
    SOR::stsb(THB,TB,pressure_levels_R,x_sizeB,y_sizeB,nlevels,ktimesB,sB)
    print("End of stsb loop")
    print("Starting int. loop.  ktimes="+ktimes+" x_size="+x_size+" y_size="+y_size)
    ;;TOTAL;;
    SOR::intstrat(dT,pressure_levels_R,x_size,y_size,nlevels,ktimes,pStrat,pTop,int_tot)
    SOR::intstrat(dTB,pressure_levels_R,x_sizeB,y_sizeB,nlevels,ktimesB,pStrat,pTop,int_totB)
    ;;ADIABATIC;;
    SOR::intadia(omega,s,pressure_levels_R,x_size,y_size,nlevels,ktimes,pStrat,pTop,int_adia)
    SOR::intadia(omegaB,sB,pressure_levels_R,x_sizeB,y_sizeB,nlevels,ktimes,pStrat,pTop,int_adiaB)
    ;;ADVECTIVE;;
    SOR::adv(Z_Array,T,pressure_levels_R,lat_rad,lon_rad,x_size,y_size,nlevels,ktimes,t_adv)
    SOR::adv(Z_ArrayB,TB,pressure_levels_R,lat_radB,lon_radB,x_sizeB,y_sizeB,nlevels,ktimes,t_advB)
    SOR::intstrat(t_adv,pressure_levels_R,x_size,y_size,nlevels,ktimes,pStrat,pTop,int_adv)
    SOR::intstrat(t_advB,pressure_levels_R,x_sizeB,y_sizeB,nlevels,ktimes,pStrat,pTop,int_advB)
    print("End of int. loop")
    ;;Put int_adv in units of m/12h
    int_adv=int_adv*4.32e4
    int_advB=int_advB*4.32e4
    wrf_smooth_2d(int_tot,9)
    wrf_smooth_2d(int_totB,9)
    wrf_smooth_2d(int_adia,9)
    wrf_smooth_2d(int_adiaB,9)
    wrf_smooth_2d(int_adv,9)
    wrf_smooth_2d(int_advB,9)

;;;Start the time loop;;;
  ;do it = TimeConst,ntimes-TimeConst,TimeConst;start_index+8,4;endTime,4
    xp=new(ktimes,float)
    yp=new(ktimes,float)
    plot_total = new((/2,ktimes/),float)
    plot_adia = new((/2,ktimes/),float)
    plot_adv = new((/2,ktimes/),float)
    plot_total=0.0
    plot_adia=0.0
    plot_adv=0.0
    stringData=name+","
    stringDataA="fnl_"+iday2+","
    xBarOnset=0
    xBarDecay=0
    do idx = 2,ktimes-TimeConst,1
        it=TimeConst*idx
        print("Working on time: " + times2(it) + " idx="+ idx )
        if(times2(it).eq."2009-01-16_00:00:00")
            xBarOnset=idx
        end if
        if(times2(it).eq."2009-01-20_00:00:00")
            xBarDecay=idx
        end if
        xloc=0
        yloc=0
        z_diff=Z_Array(:,:,5,idx)-Z_Array(:,:,5,idx-2)    
        if (it .gt. 40) then
            mxmn = 0 ; 0 for falls, 1 for rises
        else
            mxmn = 1
        end if

        ;SOR::mxloc(int_tot(:,:,idx),x_size,y_size,56,97,33,96,mxmn,xloc,yloc)
        SOR::mxloc(z_diff,x_size,y_size,56,97,33,96,mxmn,xloc,yloc)
        xp(idx)=lon1d(xloc)
        yp(idx)=lat1d(yloc)

        plot_total(:,idx)=time_avg(int_tot,int_totB,xloc,yloc,idx)
        if (plot_total(1,idx) .gt. 1000) then
            print("#########################")
            print("plot_total exceedes 1000")
            print("plot_total(1,"+idx+")="+plot_total(1,idx))
            print("#########################")
        end if    
        plot_adia(:,idx)=time_avg(int_adia,int_adiaB,xloc,yloc,idx)
        plot_adv(:,idx)=time_avg(int_adv,int_advB,xloc,yloc,idx)
    ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    end do        ; END OF TIME LOOP
    
    ;;;;;;;;;;;;;;;
    ;; PLOTTING  ;;
    ;;;;;;;;;;;;;;;

    vals     = ispan(0,ktimes-1,1)
    lag_vals=(vals-8)*6
    titleCommon=iday2+" "+mon_str+" "+iyr+" from "+pressure_levels(pStrat-1)+\
                " to "+pressure_levels(pTop-1)+" hPa"
    
    pltres                          = True
    pltres@tmXBMode                 = "Explicit"  ; explicit labels
    pltres@tmXBValues               = vals          ; location of labels
    pltres@tmXBLabels               = lag_vals              ; labels themselves
    pltres@tmXBLabelFontHeightF     = 0.015
    pltres@tmLabelAutoStride        = True         ; nice stride on labels
    if (color .eq. 1) then
        pltres@xyLineColors           = (/"blue","red"/)
        pltres@pmLegendDisplayMode  = "Always"             ; turn on legend
    else
        pltres@xyLineColors           = (/"black","black"/)
        pltres@xyMonoLineThickness  = False
        pltres@xyLineThicknesses  = (/"3.0","2.0"/)
    end if

    pltres@pmLegendSide             = "Top"                 ; Change location of legend
    pltres@pmLegendParallelPosF     = 1.1                     ; move units right
    pltres@pmLegendOrthogonalPosF   = -0.3                   ; move units down
    pltres@pmLegendWidthF           = 0.15                 ; Change width and
    pltres@pmLegendHeightF          = 0.18                    ; height of legend.
    pltres@lgPerimOn                = False                ; turn off box around
    pltres@lgLabelFontHeightF       = .02                   ; label font height
    ;pltres@xyExplicitLegendLabels = (/"FNL",legend/)      ; create explicit labels
    pltres@xyExplicitLegendLabels  = (/"RT",legend/)        ; create explicit labels
    pltres@gsnDraw                     = False
    pltres@gsnFrame                    = False
    pltres@gsnMaximize                 = True
    pltres@tiYAxisString            = "Height change m12h~S~-1~N~"
    pltres@tiXAxisString            = "Lag time (hours) from block onset"
    plxOnset=(/xBarOnset,xBarOnset/)
    plxDecay=(/xBarDecay,xBarDecay/)

    ;; Total Plot ;;
    pltres@tiMainString         = "Integrated Total " +titleCommon
    plot  = gsn_csm_xy(wks,vals,plot_total,pltres)
    ply=(/min(plot_total),max(plot_total)/)
    hlineOnset=gsn_add_polyline(wks,plot,plxOnset,ply,True)
    hlineDecay=gsn_add_polyline(wks,plot,plxDecay,ply,True)
    draw(plot)
    frame(wks)

    ;; Adiabatic Plot ;;
    pltres@tiMainString         = "Integrated Adiabatic " +titleCommon
    plot  = gsn_csm_xy(wks,vals,plot_adia,pltres)
    ply=(/min(plot_adia),max(plot_adia)/)
    hlineOnset=gsn_add_polyline(wks,plot,plxOnset,ply,True)
    hlineDecay=gsn_add_polyline(wks,plot,plxDecay,ply,True)
    draw(plot)
    frame(wks)


    ;; Advective Plot ;;
    pltres@tiMainString         = "Integrated Advective " +titleCommon
    plot  = gsn_csm_xy(wks,vals,plot_adv,pltres)
    ply=(/min(plot_adv),max(plot_adv)/)
    hlineOnset=gsn_add_polyline(wks,plot,plxOnset,ply,True)
    hlineDecay=gsn_add_polyline(wks,plot,plxDecay,ply,True)
    draw(plot)
    frame(wks)

    ;;;;;;;;;;;;;;;;;;;;;
    ;; Now plot points ;;
    ;;;;;;;;;;;;;;;;;;;;;
    res                         = True
    res@gsnDraw                = False     ; don't draw
    res@gsnFrame               = False     ; don't advance frame
    res@gsnMaximize            = True
    res@mpFillOn                = False
    res@mpPerimOn               = True
    res@pmTickMarkDisplayMode   = "Always"
    ;res@mpLimitMode           = "LatLon"  ; select subregion
    res@mpMinLatF               = 0         ;lat1d(0)
    res@mpMaxLatF               = 75        ;lat1d(y_size-1)                
    res@mpMinLonF               = -180      ;lon1d(0)-360.0
    res@mpMaxLonF               = -60
    res@tmYROn                  = False     ; turn off right and top tickmarks
    res@tmXTOn                  = False
    res@tiMainString            = "Fastest height change track"  ; title
    res@tiMainFontHeightF       = 0.02
    
    map = gsn_csm_map_ce(wks,res)          ; create map

    ;Add trajectory lines.
    plres                  = True                  ; polyline resources
    plres@gsLineThicknessF = 3.0                   ; line thickness
    plres@gsLineColor      = "black"
    ;line = gsn_add_polyline(wks,map,xp,yp,plres)  ; draw the traj

    ;draw(map)                                        
    ;frame(wks)
    ;asciiwrite("fnl_adia"+iday2+".txt",stringDataA)
    ;asciiwrite("tmp.txt",stringData)

end