;Example 1
;Ankur R Desai, 27 May 2010 - 31 Dec 2010
;Citation: Zoitz, J.A., Desai, A.R., Moore, D.J.P., Chadwick, M.,
;2011.

;To run from IDL type at the prompt:
;.compile zobitz_example1_nee.pro
;za_example1

;Requires pre-compiled or in IDL path:
;mcmc.pro

;Helper function to average values
FUNCTION average_arr,arr,int,nan=nan,double=double,sz=sz,med=med,max=max,min=min,tot=tot
  sz = size(arr)
  n = n_elements(arr)
  ngrps = long(n / long(int))
  nout = ngrps * int
  narr = reform(arr[0:nout-1],int,ngrps)
  CASE 1 OF 
    keyword_set(med) : outarr = reform(median(narr,double=double,dim=1))
    keyword_set(max) : outarr = reform(max(narr,nan=nan,dim=1))
    keyword_set(min) : outarr = reform(min(narr,nan=nan,dim=1))  
    keyword_set(tot) : outarr = reform(total(narr,1,nan=nan))
    ELSE : outarr = reform(total(narr,1,nan=nan)/total(finite(narr),1))
  ENDCASE 
  return,outarr
END


;MODELS

;Two GEP models
FUNCTION za_gep1,x,param,_EXTRA=ex
;Michaelis-Menton
  Amax = param[0]
  LUE = param[1]
  I = x[0,*]
  GEP = ((-1.0)*Amax*LUE*I) / ((LUE * I) + Amax)
  return,GEP
END

FUNCTION za_gep2,x,param,_EXTRA=ex
;Temp control
  Amax = param[0]
  LUE = param[1]
  Tmin = param[2]
  Tmax = param[3]
  I = x[0,*]
  T = x[1,*] > Tmin < Tmax
  Tfact = (T-Tmin)/(Tmax-Tmin)
  LUE*=Tfact
  GEP = ((-1.0)*Amax*LUE*I) / ((LUE * I) + Amax)
  return,GEP
END

;Two ER models
FUNCTION za_ER1,x,param,_EXTRA=ex
;Q10
  BR = param[0]
  Q10 = param[1]
  T = x[1,*]
  ER = BR * Q10^((T-10.0)/10.0)
  return,ER
END

FUNCTION za_ER2,x,param,_EXTRA=ex
  BR = param[0]
  Q10 = param[1]
  R = param[2]
  T = x[1,*]+273.15
  ER = BR * Q10^(309.0 * ( (1.0/R) - (1.0/(T-227.0)) ) )
  return,ER
END

;NEE models
FUNCTION za_nee1,x,param,_EXTRA=ex
  return, za_er1(x,param[0:1]) + za_gep1(x,param[2:3]) ;4 param
END

FUNCTION za_nee2,x,param,_EXTRA=ex
  return, za_er2(x,param[0:2]) + za_gep1(x,param[3:4]) ;5 param
END

FUNCTION za_nee3,x,param,_EXTRA=ex
  return, za_er1(x,param[0:1]) + za_gep2(x,param[2:5]) ;6 param
END

FUNCTION za_nee4,x,param,_EXTRA=ex
  return, za_er2(x,param[0:2]) + za_gep2(x,param[3:6]) ;7 param
END

;MAIN PROGRAM
PRO za_example1
;Run 1 za_nee1, (BR, Q10, Amax, LUE)
;Run 2 za_nee2, (BR, Q10, LUE, Tmin, Tmax)
;Run 3 za_nee3, (BR, Q10, R, Amax, LUE)
;Run 4 za_nee4, (BR, Q10, R, LUE, Tmin, Tmax)

;Change this to where the data files are located
  mydocs = '/Users/adesai/Documents/research/data/'
  

;Read Niwot Ridge Temperature, PAR, and NEE data
  restore,mydocs+'mcmc/nwrdata.sav'

;Read QC data
  restore,mydocs+'mcmc/nwr0708flag.sav'

;0 = 2007, 1 = 2008

;replace flagged values with NOT A NUMBER to flag as bad
  bval_o = where(ff07[7,*] NE 1,nbval7)
  bval_p = where(ff08[7,*] NE 1,nbval8)
  farr[8,bval_o,0] = !values.f_nan
  farr[8,bval_p,1] = !values.f_nan

;Convert to hourly from half-hourly, and subset for summer
;_o will be used in the assimilation
;_p will be used for evaluation
  par_o = average_arr(reform(marr[22,5760:13151,0],7392),2,/nan)
  par_p = average_arr(reform(marr[22,5760:13151,1],7392),2,/nan)
  temp_o = average_arr(reform(marr[7,5760:13151,0],7392),2,/nan)
  temp_p = average_arr(reform(marr[7,5760:13151,1],7392),2,/nan)
  nee_o = average_arr(reform(farr[8,5760:13151,0],7392),2)
  nee_p = average_arr(reform(farr[8,5760:13151,1],7392),2)

;Keep track of "good" vobservations
  gval_o = where(finite(nee_o),ngood_o)
  gval_p = where(finite(nee_p),ngood_p)
  
;Create observed and predictor datasets from the above
  x = [transpose(par_o),transpose(temp_o)]
  xp = [transpose(par_p),transpose(temp_p)]
  obs = [transpose(nee_o)]
  obs_p = [transpose(nee_p)]

;The input parameter file for mcmc 
  param1 = { name :       ['BR','Q10',   'Amax',  'LUE'], $
            value :      [1.0,  2.0,   20.0,  0.001], $
            max :        [4.0, 3.5,   50.0, 0.1], $
            min :        [0.25,  1.25, 5.0,   0.0001], $
            changeable : [1  ,  1   ,  1  ,   1  ] }

  param2 = { name :       ['BR','Q10',   'R',  'Amax', 'LUE'], $
            value :      [1.0,  2.0,   50.0, 20.0, 0.001], $
            max :        [4.0, 3.5,   100.0, 50.0, 0.1], $
            min :        [0.25,  1.25, 10.0,   5.0, 0.0001], $
            changeable : [1  ,  1   ,  1  ,   1  , 1] }

  param3 = { name :       ['BR','Q10',  'Amax', 'LUE',  'Tmin', 'Tmax'], $
            value :      [1.0,  2.0,   20.0, 0.001,  0.0, 25.0], $
            max :        [4.0, 3.5,   50.0, 0.1, 20.0, 30.0], $
            min :        [0.25,  1.25, 5.0, 0.0001,   -20, -10.0], $
            changeable : [1  ,  1   ,  1, 1  ,   1  , 1] }

  param4 = { name :       ['BR','Q10',   'R',  'Amax', 'LUE', 'Tmin', 'Tmax'], $
            value :      [1.0,  2.0,   50.0,  20.0, 0.001, 0.0, 25.0], $
            max :        [4.0, 3.5,   100.0, 50.0, 0.1, 20.0, 30.0], $
            min :        [0.25,  1.25, 10.0,  5.0, 0.0001, -20.0, -10.0], $
            changeable : [1  ,  1   ,  1  ,  1, 1  , 1, 1] }

;Call MCMC (unless a previously stored version of the output exists)
;use the standard likelihood function
print,'Running MCMC'
  stop
;type .continue to start the run
  IF file_test(mydocs+'mcmc/za_assim_ex1.sav',/read) THEN restore,mydocs+'mcmc/za_assim_ex1.sav' ELSE BEGIN 
    mcmc,x,obs,param1,/medium,outputll=outputll1,outputvalue=outputvalue1,outputy=y1,model='za_nee1',likelihood='mcmc_likelihood'
    mcmc,x,obs,param2,/medium,outputll=outputll2,outputvalue=outputvalue2,outputy=y2,model='za_nee2',likelihood='mcmc_likelihood'
    mcmc,x,obs,param3,/medium,outputll=outputll3,outputvalue=outputvalue3,outputy=y3,model='za_nee3',likelihood='mcmc_likelihood'
    mcmc,x,obs,param4,/medium,outputll=outputll4,outputvalue=outputvalue4,outputy=y4,model='za_nee4',likelihood='mcmc_likelihood'
    save,filename=mydocs+'mcmc/za_assim_ex1.sav',outputll1,outputll2,outputll3,outputll4,outputvalue1,outputvalue2,outputvalue3,outputvalue4,y1,y2,y3,y4,x,obs,param1,param2,param3,param4
  ENDELSE  

print,'Producing posterior output'
;Get best predictor
  pred1 = za_nee1(xp,outputvalue1[*,0])
  pred2 = za_nee2(xp,outputvalue2[*,0])
  pred3 = za_nee3(xp,outputvalue3[*,0])
  pred4 = za_nee4(xp,outputvalue4[*,0])

;Get range of pred for each (for each accepted parameter set)
; will take a long time if you do all accepted sets
;uncomment if you want to play with this section
;best bet is to reduce number of iterations in for loop to first
;couple thousand or so
;  allpred1 = make_array(n_elements(outputvalue1[0,*]),n_elements(xp[0,*]),/float,value=!values.f_nan)
;  allpred2 = make_array(n_elements(outputvalue2[0,*]),n_elements(xp[0,*]),/float,value=!values.f_nan)
;  allpred3 =  make_array(n_elements(outputvalue3[0,*]),n_elements(xp[0,*]),/float,value=!values.f_nan)

  pred3_ci = make_array(3,4,3696,/float,value=!values.f_nan)
  np = n_elements(outputvalue3[0,*])
  predp = make_array(3,np,/float,value=!values.f_nan)
  BR = outputvalue3[0,*]
  Q10 = outputvalue3[1,*]
  Amax = outputvalue3[2,*]
  LUE = outputvalue3[3,*]
  Tmin = outputvalue3[4,*]
  Tmax = outputvalue3[5,*]
  FOR i = 0l,3695l DO BEGIN
    IF i MOD 100 EQ 0 THEN print,i
    p_T = xp[1,i]
    predp[1,*] = BR * Q10^((p_T-10.0)/10.0)
    p_I = xp[0,i]
    p_T = xp[1,i] > Tmin < Tmax
    Tfact = (p_T-Tmin)/(Tmax-Tmin)
    p_LUE=LUE*Tfact
    predp[2,*] = ((-1.0)*Amax*p_LUE*p_I) / ((p_LUE * p_I) + Amax)
    predp[0,*] = predp[1,*] + predp[2,*]
    nee_a = predp[0,sort(predp[0,*])]    
    pred3_ci[0,*,i] = [mean(nee_a,/nan),median(nee_a),nee_a[long(0.95*np)],nee_a[long(0.05*np)]]
    er_a = predp[1,sort(predp[1,*])]    
    pred3_ci[1,*,i] = [mean(er_a,/nan),median(er_a),er_a[long(0.95*np)],er_a[long(0.05*np)]]
    gep_a = predp[2,sort(predp[2,*])]    
    pred3_ci[2,*,i] = [mean(gep_a,/nan),median(gep_a),gep_a[long(0.95*np)],gep_a[long(0.05*np)]]
  ENDFOR 

  stop

;plot of cum nee,er,gep and CI
  cnee3 = reform(total(pred3_ci[0,*,*],3,/cum))/24.*1.0368
  cer3 = reform(total(pred3_ci[1,*,*],3,/cum))/24.*1.0368
  cgep3 = reform(total(pred3_ci[2,*,*],3,/cum))/24.*1.0368
  ctm = 121+findgen(3696)/24.

  col1 = fsc_color('blue',101)
  col1a = fsc_color('cyan',102)
  col2 = fsc_Color('red',103)
  col2a = fsc_color('pink',104)

  plot,[0,0],[0,0],/nodata,xrange=[121,275],yrange=[-800,500],ytitle='gC m-2',xtitle='Day of Year',charsize=1
  errorbar,ctm,reform(cnee3[0,*]),reform(cnee3[2,*]),reform(cnee3[3,*]),/over,/noline,/noadd,fillcolor=180
  oplot,ctm,cnee3[0,*],thick=3

  errorbar,ctm,reform(cer3[0,*]),reform(cer3[2,*]),reform(cer3[3,*]),/over,/noline,/noadd,fillcolor=col1a
  oplot,ctm,cer3[0,*],thick=3,color=col1

  errorbar,ctm,reform(cgep3[0,*]),reform(cgep3[2,*]),reform(cgep3[3,*]),/over,/noline,/noadd,fillcolor=col2a
  oplot,ctm,cgep3[0,*],thick=3,color=col2


;  allpred4 = make_array(n_elements(outputvalue4[0,*]),n_elements(xp[0,*]),/float,value=!values.f_nan)
;  FOR i = 0,n_elements(outputvalue1[0,*])-1 DO allpred1[i,*] = za_nee1(xp,outputvalue1[*,i])
;  FOR i = 0,n_elements(outputvalue2[0,*])-1 DO allpred2[i,*] = za_nee2(xp,outputvalue2[*,i])
;  FOR i = 0,n_elements(outputvalue3[0,*])-1 DO BEGIN
;    IF i MOD 100 EQ 0 THEN print,i
;    allpred3[i,*] = za_nee3(xp,outputvalue3[*,i])
;  ENDFOR 

;  FOR i = 0,n_elements(outputvalue4[0,*])-1 DO allpred4[i,*] = za_nee4(xp,outputvalue4[*,i])

;Posterior parameter range (best/max/min)
;mcmc.pro sorts the parameter sets by likelihood, so first row is the
;most likely
  post1 = fltarr(3,4)
  post2 = fltarr(3,5)
  post3 = fltarr(3,6)
  post4 = fltarr(3,7)
  FOR i = 0,3 DO post1[*,i] = [outputvalue1[i,0],min(outputvalue1[i,*],/nan),max(outputvalue1[i,*],/nan)]
  FOR i = 0,4 DO post2[*,i] = [outputvalue2[i,0],min(outputvalue2[i,*],/nan),max(outputvalue2[i,*],/nan)]
  FOR i = 0,5 DO post3[*,i] = [outputvalue3[i,0],min(outputvalue3[i,*],/nan),max(outputvalue3[i,*],/nan)]
  FOR i = 0,6 DO post4[*,i] = [outputvalue4[i,0],min(outputvalue4[i,*],/nan),max(outputvalue4[i,*],/nan)]

;Run stats: r2, AIC for each run, compare to predN to obs_p 
;AIC = 2*nparam - 2 ln(L)
;BIC = -2 * ln(L) + nparam * ln(ndata)
  print,'Stats and plots'
  r2 = ([correlate(pred1,obs_p),correlate(pred2,obs_p),correlate(pred3,obs_p),correlate(pred4,obs_p)])^2
  ll = [outputll1[0],outputll2[0],outputll3[0],outputll4[0]]
  bic = ([-2.*outputll1[0]+4.0*alog(ngood_p),-2.*outputll2[0]+5.0*alog(ngood_p),-2.*outputll3[0]+6.0*alog(ngood_p),-2.*outputll4[0]+7.0*alog(ngood_p)])
  aic = ([2.*4.0-2*outputll1[0],2.*4.0-2*outputll2[0],2.*4.0-2*outputll3[0],2.*4.0-2*outputll4[0]])

;Table 1
;Run #, N param, r^2, ll,BIC
  print,'Table 1'
  print,'Nparam 4 5 6 7'
  print,'r2',r2
  print,'Log Likelihood',ll
  print,'BIC',bic
  print

;Table 2
;For each run
;Priors (min,max) and Posteriors (best,min,max) for each parameter
  print,'Table 2'
  FOR i = 0,3 DO print,param1.name[i],param1.value[i],param1.min[i],param1.max[i],post1[*,i]
  FOR i = 0,4 DO print,param2.name[i],param2.value[i],param2.min[i],param2.max[i],post2[*,i]
  FOR i = 0,5 DO print,param3.name[i],param3.value[i],param3.min[i],param3.max[i],post3[*,i]
  FOR i = 0,6 DO print,param4.name[i],param4.value[i],param4.min[i],param4.max[i],post4[*,i]

;Figures

;BR vs Q10 for each
  !p.multi = [0,2,2]
;define colors or pallete here. I used the fsc_color code from the
;David Fanning Coyote library, you can download from
;http://www.dfanning.com/documents/programs.html
;I commented this out so that it will compile without

  col1 = 0 ;fsc_color('blue',101)
  col2 = 100 ;fsc_Color('red',102)
  col3 = 180 ;fsc_Color('yellow',103)
  col4 = 220 ;fsc_color('green',104)
  plot,[0,0],[0,0],xrange=[0.25,4.25],yrange=[1.25,3.75],xtitle='BR',ytitle='Q10'
  oplot,outputvalue2[0,0:1999],outputvalue2[1,0:1999],psym=1,color=col2,symsize=0.5
  oplot,outputvalue3[0,0:1999],outputvalue3[1,0:1999],psym=1,color=col3,symsize=0.5
  oplot,outputvalue4[0,0:1999],outputvalue4[1,0:1999],psym=1,color=col4,symsize=0.5
  oplot,outputvalue1[0,0:1999],outputvalue1[1,0:1999],psym=1,color=col1,symsize=0.5
  xyouts,0.5,3.0,'Ex 1'
  xyouts,0.5,2.8,'Ex 2'
  xyouts,0.5,2.6,'Ex 3'
  xyouts,0.5,2.4,'Ex 4'
  plots,0.8,3.0,psym=1,color=col1,symsize=0.5
  plots,0.8,2.8,psym=1,color=col2,symsize=0.5
  plots,0.8,2.6,psym=1,color=col3,symsize=0.5
  plots,0.8,2.4,psym=1,color=col4,symsize=0.5

;Legend for figure
  plot,[0,0],[0,0],xrange=[0,10],yrange=[0,10],/nodata
  oplot,[2,4],[9,9],thick=2,col=0
  oplot,[2,4],[7,7],linestyle=2,thick=3,color=col1
  oplot,[2,4],[5,5],linestyle=2,thick=3,color=col2
  oplot,[2,4],[3,3],linestyle=2,thick=3,color=col3
  oplot,[2,4],[1,1],linestyle=2,thick=3,color=col4
  xyouts,1,9,'Obs'
  xyouts,1,7,'Ex 1'
  xyouts,1,5,'Ex 2'
  xyouts,1,3,'Ex 3'
  xyouts,1,1,'Ex 4'

;Daily NEE and 10 day period
  doy = findgen(153)+121
  plot,doy,average_arr(obs_p,24,/nan)*1.0368,thick=2,ytitle='gC m-2 day-1',xtitle='Day of Year',yrange=[-5,4],color=0
  oplot,doy,average_arr(pred1,24,/nan)*1.0368,linestyle=2,thick=3,color=col1
  oplot,doy,average_arR(pred2,24,/nan)*1.0368,linestyle=2,thick=3,color=col2
  oplot,doy,average_arR(pred3,24,/nan)*1.0368,linestyle=2,thick=3,color=col3
  oplot,doy,average_arR(pred4,24,/nan)*1.0368,linestyle=2,thick=3,color=col4

  tms = 182+(findgen(120)/24.)
  plot,tms,obs_p[1464:1583],yrange=[-12,8],ytitle='umol m-2 s-1',xtitle='Day of Year',thick=2,color=0
  oplot,tms,pred1[1464:1583],thick=3,linestyle=2,color=col1
  oplot,tms,pred2[1464:1583],thick=3,linestyle=2,color=col2
  oplot,tms,pred3[1464:1583],thick=3,linestyle=2,color=col3
  oplot,tms,pred4[1464:1583],thick=3,linestyle=2,color=col4


;NEW FIGURE

;model 3 -  run with all parameters
  ;show uncertainty in NEE, GEP, TER - cumulative

;1. uncomment line 
;use box - show best and range 95% CI  (sort data for each time point)


stop
END