library(SoilR)
library(FME)
library(aridec)
library(dplyr)
library(MuMIn)

############## MODEL FITTING FOR EACH ARIDEC ENTRY
# Load function for a single pool model
onepFit=function(timeSeries, initialCarbon=100, inipars){
  complete=data.frame(time=timeSeries[complete.cases(timeSeries),1],Ct=timeSeries[complete.cases(timeSeries),2])
  n=nrow(complete)
  tt=seq(from=0, to=tail(complete[,1],1), length.out = 500)
  
  Func=function(pars){
    mod=SoilR::OnepModel(t=tt,k=pars[1], C0=initialCarbon, In=0)
    Ct=SoilR::getC(mod)
    return(data.frame(time=tt, Ct=Ct))
  }
  
  costFunc=function(pars){
    output=Func(pars)
    return(modCost(model=output, obs=complete))
  }
  
  Fit=modFit(f=costFunc, p=inipars, method="Marq", lower= 0, upper=Inf)
  bestMod=Func(pars=Fit$par)
  plot(complete, ylim=c(0,1.2*max(complete[,2])))
  lines(bestMod)
  AIC=((n+2*(length(Fit$par)+1))/(n))+log(Fit$ms) # formula from Shumway & Stoffer
  AICc=log(Fit$ms)+((n+length(Fit$par)+1)/max(0,n-length(Fit$par)-2)) # for small sample sizes
  SoilRmodel=SoilR::OnepModel(t=tt,k=Fit$par[1], C0=initialCarbon, In=0)
  return(list(FMEmodel=Fit, SoilRmodel=SoilRmodel, AIC=AIC, AICc=AICc, TransitTime=1/(Fit$par)))
}

# Load function for a two pool parallel model with a full set of parameters
twoppFit=function(timeSeries, initialCarbon=100, inipars){
  complete=data.frame(time=timeSeries[complete.cases(timeSeries),1],Ct=timeSeries[complete.cases(timeSeries),2])
  n=nrow(complete)
  tt=seq(from=0, to=tail(complete[,1],1), length.out = 500)
  
  Func=function(pars){
    mod=SoilR::TwopParallelModel(t=tt,ks=pars[1:2], C0=initialCarbon*c(pars[3], 1-pars[3]), In=0, gam=0) 
    Ct=SoilR::getC(mod)
    return(data.frame(time=tt, Ct=rowSums(Ct)))
  }
  
  costFunc=function(pars){
    output=Func(pars)
    return(modCost(model=output, obs=complete))
  }
  
  Fit=modFit(f=costFunc, p=inipars, method="Marq", lower=rep(0,3), upper=c(Inf, Inf, 1))
  bestMod=Func(pars=Fit$par)
  AIC=((n+2*(length(Fit$par)+1))/(n))+log(Fit$ms) # formula from Shumway & Stoffer
  AICc=log(Fit$ms)+((n+length(Fit$par)+1)/max(0,n-length(Fit$par)-2)) # for small sample sizes
  if(AICc == Inf) stop("Time series is too short. No degrees of freedom")
  SoilRmodel=SoilR::TwopParallelModel(t=tt,ks=Fit$par[1:2], C0=initialCarbon*c(Fit$par[3], 1-Fit$par[3]), In=0, gam=0) 
  TransitTime=transitTime(A=diag(-(Fit$par[c(1,2)])),u=matrix(initialCarbon*c(Fit$par[3], 1-Fit$par[3]), ncol=1), a=0, q=0.5)
  return(list(FMEmodel=Fit, SoilRmodel=SoilRmodel, AIC=AIC, AICc=AICc, TransitTime=TransitTime))
}

# Two pool parallel model restricting parameter 3 
RtwoppFit=function(timeSeries, initialCarbon=100, lignin, inipars){
  complete=data.frame(time=timeSeries[complete.cases(timeSeries),1],Ct=timeSeries[complete.cases(timeSeries),2])
  n=nrow(complete)
  tt=seq(from=0, to=tail(complete[,1],1), length.out = 500)
  
  Func=function(pars){
    mod=SoilR::TwopParallelModel(t=tt,ks=pars[1:2], C0=initialCarbon*c(1-lignin, lignin), In=0, gam=0) 
    Ct=SoilR::getC(mod)
    return(data.frame(time=tt, Ct=rowSums(Ct)))
  }
  
  costFunc=function(pars){
    output=Func(pars)
    return(modCost(model=output, obs=complete))
  }
  
  Fit=modFit(f=costFunc, p=inipars, method="Marq", lower=rep(0,2), upper=c(Inf, Inf))
  bestMod=Func(pars=Fit$par)
  AIC=((n+2*(length(Fit$par)+1))/(n))+log(Fit$ms) # formula from Shumway & Stoffer
  AICc=log(Fit$ms)+((n+length(Fit$par)+1)/max(0,n-length(Fit$par)-2)) # for small sample sizes
  if(AICc == Inf) stop("Time series is too short. No degrees of freedom")
  SoilRmodel=SoilR::TwopParallelModel(t=tt,ks=Fit$par[1:2], C0=initialCarbon*c(1-lignin, lignin), In=0, gam=0) 
  TransitTime=transitTime(A=diag(-(Fit$par[c(1,2)])),u=matrix(initialCarbon*c(1-lignin, lignin), ncol=1), a=0, q=0.5)
  return(list(FMEmodel=Fit, SoilRmodel=SoilRmodel, AIC=AIC, AICc=AICc, TransitTime=TransitTime))
}

# Function for a two pool series model restricting parameter 4
twopsFit=function(timeSeries, initialCarbon=100, lignin, inipars){
  complete=data.frame(time=timeSeries[complete.cases(timeSeries),1],Ct=timeSeries[complete.cases(timeSeries),2])
  n=nrow(complete)
  tt=seq(from=0, to=tail(complete[,1],1), length.out = 500)
  
  Func=function(pars){
    mod=SoilR::TwopSeriesModel(t=tt,ks=pars[1:2], a21=pars[1]*pars[3], C0=initialCarbon*c(1-lignin, lignin), In=0)
    Ct=SoilR::getC(mod)
    return(data.frame(time=tt, Ct=rowSums(Ct)))
  }
  
  costFunc=function(pars){
    output=Func(pars)
    return(modCost(model=output, obs=complete))
  }
  
  Fit=modFit(f=costFunc, p=inipars, method="Marq", lower=rep(0,3), upper=c(Inf, Inf, 1))
  bestMod=Func(pars=Fit$par)
  AIC=((n+2*(length(Fit$par)+1))/(n))+log(Fit$ms) # formula from Shumway & Stoffer
  AICc=log(Fit$ms)+((n+length(Fit$par)+1)/max(0,n-length(Fit$par)-2)) # for small sample size
  if(AICc == Inf) stop("Time series is too short. No degrees of freedom")
  SoilRmodel=SoilR::TwopSeriesModel(t=tt,ks=Fit$par[1:2], a21=Fit$par[1]*Fit$par[3], C0=initialCarbon*c(1-lignin, lignin), In=0)
  TransitTime=transitTime(A=matrix(c(-Fit$par[1],Fit$par[3],0,-Fit$par[2]),ncol=2),u=matrix(initialCarbon*c(1-lignin, lignin), ncol=1), a=0, q=0.5)
  return(list(FMEmodel=Fit, SoilRmodel=SoilRmodel, AIC=AIC, AICc=AICc, TransitTime=TransitTime))
}

# load entry
db=loadEntries("~/aridec/data/")
entry=db[[150]] # select entry number
df=entry$timeSeries
if (entry$variables$V1$units == "years") { # transform all time units to months
  df=mutate(df, Time=Time*12)
} else if (entry$variables$V1$units == "weeks") { 
  df=mutate(df, Time=Time/4.3)
} else if (entry$variables$V1$units == "days") { 
  df=mutate(df, Time=Time/30)
} 
ts=df[,c(1,5)]

# fit the one pool model
M1=onepFit(ts, inipars=0.1) 
log(2)/M1$FMEmodel$par #calculate median transit time

print(entry$citationKey)
# Are more complex models identifiable?
collinearity=read.csv("~/Documents/collinearity_fulldb.csv") # collinearity indexes from Sarquis et al. 2022
# can I fit a 2 pool parallel model with all parameters?
filter(collinearity, Entry == "Teklay2007" & Model == "2pp" & N == "3")
# can I fit a 2 pool parallel model with a fixed X3 parameter?
filter(collinearity, Entry == "Teklay2007" & Model == "2pp" & X3 == "0")
# can I fit a 2 pool series model with a fixed X4 parameter?
filter(collinearity, Entry == "Teklay2007" & Model == "2ps" & N == "3" & X4 == "0")

# fit the two pool parallel model with a full set of parameters
M2=twoppFit(ts, inipars=c(0.08, 0.05, 0.5)) # adjust initial parameters as necessary
months=M2$SoilRmodel@times
Ct=getC(M2$SoilRmodel)
matplot(months,Ct, type="l",lty=1, col=2:3, ylim=c(0,100), ylab="Mass remaining (%)")
legend("topright", c("Fast","Slow"), lty=1, col=c(2:3), bty="n")
points(df[,c(1,3)])
lines(months, rowSums(Ct), lwd=2)
M2$TransitTime$quantiles # get median transit time

# Only if the models are identifiable! Are there lignin content data available? If not, skip the next models...
print(entry$variables$V5$litter) # check the species
lignindf=data.frame(species=entry$initConditions$species,lignin=entry$initConditions$lignin)
lignin=(lignindf[2,2])/100 # change the first number for species

# fit the two pool parallel model with a restricted parameter 3
M2=RtwoppFit(ts, lignin=lignin, inipars=c(0.25, 0.13))
months=M2$SoilRmodel@times
Ct=getC(M2$SoilRmodel)
matplot(months,Ct, type="l",lty=1, col=2:3, ylim=c(0,100), ylab="Mass remaining (%)")
legend("topright", c("Fast","Slow"), lty=1, col=c(2:3), bty="n")
points(df[,c(1,3)])
lines(months, rowSums(Ct), lwd=2)
M2$TransitTime$quantiles # get median transit time

# fit the two pool series model with a restricted parameter 4
M3=twopsFit(ts, lignin=lignin, inipars=c(0.06, 0.02, 0.5)) 
months=M3$SoilRmodel@times
Ct=getC(M3$SoilRmodel)
matplot(months,Ct, type="l",lty=1, col=c(2,3), ylim=c(0,100), ylab="Mass remaining (%)")
legend("topright", c("Total","Slow","Fast"), lty=1, col=c(1,3,2), bty="n")
points(df[,c(1,3)])
lines(months, rowSums(Ct), lwd=2)
M3$TransitTime$quantiles # get median transit time

# save all model outputs to one list per variable
saveRDS(list(M1=M1, M2=M2), "~/Documents/model outputs/Teklay2007_4.rds") # for studies with 2 or 3 models
saveRDS(M1, "~/Documents/model outputs/Teklay2007_5.rds") # for studies with just one model

############ CALCULATE AICc WEIGHTS

AICweights= vector("list",128) # create an empty vector for the 128 models
for (i in 1:128) {
  AICweights[[i]]=Weights(subset(finaldf, number == i, select = 'AICc'))
} 
finaldf=cbind(finaldf,unlist(AICweights)) # paste to table of pre-averaging results 

# averaging median transit time using Akaike weights
df_list <- 
  finaldf %>% 
  group_by(number) %>% 
  group_split()
avgMedianTT=unlist(lapply(df_list, function(x) {sum(x$AICweights*x$medianTT)}))
finaldf=cbind(finaldf,avgMedianTT=rep(avgMedianTT,each=3)) # paste averaged median transit time to results table

# averaging mean transit time using Akaike weights
df_list <- 
  finaldf %>% 
  group_by(number) %>% 
  group_split()
avgMeanTT=unlist(lapply(df_list, function(x) {sum(x$AICweights*x$transitTime)}))
finaldf=cbind(finaldf,avgMeanTT=rep(avgMeanTT,each=3)) # paste to results table

# calculate unconditional variance (Burnham and Anderson, page 345)
# AICweights*(ms+(medianTT-avgMedianTTs)^2) this part was done in the results spreadsheet
# ms: the mean squared residuals, i.e. ssr/length(residuals)
# this loop finishes the calculation of the variance 
for (i in 1:128) {
  print(sum(subset(finaldf, number == i, select = 'U.Q')))
} # paste results in dataframe

# calculating confidence intervals
# first get the critical value of a t-distribution for a 95% confidence interval
tdist=sapply(finaldf$degFree, function(x) {qt(.975, x)}) 
# average tdist for each aridec entry
for (i in 1:128) {
  print(mean(unlist(subset(finaldf, number == i, select = 'tdist'))))
} # paste this in the dataframe

# Create a table of only for averaged models results
finalavg=subset(finaldf, model == '1p', select = c('ID','Vn','GAI','MAT_wc','lignin','MAP','MAT_ari','radanual',
                                                   'avgMedianTT', 'avgMeanTT', 'AICvar','meantdist'))

# calculate confidence intervals
c.int= finalavg$meantdist * sqrt(finalavg$AICvar) 
finalavg=cbind(finalavg,c.int) # paste to averaged results table

# calculate the quotient median/log2*mean (figure 5)
quotient=finalavg$avgMedianTT/(log(2)*finalavg$avgMeanTT)
finalavg=cbind(finalavg,quotient) # paste to results table