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

begin

; Set up the directories where our data is located
  base_file = "data/taylor/t2_daily_State_Data_US_"
  cfsr_file = "data/taylor/t2_daily_State_Data_US_cfsr.nc"
  fname = "plt_taylor_temporal_co"

; List of our 24 ensemble members we want to plot
  CASES = (/"ck6m","ck6y","cktm","ckty","cn6m","cn6y","cntm","cnty","ct6m","ct6y","cttm","ctty",\
            "rk6m","rk6y","rktm","rkty","rn6m","rn6y","rntm","rnty","rt6m","rt6y","rttm","rtty"/)
  numCASES = dimsizes(CASES)

; Eleven years of ensemble runs
  YEARS = (/"1990","1991","1992","1993","1994","1995","1996","1997","1998","1999","2000"/)
  numYEARS = dimsizes(YEARS)

; Summer months we want in the analysis
  MONS = (/"06","07","08"/)
  numMONS = dimsizes(MONS)

; Set up arrays for the ensemble runs and the gridded observations (CFSR)
  t2_case = new((/numCASES,numYEARS,numMONS/),float)
  t2_cfsr = new((/numYEARS,numMONS/),float)

; Crate data arrays for the ensembles. We loop through each ensemble, each year, and
;  each month. Our files have 93 time steps in them so we need to pick out the days
;  for each month. Then we average over the values in the shapefile, and then average over
;  the month.
  do i=0,numCASES-1
    do j=0,numYEARS-1
      file3 = CASES(i)
      f3 = addfile(base_file+file3+"_CFSR.nc","r")
      t2 = f3->T2MAX_Colorado
      do k=0,numMONS-1
        if(k.eq.0) then
          jun = t2(j,0:29,:)
          juni = dim_avg_n(jun,1)
          junavg = avg(juni)
          t2_case(i,j,k) = junavg
        else if(k.eq.1) then
          jul = t2(j,30:60,:)
          juli = dim_avg_n(jul,1)
          julavg = avg(juli)
          t2_case(i,j,k) = julavg
        else if(k.eq.2) then
          aug = t2(j,61:91,:)
          augi = dim_avg_n(aug,1)
          augavg = avg(augi)
          t2_case(i,j,k) = augavg
        end if
        end if
        end if
      end do
    end do
  end do

  delete([/jun,juni,junavg,jul,juli,julavg,aug,augi,augavg/])

; Next we create arrays for the cfsr data. We loop through each year and each month.
;  This data set has 365 time steps. We need to pick out the correct days of the month.
;  We then average over the values in the shapefile and then average over the month.
  do i=0,numYEARS-1
    cfsr = addfile(cfsr_file,"r")
    t2c = cfsr->T2MAX_Colorado-273.15
    do j=0,numMONS-1
      if(j.eq.0) then
        jun = t2c(i,151:180,:)
        juni = dim_avg_n(jun,1)
        junavg = avg(juni)
        t2_cfsr(i,j) = junavg
      else if(j.eq.1) then
        jul = t2c(i,181:211,:)
        juli = dim_avg_n(jul,1)
        julavg = avg(juli)
        t2_cfsr(i,j) = julavg
      else if(j.eq.2) then
        aug = t2c(i,212:242,:)
        augi = dim_avg_n(aug,1)
        augavg = avg(augi)
        t2_cfsr(i,j) = augavg
      end if
      end if
      end if
    end do
    delete(t2c)
  end do

  vars = (/"T2MAX"/)
  numVars = dimsizes(vars)

; Set up our workstation and choose an NCL colortable
  wks = gsn_open_wks("X11",fname)  
  gsn_define_colormap(wks,"seaice_2")

; Set up our arrays for the taylor diagram calculations
  FValue = -99999.
  ratio      = new ((/numCASES, numVars/), float )  
  cc         = new ((/numCASES, numVars/), float ) 
  bias       = new ((/numCASES, numVars/), float ) 
  cfsr_varc   = new ((/numVars/), float ) 

; Set up the marker types and colors for the ensemble labels
  print("Num Vars : " + numVars )
  marks   = new ( numCASES , integer )
  marks = 16
  marks(11:) = 11
  numbers = new ( numCASES , integer )
  do i=0,numCASES-1
    numbers(i) = i
    numbers(i) = (1+i)*1
  end do

;#################################################################################################
; Here we calculate the variance for the CFSR data

  iv = 0

  cfsr_varc(iv) = variance(t2_cfsr)
  iv = iv+1

  print(cfsr_varc)
;#################################################################################################

; Now we want to loop through each case and calculate the variance, ratio of variences,
;  correlations, and bias. These are the numbers that go into the taylor diagram.
  do icase = 0,numCASES-1
    iv = 0
    print("Working on case " + CASES(icase) )

    varc = variance(t2_case(icase,:,:))
    ratio(icase,iv) = varc/cfsr_varc(iv)
    cc(icase,iv) = escorc(ndtooned(t2_cfsr),ndtooned(t2_case(icase,:,:)))
    bias_S = avg(t2_case(icase,:,:)-t2_cfsr)
    bias(icase,iv) = (bias_S/avg(t2_cfsr))*100.
    iv = iv+1
  end do

; Set up the resources for the Taylor Diagram
  res   = True                           ; default taylor diagram
  res@Markers      = marks
  res@Colors       = numbers
  res@varLabels    = vars
  res@caseLabels   = CASES
  res@centerDiffRMS = True
  res@varLabelsYloc = 1.5
  res@legendHeight = 0.0250*numCASES
  res@tiMainString = "Taylor Diagram for T2MAX JJA - Colorado"

; Plot the taylor diagram using this special NCL function
  plot  = taylor_diagram(wks,ratio,cc,res)

  frame(wks)
   
end