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/gsn_code.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/wrf/WRFUserARW.ncl"

begin

; for use of high resolution map database, execute
; 'setenv NCARG_RANGS /network/rit/lab/atm419lab/zhou/rangs' (on snow)
; 'setenv NCARG_RANGS /atm419lab/zhou/rangs' (on ash, reed)
; before running the script

c="./cmap_uv200.txt"
cmap=asciiread(c,(/83,3/),"float")
cmap=cmap/255

center=asciiread("../../DATA_FAILURE_START/latlon_failure_start.txt",(/82,2/),"float")
;print(center(:,1))
tinfo=asciiread("../../DATA_FAILURE_START/time_failure_start.txt",(/82,2/),"integer")
yyyymmddhh=tinfo(:,1)
;print(yyyymmddhh)
name=asciiread("../../DATA_FAILURE_START/name_failure.txt",(/82/),"string")
;print(name)

dcase=dimsizes(yyyymmddhh)
U_comp=new((/dcase,161,161/),"float")
V_comp=new((/dcase,161,161/),"float")
Z_comp=new((/dcase,161,161/),"float")

do case=0,dcase-1,1

print("working on "+yyyymmddhh(case)+" -- "+name(case))

data=addfile("../../DATA_FAILURE_START/"+yyyymmddhh(case)+"_UV.grib","r")
;names=getfilevarnames(data)
;print(names)

U:=data->U_GDS0_ISBL
U!0="lev"
U!1="lat"
U!2="lon"

V:=data->V_GDS0_ISBL
V!0="lev"
V!1="lat"
V!2="lon"

Z:=data->Z_GDS0_ISBL
;printVarSummary(Z)
Z!0="lev"
Z!1="lat"
Z!2="lon"

;print(U&lev) ; levs are 200, 500, 850
;print(U&lat)

lat=U&lat
lon=U&lon
centlat_ind=ind_nearest_coord(center(case,0),lat,0)
centlon_ind=ind_nearest_coord(center(case,1),lon,0)
centlat=lat(centlat_ind)
centlon=lon(centlon_ind)

U200=U(0,:,:)
V200=V(0,:,:)
MAG=sqrt(U200^2+V200^2)
Z200=Z(0,:,:)/9.8/10.0

copy_VarMeta(U(0,:,:),U200)
copy_VarMeta(U(0,:,:),V200)
copy_VarMeta(U(0,:,:),MAG)
copy_VarMeta(Z(0,:,:),Z200)

pi=4.0*atan(1.0)
mlist=(/2,3,4,2,5,6,7,2,8,9,2/)
dm=dimsizes(mlist)

U200_1deg=U200(::4,::4)
V200_1deg=V200(::4,::4)
Z200_1deg=Z200(::4,::4)
;printVarSummary(U200_1deg)
;print(U200_1deg&lon)
lat_1deg=U200_1deg&lat
lon_1deg=U200_1deg&lon
dlat_1deg=dimsizes(lat_1deg)
dlon_1deg=dimsizes(lon_1deg)

U200_1deg_smthx=U200_1deg
V200_1deg_smthx=V200_1deg
Z200_1deg_smthx=Z200_1deg

do m=0,dm-1,1
  ;print("max u ="+max(U200_1deg_smthx)+"min u ="+min(U200_1deg_smthx))
  K=0.5*(1.0-cos(2.0*pi/mlist(m)))^(-1.0)

  U200_1deg_smthx_inter=U200_1deg_smthx
  V200_1deg_smthx_inter=V200_1deg_smthx
  Z200_1deg_smthx_inter=Z200_1deg_smthx

    do i=1,dlon_1deg-2,1

      U200_1deg_smthx(:,i)=U200_1deg_smthx_inter(:,i)+K*(U200_1deg_smthx_inter(:,i-1)+U200_1deg_smthx_inter(:,i+1)-2.0*U200_1deg_smthx_inter(:,i))
      V200_1deg_smthx(:,i)=V200_1deg_smthx_inter(:,i)+K*(V200_1deg_smthx_inter(:,i-1)+V200_1deg_smthx_inter(:,i+1)-2.0*V200_1deg_smthx_inter(:,i))
      Z200_1deg_smthx(:,i)=Z200_1deg_smthx_inter(:,i)+K*(Z200_1deg_smthx_inter(:,i-1)+Z200_1deg_smthx_inter(:,i+1)-2.0*Z200_1deg_smthx_inter(:,i))

    end do

end do

U200_1deg_smthy=U200_1deg_smthx
V200_1deg_smthy=V200_1deg_smthx
Z200_1deg_smthy=Z200_1deg_smthx

do m=0,dm-1,1
  ;print("max u ="+max(U200_1deg_smthy)+"min u ="+min(U200_1deg_smthy))
  K=0.5*(1.0-cos(2.0*pi/mlist(m)))^(-1)

  U200_1deg_smthy_inter=U200_1deg_smthy
  V200_1deg_smthy_inter=V200_1deg_smthy
  Z200_1deg_smthy_inter=Z200_1deg_smthy

    do j=1,dlat_1deg-2,1

      U200_1deg_smthy(j,:)=U200_1deg_smthy_inter(j,:)+K*(U200_1deg_smthy_inter(j-1,:)+U200_1deg_smthy_inter(j+1,:)-2.0*U200_1deg_smthy_inter(j,:))
      V200_1deg_smthy(j,:)=V200_1deg_smthy_inter(j,:)+K*(V200_1deg_smthy_inter(j-1,:)+V200_1deg_smthy_inter(j+1,:)-2.0*V200_1deg_smthy_inter(j,:))
      Z200_1deg_smthy(j,:)=Z200_1deg_smthy_inter(j,:)+K*(Z200_1deg_smthy_inter(j-1,:)+Z200_1deg_smthy_inter(j+1,:)-2.0*Z200_1deg_smthy_inter(j,:))

    end do

end do

U200_1deg_smth=U200_1deg_smthy
V200_1deg_smth=V200_1deg_smthy
Z200_1deg_smth=Z200_1deg_smthy

;print("max u ="+max(U200_1deg_smth)+"min u ="+min(U200_1deg_smth))
;print("max v ="+max(V200_1deg_smth)+"min v ="+min(V200_1deg_smth))
;print("max z ="+max(Z200_1deg_smth)+"min z ="+min(Z200_1deg_smth))

U200_smth=linint2(lon_1deg,lat_1deg(::-1),U200_1deg_smth,True,lon,lat(::-1),0) ;lat should be increasing
V200_smth=linint2(lon_1deg,lat_1deg(::-1),V200_1deg_smth,True,lon,lat(::-1),0) ;as long as two sets of lat are aligned,
Z200_smth=linint2(lon_1deg,lat_1deg(::-1),Z200_1deg_smth,True,lon,lat(::-1),0) ;original data sequence doesn't matter

;print(num(ismissing(U200_smth)))
;print(num(ismissing(V200_smth)))
;print(num(ismissing(Z200_smth)))
;printVarSummary(U200_smth)

copy_VarMeta(U200,U200_smth) ;otherwise, lat will be flipped
copy_VarMeta(V200,V200_smth)
copy_VarMeta(Z200,Z200_smth)

U200_dev=U200-U200_smth
V200_dev=V200-V200_smth
Z200_dev=Z200-Z200_smth

copy_VarMeta(U200,U200_dev) 
copy_VarMeta(V200,V200_dev)
copy_VarMeta(Z200,Z200_dev)

;print("max Udev = "+max(U200_dev)+" min Udev = "+min(U200_dev))

rn=31 ;radius of interpolation domain, set bigger than needed (7.5 deg instead of 5 deg)
tn=72 
dr=0.25 ; deg lat/lon
dtheta=5 ;deg
theta=ispan(0,360-1,dtheta)

lon_cyl2d=new((/rn,tn/),float)
lat_cyl2d=new((/rn,tn/),float)

do r=0,rn-1,1

  lon_cyl2d(r,:)=lon(centlon_ind)+r*dr*cos(theta*pi/180)
  lat_cyl2d(r,:)=lat(centlat_ind)+r*dr*sin(theta*pi/180)

end do

lon_cyl1d=ndtooned(lon_cyl2d)
lat_cyl1d=ndtooned(lat_cyl2d)

;print(lon_cyl1d)

U200_dev_cyl1d=linint2_points(lon,lat(::-1),U200_dev(::-1,:),True,lon_cyl1d,lat_cyl1d,0) ;need to reverse original data 
V200_dev_cyl1d=linint2_points(lon,lat(::-1),V200_dev(::-1,:),True,lon_cyl1d,lat_cyl1d,0) ;corresponding to lat array
Z200_dev_cyl1d=linint2_points(lon,lat(::-1),Z200_dev(::-1,:),True,lon_cyl1d,lat_cyl1d,0)

U200_dev_cyl2d=reshape(U200_dev_cyl1d,(/rn,tn/))
V200_dev_cyl2d=reshape(V200_dev_cyl1d,(/rn,tn/))
Z200_dev_cyl2d=reshape(Z200_dev_cyl1d,(/rn,tn/))

;print("max Udev = "+max(U200_dev)+" min Udev = "+min(U200_dev))
;print("max Udev.interp = "+max(U200_dev_cyl1d)+" min Udev.interp = "+min(U200_dev_cyl1d))

r0=21 ;radius of filtering domain (5 deg)
l=r0/5.0 ;parameter controling filtering shape

U200_vortex_cyl2d=U200_dev_cyl2d
V200_vortex_cyl2d=V200_dev_cyl2d
Z200_vortex_cyl2d=Z200_dev_cyl2d

;print("max Uvortex = "+max(U200_vortex_cyl2d)+" min Uvortex = "+min(U200_vortex_cyl2d))

do r=0,rn-1,1
  
  if (r.le.r0) then

    E=(exp(-((r0-r)^2/l^2))-exp(-((r0)^2/l^2)))/(1.0-exp(-((r0)^2/l^2)))
    ;print("E = "+E)
    U200_vortex_cyl2d(r,:)=U200_dev_cyl2d(r,:)-U200_dev_cyl2d(r0,:)*E-avg(U200_dev_cyl2d(r0,:))*(1-E)
    V200_vortex_cyl2d(r,:)=V200_dev_cyl2d(r,:)-V200_dev_cyl2d(r0,:)*E-avg(V200_dev_cyl2d(r0,:))*(1-E)
    Z200_vortex_cyl2d(r,:)=Z200_dev_cyl2d(r,:)-Z200_dev_cyl2d(r0,:)*E-avg(Z200_dev_cyl2d(r0,:))*(1-E)

  else

    U200_vortex_cyl2d(r,:)=0
    V200_vortex_cyl2d(r,:)=0
    Z200_vortex_cyl2d(r,:)=0

  end if

end do

;print("max Uvortex = "+max(U200_vortex_cyl2d)+" min Uvortex = "+min(U200_vortex_cyl2d))
;small discrepancy should be a result of TC location difference

lon_sub=lon(centlon_ind-20:centlon_ind+20)
lat_sub=lat(centlat_ind+20:centlat_ind-20)

U200_vortex_sub=transpose(dsgrid2(lon_cyl1d,lat_cyl1d,ndtooned(U200_vortex_cyl2d),lon_sub,lat_sub)) ;dsgrid returns dx*dy array
V200_vortex_sub=transpose(dsgrid2(lon_cyl1d,lat_cyl1d,ndtooned(V200_vortex_cyl2d),lon_sub,lat_sub)) ;instead of dy*dx !!!!!!!!!
Z200_vortex_sub=transpose(dsgrid2(lon_cyl1d,lat_cyl1d,ndtooned(Z200_vortex_cyl2d),lon_sub,lat_sub))

;print("max Uvortex_sub = "+max(U200_vortex_sub)+" min Uvortex_sub = "+min(U200_vortex_sub))
;printVarSummary(U200_vortex_sub)

copy_VarMeta(U200(centlat_ind+20:centlat_ind-20,centlon_ind-20:centlon_ind+20),U200_vortex_sub)
copy_VarMeta(V200(centlat_ind+20:centlat_ind-20,centlon_ind-20:centlon_ind+20),V200_vortex_sub)
copy_VarMeta(Z200(centlat_ind+20:centlat_ind-20,centlon_ind-20:centlon_ind+20),Z200_vortex_sub)

;printVarSummary(U200_vortex_sub)

U200_vortex_rm=U200
V200_vortex_rm=V200
Z200_vortex_rm=Z200

;print("max Uvortex_rm = "+max(U200_vortex_rm)+" min Uvortex_rm = "+min(U200_vortex_rm))
;printVarSummary(U200_vortex_rm)
;printVarSummary(U200_vortex_rm(centlat_ind+20:centlat_ind-20,centlon_ind-20:centlon_ind+20))
;printVarSummary(U200_vortex_sub)

U200_vortex_rm(centlat_ind+20:centlat_ind-20,centlon_ind-20:centlon_ind+20)= \
  U200_vortex_rm(centlat_ind+20:centlat_ind-20,centlon_ind-20:centlon_ind+20)-U200_vortex_sub
V200_vortex_rm(centlat_ind+20:centlat_ind-20,centlon_ind-20:centlon_ind+20)= \
  V200_vortex_rm(centlat_ind+20:centlat_ind-20,centlon_ind-20:centlon_ind+20)-V200_vortex_sub
Z200_vortex_rm(centlat_ind+20:centlat_ind-20,centlon_ind-20:centlon_ind+20)= \
  Z200_vortex_rm(centlat_ind+20:centlat_ind-20,centlon_ind-20:centlon_ind+20)-Z200_vortex_sub

MAG_vortex_rm=sqrt(U200_vortex_rm^2+V200_vortex_rm^2)
copy_VarMeta(U200_vortex_rm,MAG_vortex_rm)

;printVarSummary(U200_vortex_rm)

;print(" "+centlat+" "+lat(centlat_ind+28)+" "+lat(centlat_ind-28))
;print(" "+centlon+" "+lon(centlon_ind+28)+" "+lon(centlon_ind-28))

U_comp(case,:,:)=U200_vortex_rm(centlat_ind+80:centlat_ind-80,centlon_ind-80:centlon_ind+80)
V_comp(case,:,:)=V200_vortex_rm(centlat_ind+80:centlat_ind-80,centlon_ind-80:centlon_ind+80)
Z_comp(case,:,:)=Z200_vortex_rm(centlat_ind+80:centlat_ind-80,centlon_ind-80:centlon_ind+80)

wks=gsn_open_wks("png","failure_start_case_"+case)
gsn_define_colormap(wks,cmap)

res=True

res@gsnAddCyclic=False
res@gsnFrame=False
res@gsnDraw=False
res@gsnMaximize=False
res@gsnStringFont=21
res@gsnLeftString="200hPa Wind (kt), Geopotential Height (gpdm)"
res@gsnLeftStringOrthogonalPosF=0.01
res@gsnLeftStringFontHeightF=0.015
res@gsnRightString=name(case)+" @ "+sprinti("%10.0i",yyyymmddhh(case))
;res@gsnRightString="+"+sprinti("%0.2i",counter)+"h"
;res@gsnRightString="+"+sprintf("%4.1f",counter*0.5)+"h"
res@gsnRightStringOrthogonalPosF=0.01
res@gsnRightStringFontHeightF=0.015
res@gsnScalarContour=True

res@vpXF=0.15
res@vpYF=0.95
res@vpHeightF=0.7
res@vpWidthF=0.7

;res@stLineColor=0
;res@stArrowLengthF=0.004
;res@stMinArrowSpacingF=0.004
;res@stMinDistanceF=0.1
;res@stLineThicknessF=2.0
;res@stStepSizeF=0.01

res@mpDataResolution="FinestResolution"
res@mpDataBaseVersion="HighRes"
res@mpLimitMode="LatLon"
;res@mpCenterLatF=centlat
;res@mpCenterLonF=centlon
res@mpMinLatF=centlat-20
res@mpMaxLatF=centlat+20
res@mpMinLonF=centlon-20
res@mpMaxLonF=centlon+20
res@mpMonoFillColor=False
res@mpFillColors=(/-1,-1,-1/)
res@mpOutlineOn=True

res@tmXTOn                = False
res@tmYROn                = False

res@tmXBMode              = "Explicit"
res@tmXBValues            = ispan(100,180,10)
res@tmXBLabels            = (/"100~S~o~N~E","110~S~o~N~","120~S~o~N~E",\
                              "130~S~o~N~","140~S~o~N~E","150~S~o~N~",\
                              "160~S~o~N~E","170~S~o~N~","180~S~o~N~E"/)
res@tmXBMajorLengthF      = 0.01
res@tmXBMajorOutwardLengthF=0.0
res@tmXBMajorThicknessF   = 3.0
;res@tmXBMinorOn           = False
res@tmXBMinorPerMajor     = 4
res@tmXBMinorLengthF      = 0.004
res@tmXBMinorOutwardLengthF=0.0
res@tmXBMinorThicknessF   = 2.0
res@tmXBLabelFont         = 21
res@tmXBLabelFontHeightF  = 0.015
res@tmXBLabelDeltaF       = -0.6

res@tmYLMode              = "Explicit"
res@tmYLValues            = ispan(0,50,10)
res@tmYLLabels            = (/"0~S~o~N~","10~S~o~N~N","20~S~o~N~",\
                              "30~S~o~N~N","40~S~o~N~","50~S~o~N~N"/)
res@tmYLMajorLengthF      = 0.01
res@tmYLMajorOutwardLengthF=0.0
res@tmYLMajorThicknessF   = 3.0
;res@tmYLMinorOn           = False
res@tmYLMinorPerMajor     = 4
res@tmYLMinorLengthF      = 0.004
res@tmYLMinorOutwardLengthF=0.0
res@tmYLMinorThicknessF   = 2.0
res@tmYLLabelFont         = 21
res@tmYLLabelFontHeightF  = 0.015
res@tmYLLabelDeltaF       = -0.3

res@cnFillOn=True
res@cnLevelSelectionMode="ExplicitLevels"
;res@cnLevels:=fspan(0,80,81)
res@cnLevels:=fspan(0,160,81)
res@cnFillOpacityF=1
res@cnLinesOn=False

res@lbBoxSeparatorLinesOn=False
res@lbLabelStride=10

res@pmLabelBarHeightF=0.1

MAG_vortex_rm_kts=MAG_vortex_rm/0.51444

copy_VarMeta(MAG_vortex_rm,MAG_vortex_rm_kts)

plot1=gsn_csm_contour_map(wks,MAG_vortex_rm_kts,res)

res@cnFillOn=False
res@cnLinesOn=True
res@cnInfoLabelOn=False

res@cnLevels:=ispan(floattoint(min(Z200_vortex_rm)),floattoint(max(Z200_vortex_rm)),5)
res@cnLineLabelPlacementMode="constant"
res@cnLineDashSegLenF=0.4
res@cnLineColor="black"
res@cnLineLabelFontColor="black"
res@cnLineLabelFontHeightF=0.010
res@cnLineLabelFontThicknessF=2.0
res@cnLineLabelInterval=1
res@cnLineThicknessF=5.0

plot2=gsn_csm_contour_map(wks,smth9_Wrap(Z200_vortex_rm,0.5,0.25,False),res)

U200_vortex_rm_kts=U200_vortex_rm/0.51444
V200_vortex_rm_kts=V200_vortex_rm/0.51444

copy_VarMeta(U200_vortex_rm,U200_vortex_rm_kts)
copy_VarMeta(V200_vortex_rm,V200_vortex_rm_kts)

res@vcRefMagnitudeF=65
res@vcRefLengthF=0.025
;res@vcGlyphStyle="CurlyVector"
res@vcGlyphStyle="WindBarb"
;res@vcLineArrowThicknessF=3.0
res@vcWindBarbLineThicknessF=2.0
res@vcWindBarbTickLengthF=0.4
res@vcRefAnnoOrthogonalPosF=-0.079
res@vcRefAnnoParallelPosF=1.0
;res@vcMinDistanceF=0.03
res@vcMonoLineArrowColor=True
;res@vcLineArrowColor=1
res@vcRefAnnoArrowUseVecColor=False
res@vcRefAnnoArrowLineColor=1

plot3=gsn_csm_vector_map(wks,U200_vortex_rm_kts(::10,::10),V200_vortex_rm_kts(::10,::10),res)

draw(plot1)
draw(plot2)
draw(plot3)

frame(wks)

end do

rm=systemfunc("rm -rf UVZ_composite_all.nc")

fout=addfile("./UVZ_composite_all.nc","c")
fout->U=U_comp
fout->V=V_comp
fout->Z=Z_comp

end