library(ggplot2)
library(stringr)
library(dplyr)
setwd("C:/Research and study 2021 spring/Updated R for drydep/for_Bin/ESTAR case update")

# This script runs the Surface Tiled Aerosol and Gas Exchange (STAGE) deposition model as it is parameterized
# in the CMAQ v5.3 release. The structure of the code is different here than CMAQ due to the intened use of this code 
# for modeling of field scale fluxes and algorithm development. 

# STAGE land use specific variables adapted from Massad et al. 2010
# land use    | a_cut  | l_leaf
#----------------------------------
# Agiculture  | 0.148  | 0.02
# Hay         | 0.148  | 0.01
# Deciduous   | 0.0318 | 0.05
# Evergreen   | 0.0318 | 0.005
# Mixed Forest| 0.0318 | 0.028
# Shrub       | 0.120  | 0.02
# Grass       | 0.176  | 0.01

# read observation data

#source("read.obs.IT-ISP.r"); Site <- "Ispra"; units <- expression(ppb~m~s^{-1})#; Rst_min <- 153
#source("read.obs.UK_AMo.r"); Site <- "Auchencorth Moss"; units <- expression(ppb~m~s^{-1})#; Rst_min <- 101
#source("read.obs.Bugacpuszta.r"); Site <- "Bugacpuszta"; units <- expression(ppb~m~s^{-1})#; Rst_min <- 75
#source("read.obs.EasterBush.r"); Site <- "Easter Bush"; units <- expression(ppb~m~s^{-1})#; Rst_min <- 63
source("read.obs.harvard_forest.r"); Site <- "Harvard Forest"; units <- expression(ppb~m~s^{-1})#;Rst_min <- 214
#source("read.obs.hyytiala.r"); Site <- "Hyytiala"; units <- expression(ppb~m~s^{-1})#; Rst_min <- 159
#source("read.obs.ramat.r");  Site <- "Ramat"; units <- expression(ppb~m~s^{-1})#; Rst_min <- 74
#source("read.obs.Borden_forest.r");  Site <- "Borden Forest"; units <- expression(ppb~m~s^{-1}); Rst_min <- 101

version <- 'v5.3.2'

plot_site <- str_replace(Site," ","_")
# loop through time and species

# get resistances
source("STAGE.functions.r")
#new lai values Bin Cheng for sensitivity studies
#lainew <- lai*1.2
get.rb_leaf <- function(Name,ustar, lai, Ta, l_leaf) get.rb_Massad(Name,ustar, lai, Ta, l_leaf)


F.mod <- F.obs*0
Ra      <- get.ra(ustar,Ubar,VonK)

LE      <- get.le(soilm, wwlt, wfc, Ta, Ts, Tg, Rs, Rst_min, lai, ustar, Ubar, RH, l_leaf,s_wet,snow,wsat,wres,bslp,pH)
for(i in 1:length(Name)){
   c_atm <- concs[i,]
# note surface temperature is not available so subsitute with air temperature   
   Rb_leaf <- get.rb_leaf(Name[i],ustar, lai, Ta, l_leaf) # Bin Cheng
   Rst     <- get.r_stom(Name[i],soilm, wwlt, wfc, Ta, Ts, Rs, Rst_min, lai, ustar, Ubar, RH)
   Rcut    <- get.r_cut(Name[i],Ta,Ts,RH,pH,l_wet, a_cut, snow, lai)
   Rgrnd   <- get.r_grnd(Name[i],ustar,Ubar,Ta,Ts,Tg,lai,s_wet,snow,soilm,wsat,wfc,wwlt,wres,bslp,pH) 
   Rb_soil <- get.rb_soil(Name,ustar,Ta,lai) # Bin Cheng
   Rinc    <- get.r_inc(ustar,Ubar,lai) #Bin Cheng
   Rsoil   <- Rgrnd - Rb_soil - Rinc
   
# get compensation points
   comp <- get.comp(Name[i],Tg,gamma.g,gamma.st,gamma.cut, pH, soilm, bulk_d)
   c_stom <- comp$Stom
   c_cut  <- comp$Cut
   c_soil <- comp$Soil

#Calculate leaf compensation point follwing Nimitz et al 2001 modified to account for a cuticular compsensation point
   c_leaf <- (c_atm/(Ra*Rb_leaf)+                                                                   # Atmospheric Component
              c_stom*(1.0/(Ra*Rst)+1.0/(Rb_leaf*Rst)+1.0/(Rgrnd*Rst))+                              # Stomatal Component
              c_cut*(1.0/(Ra*Rcut)+1.0/(Rb_leaf*Rcut)+1.0/(Rgrnd*Rcut))+                            # Cuticular Component
              c_soil/(Rb_leaf*Rgrnd))/                                                              # Soil Component
             (1.0/(Ra*Rb_leaf) +1.0/(Ra*Rst) +1.0/(Ra*Rcut)+1.0/(Rb_leaf*Rgrnd)+1.0/(Rb_leaf*Rcut)+
              1.0/(Rb_leaf*Rst)+1.0/(Rgrnd*Rst)+1.0/(Rgrnd*Rcut)) # Least common denominator

#Calculate the canopy compensation point follwing Nimitz et al 2001 modified to account for a cuticular compsensation point
   c_z0     = (c_atm/Ra+c_leaf/Rb_leaf+c_soil/Rgrnd)/(1.0/Ra+1.0/Rb_leaf+1.0/Rgrnd)           

# calculate the flux
   flux.mod  <- -veg  * (c_atm-c_z0)/Ra -                 # air-vegetation flux
              (1-veg) * (c_atm-c_soil)/( Ra + Rgrnd )     # air-soil flux
   flux.mod  <- flux.mod*flux.fact # ug/m2/s -> ng/m2/s but depened on the units of the input file
   F.mod[i,] <- flux.mod
# component fluxes (flux.mod = flux.stom + flux.cut + flux.soil following Kirshhoff's law)
   flux.stom <- -(c_leaf-c_stom)/Rst*flux.fact # ng/m2/s
   flux.cut  <- -(c_leaf-c_cut)/Rcut*flux.fact # ng/m2/s
   flux.soil <- -(c_z0-c_soil)/Rgrnd*flux.fact # ng/m2/s

   pie.data <- data.frame(sink = c("Stomatal","Cuticular","Soil"),
                          fluxes   = c(sum(flux.stom),sum(flux.cut),sum(flux.soil))/sum(flux.mod))
   pie.data <- pie.data %>% 
               arrange(desc(sink)) %>%
               mutate(prop = fluxes / sum(pie.data$fluxes) *100) %>%
               mutate(ypos = cumsum(prop)- 0.5*prop )

   pie <- ggplot(pie.data, aes(x="", y=prop, fill=sink))+
          geom_bar(width = 1, stat = "identity",color="white") + 
          coord_polar("y", start=0)+ 
          ggtitle(paste(Site,"Component Fluxes")) +
          theme_void() +
          geom_text(aes(y = ypos, label = paste(round(fluxes*100),"%",sep='')), color = "white", size=6) 
# observed flux
   flux.obs  <- F.obs[i,]
# calcuate the deposition velocity
   vd     <-  veg * (1.0-c_z0/c_atm)/Ra + (1-veg) /( Ra + Rgrnd )
   vd_tst <- veg/(Ra+1/(1/(Rb_leaf+1/(1/Rst+1/Rcut))+1/Rgrnd)) + (1-veg) /( Ra + Rgrnd )
   vd_sen <- veg/(1/(1/(Rb_leaf+1/(1/Rst+1/Rcut))+1/Rgrnd)) + (1-veg) /(Rgrnd )
   vd_no_ra <- (1/(Rb_leaf+1/(1/Rst+1/Rcut))+1/Rgrnd)
# calculate O3 flux based on vd_***
   Flux_tst <- -c_atm * vd_tst
   Flux_sen <- -c_atm * vd_sen
   flux.no.ra <- -c_atm * vd_no_ra
#head(Flux_tst)
#head(Flux_sen)
#tail(Flux_tst)
#tail(Flux_sen)
   
   hours <- paste(date.time$h)
   valid <- which(nchar(hours)==1)
   hours[valid] <- paste('0',hours[valid],sep='')
   months <- paste(date.time$mon+1)
   valid <- which(nchar(months)==1)
   months[valid] <- paste('0',months[valid],sep='')
   
   plot.data <- data.frame(hour    = c(hours,hours),
                           mon     = c(months,months),
                           value   = c(F.obs[i,],F.mod[i,]),
                           model   = c(rep('Obs',length(date.time)),rep('STAGE',length(date.time))))


   vd.data <- data.frame(hour    = c(hours,hours),
                         mon     = c(months,months),
                         value   = c(vd_o3,vd),
                         model   = c(rep('Obs',length(date.time)),rep('STAGE',length(date.time))))


   comp.data <- data.frame(hour    = c(hours,hours,hours),
                           mon     = c(months,months,months),
                           flux    = c(flux.stom,flux.cut,flux.soil),
                           model   = c(rep('Stomatal',length(date.time)),rep('Cuticular',length(date.time)),rep('Soil',length(date.time))))

   hourly.plot <- ggplot(plot.data,aes(x=hour,y=value,fill=model))+
        geom_boxplot(outlier.shape=NA) + 
        stat_summary(fun=mean, geom="point", shape=16, size=1)+
        ggtitle(Name[i]) +
        xlab("Time (EST)") + 
        ylab(units) +
        ylim(min(quantile(F.mod[i,],0.01),quantile(F.obs[i,],0.01)),min(quantile(F.mod[i,],0.99),quantile(F.obs[i,],0.99)))

   monthly.plot <- ggplot(plot.data,aes(x=mon,y=value,fill=model))+
        geom_boxplot(outlier.shape=NA) + 
        stat_summary(fun=mean, geom="point", shape=16, size=1)+
        ggtitle(Name[i]) +
        xlab("Month") + 
        ylab(units) +
        ylim(min(quantile(F.mod[i,],0.01),quantile(F.obs[i,],0.01)),min(quantile(F.mod[i,],0.99),quantile(F.obs[i,],0.99)))
   
   model.plot <- ggplot(plot.data,aes(x=model,y=value,fill=model))+
        geom_boxplot(outlier.shape=NA) + 
        stat_summary(fun=mean, geom="point", shape=16, size=1)+
        ggtitle(Name[i]) + 
        ylab(units) +
        ylim(min(quantile(F.mod[i,],0.01),quantile(F.obs[i,],0.01)),min(quantile(F.mod[i,],0.99),quantile(F.obs[i,],0.99)))
   if(is.nan(max(L))==F){
      L.val <- which(L > -100)
      L     <- L[L.val]
      LE    <- LE[L.val]
      le.data   <- data.frame(hour = c(hours[L.val],hours[L.val]),
                           mon  = c(months[L.val],months[L.val]),
                           LE      = c(L,LE), # Place holder for latent heat flux
                           model  = c(rep('Obs',length(L)),rep('STAGE',length(LE))))
      le.hourly <- ggplot(le.data,aes(x=hour,y=LE,fill=model))+
        geom_boxplot(outlier.shape=NA) + 
        xlab("Time (EST)") + 
        ylab(expression(W~m^{-2})) +
        stat_summary(fun=mean, geom="point", shape=16, size=1)

      le.monthly <- ggplot(le.data,aes(x=mon,y=LE,fill=model))+
        geom_boxplot(outlier.shape=NA) + 
        xlab("Month") + 
        ylab(expression(W~m^{-2})) +
        stat_summary(fun=mean, geom="point", shape=16, size=1)

      le.plot <- ggplot(le.data,aes(x=model,y=LE,fill=model))+
        geom_boxplot(outlier.shape=NA) + 
        stat_summary(fun=mean, geom="point", shape=16, size=1)+
        ggtitle("LE") + 
        ylab(expression(W~m^{-2}))
      
      out.name <- paste("LE_",plot_site,"_hourly.png",sep="")
      png(file=out.name,bg='white',width=10,height=6,units='in',res=150) 
      plot(le.hourly)
      dev.off()
      out.name <- paste("LE_",plot_site,"_data.png",sep="")
      png(file=out.name,bg='white',width=10,height=6,units='in',res=150) 
      plot(le.data)
      dev.off()
      out.name <- paste("LE_",plot_site,"_monthly.png",sep="")
      png(file=out.name,bg='white',width=10,height=6,units='in',res=150) 
      plot(le.monthly)
      dev.off()
      out.name <- paste("LE_",plot_site,"_plot.png",sep="")
      png(file=out.name,bg='white',width=10,height=6,units='in',res=150) 
      plot(le.plot)
      dev.off()
      
      df = data.frame(hours[L.val], months[L.val], L, LE)
      write.csv(x=df, file="C:/Research and study 2021 spring/Updated R for drydep/for_Bin/LE analysis/ESTAR LE.csv")
   } 
   out.name <- paste(Name[i],"_",plot_site,"_hourly.png",sep="")
   png(file=out.name,bg='white',width=10,height=6,units='in',res=150) 
   plot(hourly.plot)
   dev.off()
   out.name <- paste(Name[i],"_",plot_site,"_monthly.png",sep="")
   png(file=out.name,bg='white',width=10,height=6,units='in',res=150) 
   plot(monthly.plot)
   dev.off()
   out.name <- paste(Name[i],"_",plot_site,"_box.png",sep="")
   png(file=out.name,bg='white',width=5,height=6,units='in',res=150) 
   plot(model.plot)
   dev.off()
   out.name <- paste(Name[i],"_",plot_site,"_pie.png",sep="")
   png(file=out.name,bg='white',width=5,height=6,units='in',res=150) 
   plot(pie)
   dev.off()

#   comp.flux <- summary(lm(F.obs[1,]~flux.stom+flux.cut+flux.soil-1))

# Write output
   out.name <- paste(Name[i],"_",plot_site,"_STAGE_fluxes.csv",sep="")
   out.dat  <- list('Date and Time'                                    = date.time,
                    'Observed Flux'                                    = flux.obs,
                    'Modeled Flux'                                     = flux.mod,
                    'Modeled Sotmatal Flux'                            = flux.stom,
                    'Modeled Cuticular Flux'                           = flux.cut,
                    'Modeled Soil Flux'                                = flux.soil,
                    'Concentration'                                    = c_atm,
                    'Compensation point'                               = c_z0, 
                    'Aerodynamic resistance s/m'                       = Ra,
                    'Leaf quasi-laminar boundary layer resistance s/m' = Rb_leaf,
                    'Soil quasi-laminar boundary layer resistance s/m' = Rb_soil,
                    'In canopy resistance s/m'                         = Rinc,
                    'Stomatal resistance s/m'                          = Rst,
                    'Cuticular resistance s/m'                         = Rcut,
                    'Bluk soil resistance s/m'                         = Rgrnd,
                    'Soil resistance s/m'                              = Rsoil,
                    'Beta square or Vonk values'                       = VonK,
                    'H value'                                          = H,
                    'OBS Deposition velocity m/s'                      = vd_o3,
                    'ESTAR Deposition velocity m/s'                    = vd,
                    'Vd sensitivity removing Ra'                       = vd_sen,
                    'O3 flux sensitivity removing Ra'                  = Flux_sen,
                    'vd_no_ra'                                         = vd_no_ra,
                    'flux_no_ra'                                       = flux.no.ra)                  
   write.csv(out.dat,out.name, row.names = F)
   print(paste(Name[i],'NMB',mean(F.mod[i,]-F.obs[i,])/mean(F.obs[i,])))
   print(paste(Name[i],'NME',mean(abs(F.mod[i,]-F.obs[i,]))/mean(F.obs[i,])))
   print(paste(Name[i],'R',cor(F.mod[i,],F.obs[i,])))
#   comp.flux
   print(paste('Latent Heat','NMB',mean(LE-L)/mean(L)))
   print(paste('Latent Heat','NME',mean(abs(LE-L))/mean(L)))
   print(paste('Latent Heat','R',cor(LE,L)))
   
}


#out.dat  <- list('Date and Time'                                    = date.time[L.val],
#                 'L'                                                = L,
#                 'LE'                                               = LE)
#write.csv(x= out.dat,file="C:/Research and study 2021 spring/Updated R for drydep/for_Bin/LE analysis/ESTAR case LE.csv")

