# https://cran.r-project.org/web/packages/RRphylo/vignettes/RRphylo.html
rm(list=ls())
setwd("~/Dropbox/FCU/Teaching/Mentoring/2022Spring-Current_Professor/2022Fall/YourueiMin/rcode_dcj/empricaldataanalysis/matchedtreetraitTony")

library(geiger)
library(ape)
#install.packages("extraDistr")
library(extraDistr)
#install.packages("LaplacesDemon")
library(LaplacesDemon)
#install.packages("RRphylo")
library(RRphylo)
#install.packages("phytools")
library(phytools)
#?mle
library(stats4)
library(qpcR)

compute_betas_rrphylo <- function(tree, y) { 
  # Supporting Functions
  makeL(tree)-> L
  
  makeL1(tree)->L1
  #y<-fastBM(tree,internal=F)
  optL <- function(lambda){
    #lambda<-0.001
    y <- scale(y)#(y-mean(y))/sd(y)
    betas <- (solve(t(L) %*% L + lambda * diag(ncol(L))) %*%t(L)) %*% (as.matrix(y) - rootV)
    y.hat <- (L %*% betas) + rootV
    Rvar <- array()
    for (i in 1:Ntip(t)) {
      #    i<-2
      ace.tip <- betas[match(names(which(L[i, ] != 0)),rownames(betas)), ]
      mat = as.matrix(dist(ace.tip))
      Rvar[i] <- sum(mat[row(mat) == col(mat) + 1])
    }
    #  the rate variation within clades
    abs(1 - (var(Rvar)))  #+ (mean(as.matrix(y))/mean(y.hat))) #y is scaled so zero mean
  }
  
  # Processing
  t <- tree
  toriginal <- t
  yoriginal <- y <- treedataMatch(tree, y)[[1]]
  Loriginal <- L <- makeL(t)
  L1original <- L1 <- makeL1(t)
  u <- data.frame(yoriginal, (1/diag(vcv(toriginal))^2))
  u <- u[order(u[, ncol(u)], decreasing = TRUE), ]
  u1 <- u[1:(nrow(u) * 0.1), , drop = FALSE]
  rootV <- c(apply(u1[, 1:(ncol(u1) - 1), drop = FALSE], 2, 
                   function(x) weighted.mean(x, u1[, dim(u1)[2]])))
  
  
  h <- mle(optL, start = list(lambda = 1), method = "L-BFGS-B", upper = 10, lower = 0.001)
  lambda <- h@coef
  
  betas.rrphylo <- (solve(t(L) %*% L + lambda * diag(ncol(L))) %*% t(L)) %*% (as.matrix(y) - rootV)
  betas.rrphylo

  plot(tree)
#  edgelabels()
  tiplabels()
  nodelabels()

  nodeset<-rownames(betas.rrphylo)[1:(Ntip(tree)-1)]
  tipset<-rownames(betas.rrphylo)[(Ntip(tree):(2*Ntip(tree)-1))]
  tipset<-gsub("t", "",tipset)
  #rownames(betas.rrphylo)
  rownames(betas.rrphylo)<-c(nodeset,tipset)

  
  return(list(h=h,betas.rrphylo=betas.rrphylo))
}

ar1.negloglike <- function(param, betas.rrphylo=betas.rrphylo,tree=tree){
  phi<-param["phi"]
  sigsq<-param["sigsq"]
  
  tree<-reorder(tree,"postorder")
  tree$root.edge<-0
  N<-dim(tree$edge)[1]
  anc<-tree$edge[,1]
  des<-tree$edge[,2]
  
  neglogl<-0
  for(index in N:1){
    desindex<-which(rownames(betas.rrphylo)==des[index])
    ancindex<-which(rownames(betas.rrphylo)==anc[index])
    neglogl<-neglogl+  (betas.rrphylo[desindex]- phi*betas.rrphylo[ancindex])^2/2/ sigsq  
  }
  neglogl<-neglogl+(N-1)*log(2*pi*sigsq)/2
  names(neglogl)<-"neglogl"
  return(neglogl)
}

mle_ar1 <- function(betas.rrphylo=betas.rrphylo, tree=tree) {
  yw <- arima(betas.rrphylo, order = c(1, 0, 0), method = "ML")
  start_params <- c(yw$coef[1], yw$sigma2)
  names(start_params) <- c("phi", "sigsq") # Only phi for AR(1)

  lower_bounds <- c( -Inf, 1e-6) # sigsq should be positive
  opt <- optim(start_params, ar1.negloglike, betas.rrphylo = betas.rrphylo, tree = tree, method = "L-BFGS-B", lower = lower_bounds)
  
  
  list(coefficients = opt$par, loglik = -opt$value, hessian = opt$hessian)
}



ar2.negloglike <- function(param, betas.rrphylo=betas.rrphylo,tree=tree){
  phi<-param["phi"]
  phi2<-param["phi2"]
  sigsq<-param["sigsq"]
  # phi1<-0.01
  # phi2<-0.05
  # sigsq<-0.1
  # betas.rrphylo=sim.rrphylo.betas.est

  tree<-reorder(tree,"postorder")
  tree$root.edge<-0
  N<-dim(tree$edge)[1]
  anc<-tree$edge[,1]
  des<-tree$edge[,2]
  # nodeset<-rownames(betas.rrphylo)[1:(Ntip(tree)-1)]
  # tipset<-rownames(betas.rrphylo)[(Ntip(tree):(2*Ntip(tree)-1))]
  # tipset<-gsub("t", "",tipset)
  # rownames(betas.rrphylo)
  # rownames(betas.rrphylo)<-c(nodeset,tipset)
  # Print the updated data frame or matrix
  #print(betas.rrphylo)

  
  treeedge<-tree$edge
  colnames(treeedge)<-c("anc","des")
  treeedge
  neglogl<-0
  for(index in N:1){
    #index <- N-4
    desindex<-which(rownames(betas.rrphylo)==des[index])
    ancindex1<-which(rownames(betas.rrphylo)==anc[index])
    if(anc[index]==Ntip(tree)+1){#this only allow 1 branch so AR1
    betas.rrphylo
    neglogl<-neglogl+(betas.rrphylo[desindex]- phi*betas.rrphylo[ancindex1])^2/2/sigsq 
    }else{
    
    ancindex2<-which(rownames(betas.rrphylo)==anc[des==anc[index]])
neglogl<-neglogl+(betas.rrphylo[desindex]- phi*betas.rrphylo[ancindex1] - phi2*betas.rrphylo[ancindex2] )^2/2/sigsq
  treeedge
    } 
  }
  neglogl<-neglogl+(N-2)*log(2*pi*sigsq)/2
  names(neglogl)<-"neglogl"
  return(neglogl)
}

mle_ar2 <- function(betas.rrphylo=betas.rrphylo, tree=tree) {
  yw <- arima(betas.rrphylo, order = c(2, 0, 0), method = "ML")
  start_params <- c(yw$coef[1], yw$coef[2], yw$sigma2)
  names(start_params) <- c("phi", "phi2", "sigsq") # phi and phi2 for AR(2)
  
  lower_bounds <- c(-Inf, -Inf, 1e-6) # sigsq should be positive
  opt <- optim(start_params, ar2.negloglike, betas.rrphylo = betas.rrphylo, tree = tree, method = "L-BFGS-B", lower = lower_bounds)
  
  list(coefficients = opt$par, loglik = -opt$value, hessian = opt$hessian)
}


arma11.negloglike <- function(param, betas.rrphylo=betas.rrphylo,tree=tree){
  #betas.rrphylo <-sim.y.betas$betas.rrphylo.sim
#   phi<-0.1
#   theta <- 0.2
#   sigsq<-1
  phi<-param["phi"]
  theta <- param["theta"]
  sigsq<-param["sigsq"]
  #sigsq<-1
  #betas.rrphylo<-betas.rrphylo.sim

  
  tree<-reorder(tree,"postorder")
  tree$root.edge<-0
  N<-dim(tree$edge)[1]
  anc<-tree$edge[,1]
  des<-tree$edge[,2]
  
  neglogl<-0

  eps <- rep(0, N)
  eta <- rep(0, N)
  eps[Ntip(tree)+1] <- betas.rrphylo[Ntip(tree)+1] / (1-phi)  # Approximation
  eta[Ntip(tree)+1] <- eps[Ntip(tree)+1]

  for(index in N:1){
    #N<-1
    desindex<-which(rownames(betas.rrphylo)==des[index])
    ancindex<-which(rownames(betas.rrphylo)==anc[index])
    eta[desindex] <- betas.rrphylo[desindex] - phi * betas.rrphylo[ancindex] - theta * eps[ancindex]
    eps[desindex] <- eta[desindex]

    neglogl<-neglogl+  eta[desindex]^2/2/ sigsq  
  }
  neglogl<-neglogl+(N-1)*log(2*pi*sigsq)/2
  names(neglogl)<-"neglogl"
  return(neglogl)
}


mle_arma11 <- function(params,betas.rrphylo=betas.rrphylo,tree=tree) {
  # Log-likelihood function for ARMA(1, 1)
  # Initial parameter estimates using Yule-Walker equations
  #y<-sim_data
  
  yw <- arima(betas.rrphylo, order = c(1, 0, 1), method = "ML")
  start_params <- c(yw$coef[1], yw$coef[2], yw$sigma2)
names(start_params)<-c("phi","theta","sigsq")
  # Optimization to find MLE

lower_bounds <- c(-Inf, -Inf, 1e-6) # sigsq should be positive
opt <- optim(start_params, arma11.negloglike, betas.rrphylo = betas.rrphylo, tree = tree, method = "L-BFGS-B", lower = lower_bounds)

  # Return the results
  list(coefficients = opt$par, loglik = -opt$value, hessian = opt$hessian)
}


arma21.negloglike <- function(param, betas.rrphylo=betas.rrphylo, tree=tree) {
  #phi2<-0.12
  phi <- param["phi"]
  phi2 <- param["phi2"] # New AR term
  theta <- param["theta"]
  sigsq <- param["sigsq"]

  tree <- reorder(tree, "postorder")
  tree$root.edge <- 0
  N <- dim(tree$edge)[1]
  anc <- tree$edge[,1]
  des <- tree$edge[,2]

  neglogl <- 0

  eps <- rep(0, N)
  eta <- rep(0, N)
  eps[Ntip(tree)+1] <- betas.rrphylo[Ntip(tree)+1] / (1 - phi - phi2) # Adjusted for ARMA(2,1)
  eta[Ntip(tree)+1] <- eps[Ntip(tree)+1]

  for(index in N:1) {
    desindex <- which(rownames(betas.rrphylo) == des[index])
    ancindex <- which(rownames(betas.rrphylo) == anc[index])
    eta[desindex] <- betas.rrphylo[desindex] - phi * betas.rrphylo[ancindex] - phi2 * betas.rrphylo[ancindex] - theta * eps[ancindex] # Adjusted for ARMA(2,1)
    eps[desindex] <- eta[desindex]

    neglogl <- neglogl + eta[desindex]^2 / 2 / sigsq
  }
  neglogl <- neglogl + (N - 1) * log(2 * pi * sigsq) / 2
  names(neglogl) <- "neglogl"
  return(neglogl)
}


mle_arma21 <- function(params, betas.rrphylo=betas.rrphylo, tree=tree) {
  yw <- arima(betas.rrphylo, order = c(2, 0, 1), method = "ML")
  start_params <- c(yw$coef[1], yw$coef[2], yw$coef[3], yw$sigma2)
  names(start_params) <- c("phi", "phi2", "theta", "sigsq") # Adjusted for ARMA(2,1)


  lower_bounds <- c(-Inf, -Inf,-Inf, 1e-6) # sigsq should be positive
  opt <- optim(start_params, arma21.negloglike, betas.rrphylo = betas.rrphylo, tree = tree, method = "L-BFGS-B", lower = lower_bounds)
  
    
  list(coefficients = opt$par, loglik = -opt$value, hessian = opt$hessian)
}


arma22.negloglike <- function(param, betas.rrphylo=betas.rrphylo, tree=tree) {
  #theta2<-0.1
  phi <- param["phi"]
  phi2 <- param["phi2"]
  theta <- param["theta"]
  theta2 <- param["theta2"] # New MA term
  sigsq <- param["sigsq"]

  tree <- reorder(tree, "postorder")
  tree$root.edge <- 0
  N <- dim(tree$edge)[1]
  anc <- tree$edge[,1]
  des <- tree$edge[,2]

  neglogl <- 0

  eps <- rep(0, N)
  eta <- rep(0, N)
  eps[Ntip(tree)+1] <- betas.rrphylo[Ntip(tree)+1] / (1 - phi - phi2) # Adjusted for ARMA(2,2)
  eta[Ntip(tree)+1] <- eps[Ntip(tree)+1]

  for(index in N:1) {
    desindex <- which(rownames(betas.rrphylo) == des[index])
    ancindex <- which(rownames(betas.rrphylo) == anc[index])
    # Adjusted eta calculation for ARMA(2,2)
    eta[desindex] <- betas.rrphylo[desindex] - phi * betas.rrphylo[ancindex] - phi2 * betas.rrphylo[ancindex] - theta * eps[ancindex] - theta2 * eps[ancindex]
    eps[desindex] <- eta[desindex]

    neglogl <- neglogl + eta[desindex]^2 / 2 / sigsq
  }
  neglogl <- neglogl + (N - 1) * log(2 * pi * sigsq) / 2
  names(neglogl) <- "neglogl"
  return(neglogl)
}

mle_arma22 <- function(params, betas.rrphylo=betas.rrphylo, tree=tree) {
  yw <- arima(betas.rrphylo, order = c(2, 0, 2), method = "ML")
  start_params <- c(yw$coef[1], yw$coef[2], yw$coef[3], yw$coef[4], yw$sigma2)
  names(start_params) <- c("phi", "phi2", "theta", "theta2", "sigsq") # Adjusted for ARMA(2,2)

  lower_bounds <- c(-Inf, -Inf,-Inf,-Inf, 1e-6) # sigsq should be positive
  opt <- optim(start_params, arma22.negloglike, betas.rrphylo = betas.rrphylo, tree = tree, method = "L-BFGS-B", lower = lower_bounds)
  
  
  list(coefficients = opt$par, loglik = -opt$value, hessian = opt$hessian)
}

# Function to calculate AIC
calculate_aic <- function(output) {
  k <- length(output$coefficients) # Number of parameters
  loglik <- output$loglik
  aic <- 2 * k - 2 * loglik
  return(aic)
}


# get the mle of the parameter
ar1.mle <- function(betas.rrphylo=betas.rrphylo,tree=tree){
  #tree<-reorder(tree,"postorder")
  tree$root.edge<-0
  N<-dim(tree$edge)[1]
  anc<-tree$edge[,1]
  des<-tree$edge[,2]
  
  phi.mle.num<-0
  phi.mle.den<-0
  
  for(index in 1:N){
    desindex<-which(rownames(betas.rrphylo)==des[index])
    ancindex<-which(rownames(betas.rrphylo)==anc[index])
    phi.mle.num<-phi.mle.num + (betas.rrphylo[desindex]*betas.rrphylo[ancindex])  
    phi.mle.den<-phi.mle.den + (betas.rrphylo[ancindex]^2)
  } 
  phi.mle <- phi.mle.num/phi.mle.den
#  phi.mle
  sigsq.mle<-0
  for(index in 1:N){
    desindex<-which(rownames(betas.rrphylo)==des[index])
    ancindex<-which(rownames(betas.rrphylo)==anc[index])
    sigsq.mle<-sigsq.mle + (betas.rrphylo[desindex]- phi.mle*betas.rrphylo[ancindex])^2
  }
  sigsq.mle<-sigsq.mle/N
  
  return(list(phi=phi.mle,sigsq=sigsq.mle))
}


# get the mle of the parameter
ar2.mle <- function(betas.rrphylo=betas.rrphylo,tree=tree){
  #tree<-reorder(tree,"postorder")
  tree$root.edge<-0
  N<-dim(tree$edge)[1]
  anc<-tree$edge[,1]
  des<-tree$edge[,2]
  
  A<-0;B<-0;C<-0;D<-0;E<-0
  tree$edge
  for(index in N:1){
  #  index<-3
    desindex<-which(rownames(betas.rrphylo)==des[index])
    ancindex1<-which(rownames(betas.rrphylo)==anc[index])

   if(anc[index]==Ntip(tree)+1){#this only allow 1 branch so AR1
   A<-A+betas.rrphylo[ancindex1]^2
   D<-D+(betas.rrphylo[desindex]*betas.rrphylo[ancindex1])  
  }else{
   ancindex2<-which(rownames(betas.rrphylo)==anc[des==anc[index]])
   A<-A+betas.rrphylo[ancindex1]^2
   B<-B+(betas.rrphylo[ancindex1]*betas.rrphylo[ancindex2])  
   C<-C+ betas.rrphylo[ancindex2]^2  
   D<-D+(betas.rrphylo[desindex]*betas.rrphylo[ancindex1])  
   E<-E+(betas.rrphylo[desindex]*betas.rrphylo[ancindex2])  
  } 
  }
 # A<-1;B<-2;C<-3;D<-1;E<-0.2

  phi12.mle <- solve(matrix(c(A,B,B,C),nrow=2))%*%array(c(D,E),c(2,1))
  phi1.mle<-phi12.mle[1]
  phi2.mle<-phi12.mle[2]
#  phi.mle
  sigsq.mle<-0
  for(index in N:1){
    desindex<-which(rownames(betas.rrphylo)==des[index])
    ancindex1<-which(rownames(betas.rrphylo)==anc[index])
     if(anc[index]==Ntip(tree)+1){#this only allow 1 branch so AR1
    sigsq.mle<-sigsq.mle + (betas.rrphylo[desindex]- phi1.mle*betas.rrphylo[ancindex1])^2
  }else{
       ancindex2<-which(rownames(betas.rrphylo)==anc[des==anc[index]])
       sigsq.mle<-sigsq.mle + (betas.rrphylo[desindex]- phi1.mle*betas.rrphylo[ancindex1]-phi2.mle*betas.rrphylo[ancindex2])^2
      } 
  }
  sigsq.mle<-sigsq.mle/(N-2)
  
  return(list(phi1=phi1.mle,phi2=phi2.mle,sigsq=sigsq.mle))
}

test.phi<-function(param, betas.rrphylo=betas.rrphylo,tree=tree){
  phi<-param["phi"]
  sigsq<-param["sigsq"]
  phi.mle.den<-0
  anc<-tree$edge[,1]
  des<-tree$edge[,2]
  N<-dim(tree$edge)[1]
  for(index in 1:N){
    phi.mle.den<-phi.mle.den + (betas.rrphylo[anc[index]]^2)
  } 
  var.phi<-sigsq/phi.mle.den
  se.phi<-sqrt(var.phi)
  z=phi/se.phi
  #compute the p-value for z
  pval<-2*pnorm(-abs(z))
  names(pval)<-"pval"
  return(list(z=z,pval=pval))
}

# split the character and number for "Node_10"
get_subtree_numbers <- function(subtree) {
  node_labels <- subtree$node.label
  split_labels <- strsplit(node_labels, "_")
  number_parts <- sapply(split_labels, function(x) as.numeric(x[2]))
  return(number_parts)
}

getMainCladesFromRoot <- function(tree) {
  # Determine the number of internal nodes
  n_internal_nodes <- Nnode(tree)
  
  # Create labels for the internal nodes
  internal_labels <- paste0("Node_", (Ntip(tree) + 1):(Ntip(tree) + n_internal_nodes))
  
  # Assign the labels to tree$node.label
  tree$node.label <- internal_labels
  
  # Find the edges stemming directly from the root
  edges_from_root <- tree$edge[tree$edge[,1] == Ntip(tree) +1, 2]   
  
  clades <- list()
  for (i in seq_along(edges_from_root)) {
    # i<-2
    # Extract the clade
    clade <- extract.clade(tree, node = edges_from_root[i])
    
    clades[[i]] <- clade
  }
  
  return(clades)
}

modelfit <- function(trait=trait, tree=tree) {
  betas.rrphylo <- compute_betas_rrphylo(tree, y)$betas.rrphylo
  # Run MLE estimation for each model
  output_ar1 <- mle_ar1(betas.rrphylo = betas.rrphylo, tree = tree)
  output_ar2 <- mle_ar2(betas.rrphylo = betas.rrphylo, tree = tree)
  output_arma11 <- mle_arma11(betas.rrphylo = betas.rrphylo, tree = tree)
  output_arma21 <- mle_arma21(betas.rrphylo = betas.rrphylo, tree = tree)
  output_arma22 <- mle_arma22(betas.rrphylo = betas.rrphylo, tree = tree)
  
  # Calculate AIC for each model
  aic_ar1 <- calculate_aic(output_ar1)
  aic_ar2 <- calculate_aic(output_ar2)
  aic_arma11 <- calculate_aic(output_arma11)
  aic_arma21 <- calculate_aic(output_arma21)
  aic_arma22 <- calculate_aic(output_arma22)
  
  # Create a list to store the MLE and AIC results for each model
  results <- list(
    AR1 = list(mle = output_ar1, aic = aic_ar1),
    AR2 = list(mle = output_ar2, aic = aic_ar2),
    ARMA11 = list(mle = output_arma11, aic = aic_arma11),
    ARMA21 = list(mle = output_arma21, aic = aic_arma21),
    ARMA22 = list(mle = output_arma22, aic = aic_arma22)
  )
  
  return(results)
}

tree2subtrees <- function(tree) {
  # Set up the plot layout
  par(mfrow = c(1, 3))
  
  # Plot the raw tree
  plot(tree,  show.tip.label = FALSE , edge.width = 2, cex = 3)
  title(main = "Raw Tree", cex.main = 2)
  nodelabels(tree$node.label, cex = 2)
  tiplabels(tree$tip.label, cex = 2)
  
  # Split the tree from the root into two clades
  clades <- getMainCladesFromRoot(tree)
  
  # Plot each clade
  for (i in seq_along(clades)) {
    subtree<-clades[[i]]
    node_numbers <- gsub("Node_", "", subtree$node.label)
    # Rename node labels
    subtree$node.label <- node_numbers#paste("t", node_numbers, sep = "")
    plot(subtree, show.tip.label = FALSE, edge.width = 2, cex = 3)
    title(main = paste("Clade", i), cex.main = 2)
    nodelabels(subtree$node.label, cex = 2)
    tiplabels(subtree$tip.label, cex = 2)
  }
}
#######################################################################
#######################################################################

# Part III 
#rm(list=ls())
## Let yr to finish the rest of the code  so she can be sure she understand the code to get her own credit.---it is ok tony...so balanced and reconcile it then everyone happy
#sim trait data  

# mfit<-modelfit(trait=y, tree=tree) 
# #AR(1) testing
# testout<-test.phi(param=mfit$AR1$mle$coefficients, betas.rrphylo=betas.rrphylo,tree=tree)
# testout$z
# testout$pval
# #   pval 
# #0.5774715 
# 
# modelset<-c("AR1","AR2","ARMA1","ARMA21","ARMA22")
# loglikset<-c(mfit$AR1$mle$loglik,mfit$AR2$mle$loglik,mfit$ARMA11$mle$loglik,mfit$ARMA21$mle$loglik,
#              mfit$ARMA22$mle$loglik)
# aicset<-c(mfit$AR1$aic,mfit$AR2$aic,mfit$ARMA11$aic,mfit$ARMA21$aic,mfit$ARMA22$aic)
# paranumset<-c(2,3,3,4,5)
# aicweightset<-akaike.weights(aicset)$weights 
# 
# fitset<-data.frame(
# modelset=modelset,
# paranumset=paranumset,
# loglikset=round(loglikset,2),
# aicset=round(aicset,2),
# aicweightset=round(aicweightset,2),
# rankset=round(rank(aicset))
# )
# 
# fitset



# # Get the maximum length of the vectors
# max_length <- max(length(mfit$AR1$mle$coefficients), length(mfit$AR2$mle$coefficients), length(mfit$ARMA11$mle$coefficients), length(mfit$ARMA21$mle$coefficients), length(mfit$ARMA22$mle$coefficients))
# 
# paramnames<-c("phi","phi2","theta","theta2","sigsq")
# # Pad the vectors with NA values to match the maximum length
# AR1 <- c(mfit$AR1$mle$coefficients, rep(NA, max_length - length(mfit$AR1$mle$coefficients)))
# AR2 <- c(mfit$AR2$mle$coefficients, rep(NA, max_length - length(mfit$AR2$mle$coefficients)))
# ARMA11 <- c(mfit$ARMA11$mle$coefficients, rep(NA, max_length - length(mfit$ARMA11$mle$coefficients)))
# ARMA21 <- c(mfit$ARMA21$mle$coefficients, rep(NA, max_length - length(mfit$ARMA21$mle$coefficients)))
# ARMA22 <- c(mfit$ARMA22$mle$coefficients, rep(NA, max_length - length(mfit$ARMA22$mle$coefficients)))
# 
# paramnames<-c("phi","phi2","theta","theta2","sigsq")
# # Combine the padded vectors into a data frame
# parammleset <- rbind(AR1[paramnames], AR2[paramnames], ARMA11[paramnames], ARMA21[paramnames], ARMA22[paramnames])
# colnames(parammleset)<-paramnames
# parammleset
# rownames(parammleset)<-modelset
# parammleset

empmodel <- function(mfit) {
  # Define the model set
  modelset <- c("AR1", "AR2", "ARMA1", "ARMA21", "ARMA22")
  
  # Calculate log likelihoods, AICs, parameter numbers, and AIC weights
  loglikset <- c(mfit$AR1$mle$loglik, mfit$AR2$mle$loglik, mfit$ARMA11$mle$loglik, mfit$ARMA21$mle$loglik, mfit$ARMA22$mle$loglik)
  aicset <- c(mfit$AR1$aic, mfit$AR2$aic, mfit$ARMA11$aic, mfit$ARMA21$aic, mfit$ARMA22$aic)
  paranumset <- c(2, 3, 3, 4, 5)
  aicweightset <- akaike.weights(aicset)$weights
  
  # Create a data frame with the calculated values
  fitset <- data.frame(
    modelset = modelset,
    paranumset = paranumset,
    loglikset = round(loglikset, 2),
    aicset = round(aicset, 2),
    aicweightset = round(aicweightset, 2),
    rankset = round(rank(aicset))
  )
  
  # Get the maximum length of the vectors
  max_length <- max(length(mfit$AR1$mle$coefficients), length(mfit$AR2$mle$coefficients), length(mfit$ARMA11$mle$coefficients), length(mfit$ARMA21$mle$coefficients), length(mfit$ARMA22$mle$coefficients))
  
  # Define parameter names
  paramnames <- c("phi", "phi2", "theta", "theta2", "sigsq")
  
  # Pad the vectors with NA values to match the maximum length
  AR1 <- c(mfit$AR1$mle$coefficients, rep(NA, max_length - length(mfit$AR1$mle$coefficients)))
  AR2 <- c(mfit$AR2$mle$coefficients, rep(NA, max_length - length(mfit$AR2$mle$coefficients)))
  ARMA11 <- c(mfit$ARMA11$mle$coefficients, rep(NA, max_length - length(mfit$ARMA11$mle$coefficients)))
  ARMA21 <- c(mfit$ARMA21$mle$coefficients, rep(NA, max_length - length(mfit$ARMA21$mle$coefficients)))
  ARMA22 <- c(mfit$ARMA22$mle$coefficients, rep(NA, max_length - length(mfit$ARMA22$mle$coefficients)))
  
  # Combine the padded vectors into a data frame
  parammleset <- rbind(AR1[paramnames], AR2[paramnames], ARMA11[paramnames], ARMA21[paramnames], ARMA22[paramnames])
  colnames(parammleset) <- paramnames
  rownames(parammleset) <- modelset
  
  # Calculate weighted parameter estimates
  weighted_param_estimates <- sapply(paramnames, function(param) sum(parammleset[, param] * fitset$aicweightset, na.rm = TRUE))
  
  # Add the weighted parameter estimates to the data frame
  parammleset <- rbind(parammleset, weighted_param_estimates)
  rownames(parammleset)[nrow(parammleset)] <- "Weighted Estimate"
  
  # Return the data frames
  list(fitset = fitset, parammleset = parammleset)
}


testrate2subtrees <- function(tree, y) {
  # Split the tree into subtrees, compute statistics, perform likelihood ratio test...
  clades <- getMainCladesFromRoot(tree)
  
  ## get the first tree
  subtree1<-clades[[1]]
  subtree1$node.label
  # Extract the numbers from the original node labels
  node_numbers <- gsub("Node_", "", subtree1$node.label)
  # Rename node labels
  subtree1$node.label <- paste("t", node_numbers, sep = "")
  # Print the new node labels
  print(subtree1$node.label)
  subtree1$node.label<-subtree1$node.label
  subtree1$node.label
  subtree1$tip.label
  subtree1$edge
  y1<-y[names(y) %in% subtree1$tip.label]
  y1
  ## get the rate branch estimates from RRphylo
  betas.rrphylo.1 <- compute_betas_rrphylo(subtree1,y1)$betas.rrphylo
  
  #rownames(betas.rrphylo.2)[1:(Ntip(subtree1)-1)]<-get_subtree_numbers(subtree2)
  names(betas.rrphylo.1)<-c( (length(subtree1$tip.label)+1): (length(subtree1$tip.label)+  length((subtree1$node.label))), 1:length(subtree1$tip.label))
  
  #names(betas.rrphylo.1)<-c(subtree1$node.label,                          subtree1$tip.label)
  betas.rrphylo.1
  
  # Assuming your data is in a vector named 'data'
  names(betas.rrphylo.1) <- gsub("t", "", names(betas.rrphylo.1))
  
  # Print the new names
  print(names(betas.rrphylo.1))
  
  
  ################################
  ###### get the second tree #####
  ################################
  subtree2<-clades[[2]]
  y2<-y[names(y) %in% subtree2$tip.label]
  
  subtree2$tip.label
  subtree2$node.label
  # Extract the numbers from the original node labels
  node_numbers <- gsub("Node_", "", subtree2$node.label)
  # Rename node labels
  subtree2$node.label <- paste("t", node_numbers, sep = "")
  subtree2$node.label
  
  # Print the new node labels
  print(subtree2$node.label)
  y2<-y[names(y) %in% subtree2$tip.label]
  y2
  ## get the rate branch estimates from RRphylo
  betas.rrphylo.2 <- compute_betas_rrphylo(subtree2,y2)$betas.rrphylo
  betas.rrphylo.2
  #rownames(betas.rrphylo.2)[1:(Ntip(subtree1)-1)]<-get_subtree_numbers(subtree2)
  
  names(betas.rrphylo.2)<-c( (length(subtree2$tip.label)+1): (length(subtree2$tip.label)+  length((subtree2$node.label))), 1:length(subtree2$tip.label))
  betas.rrphylo.2
  print(names(betas.rrphylo.2))
  
  ##################
  ##################
  param.mle <- unlist(ar1.mle(betas.rrphylo = betas.rrphylo, tree = tree))
  names(param.mle) <- c("phi", "sigsq")
  param.mle
  ar1.negloglike(param.mle, betas.rrphylo=betas.rrphylo, tree=tree)
  
  
  ####
  
  mle_ar1(betas.rrphylo=betas.rrphylo.1, tree=subtree1)
  ar1.mle(betas.rrphylo=betas.rrphylo.1,tree=subtree1)
  param.mle.1<-unlist(ar1.mle(betas.rrphylo=betas.rrphylo.1,tree=subtree1))
  names(param.mle.1)<-c("phi","sigsq")
  param.mle.1
  ar1.negloglike(param.mle.1, betas.rrphylo=betas.rrphylo.1, tree=subtree1)
  
  rownames(betas.rrphylo.2)<-c( (length(subtree2$tip.label)+1):(length(subtree2$tip.label)+length(subtree2$node.label)), 1:length(subtree2$tip.label))
  betas.rrphylo.2
  param.mle.2<-unlist(ar1.mle(betas.rrphylo=betas.rrphylo.2,tree=subtree2))
  names(param.mle.2)<-c("phi","sigsq")
  param.mle.2
  ar1.negloglike(param.mle.2, betas.rrphylo=betas.rrphylo.2, tree=subtree1)
  
  
  length(betas.rrphylo)
  length(betas.rrphylo.1)
  length(betas.rrphylo.2)
  
  
  ## need to wrap above into a function asking user to enter tree, and y the it would return the mle, and the likelihood
  # do the testing homogeity of the two parameter H0: phi1-phi2 = 0
  # Compute the negative log-likelihoods
  LL_global <- ar1.negloglike(param.mle, betas.rrphylo=betas.rrphylo, tree=tree)
  LL_sub1 <- ar1.negloglike(param.mle.1, betas.rrphylo=betas.rrphylo.1, tree=subtree1)
  LL_sub2 <- ar1.negloglike(param.mle.2, betas.rrphylo=betas.rrphylo.2, tree=subtree2)
  LL_global
  LL_sub1
  LL_sub2
  
  # Compute the test statistic
  test_statistic <- 2 * ((LL_sub1 + LL_sub2)-LL_global)
  names(test_statistic)<-c("test.stat")
  test_statistic
  
  # Perform the chi-square test
  p_value <- 1 - pchisq(test_statistic, df=2) # two AR(1) - one AR(1)
  names(p_value)<-c("p_value")
  p_value # zero seems make sense as the rate is from the y simulated from 
  
  # Return a list containing the results
  return(list(
    subtrees = list(tree, subtree1, subtree2),
    trait_values = list(y, y1, y2),
    mles = list(param.mle, param.mle.1, param.mle.2),
    likelihoods = list(LL_global, LL_sub1, LL_sub2),
    test_statistic = test_statistic,
    p_value = p_value
  ))
}

treeARanova <- function(testrate2subtrees_output) {
  treecase <- c("treeAR1", "subtree1AR1","subtree2AR1")
  
  mles <- testrate2subtrees_output$mles[[1]]
  mles.1 <- testrate2subtrees_output$mles[[2]]
  mles.2 <- testrate2subtrees_output$mles[[3]]
  
  numparam <- c(2,2,2)
  mlesoutphi <- c(mles['phi'], mles.1['phi'], mles.2['phi'])
  mlesoutsigsq <- c(mles['sigsq'], mles.1['sigsq'], mles.2['sigsq'])
  
  likelihoods <- testrate2subtrees_output$likelihoods[[1]]
  likelihoods1 <- testrate2subtrees_output$likelihoods[[2]]
  likelihoods2 <- testrate2subtrees_output$likelihoods[[3]]
  
  likelihoodsout <- c(likelihoods, likelihoods1, likelihoods2)
  
  teststat <- c(testrate2subtrees_output$test_statistic, NA, NA)
  deg <- c(2, NA, NA)
  pval <- c(format(testrate2subtrees_output$p_value, scientific = TRUE), NA, NA)
  
  tworatedf <- data.frame(numparam = numparam, mlesoutphi = round(mlesoutphi,3), mlesoutsigsq = round(mlesoutsigsq,3), likelihoodsout = round(likelihoodsout,3), teststat = round(teststat,3), deg = deg, pval = pval)
  rownames(tworatedf) <- treecase
  
  return(tworatedf)
}



#####################################
##############Data Analysis##########
#####################################
ntaxa<-8
stree<-stree(ntaxa)
tree<-compute.brlen(stree(ntaxa,type="balanced"))
#tree<-reorder(tree,"postorder")
fastBM(tree,internal=T,sig2=1)->phen
phen[1:ntaxa]->y
trait<-round(y,2)
trait
edgelength<-tree$edge.length
names(edgelength)<-edge_labels
edgelength
betas.rrphylo <- compute_betas_rrphylo(tree, y)$betas.rrphylo
betas.rrphylo

# Rename node labels
node_labels <- paste("nd", (ntaxa+1):(2*ntaxa-1),sep="")
plot(tree,  show.tip.label = FALSE , edge.width = 2, cex = 3)
title(main = "Raw Tree", cex.main = 1.5)
nodelabels(node_labels, cex = 1.5)
# Add tip labels
tiplabels(paste("y",1:ntaxa,sep=""), cex = 1.5)
# Get the number of edges
num_edges <- Nedge(tree)
# Create edge labels
edge_labels <- paste0("eg", seq_len(num_edges),sep="")
edgelabels(edge_labels, cex = 1.5)

df.trait<-data.frame(trait=t(trait))
colnames(df.trait)<-paste("y",1:length(trait),sep="")
df.trait
xtable(df.trait)

df.edgelen<-data.frame(edgelength=round(t(edgelength),3))
colnames(df.edgelen)<-paste("eg",1:length(df.edgelen),sep="")
df.edgelen
xtable(df.edgelen,digits=3)

df.betas.rrphylo<-data.frame(betas.rrphylo=t(round(betas.rrphylo,3)))
colnames(df.betas.rrphylo)<-paste("nd",1:length(df.betas.rrphylo),sep="")
df.betas.rrphylo
xtable(df.betas.rrphylo,digits=3)




#######################################
############ Formal Analysis ##########
#######################################

mfit<-modelfit(trait=y, tree=tree) 
empmodelout<-empmodel(mfit)

empmodelout$parammleset
empmodelout$fitset

xtable(empmodelout$parammleset,digits=3)
xtable(empmodelout$fitset,digits=3)

testrate2subtrees_output <- testrate2subtrees(tree, y) 
tworatedf<- treeARanova(testrate2subtrees_output)
tworatedf
xtable(tworatedf,digits=3)

#make the table for the models report the mle, loglik and aic
# ## shall get the rate first then do postorder
# betas.rrphylo <- compute_betas_rrphylo(tree, y)$betas.rrphylo
# betas.rrphylo
# 
# length(betas.rrphylo)
# 
# mle_ar1(betas.rrphylo=betas.rrphylo, tree=tree)
# mle_ar2(betas.rrphylo=betas.rrphylo, tree=tree)
# mle_arma11(betas.rrphylo=betas.rrphylo, tree=tree)
# mle_arma21(betas.rrphylo=betas.rrphylo, tree=tree)
# mle_arma22(betas.rrphylo=betas.rrphylo, tree=tree)
# 
# # Assuming the output from each function is stored in the respective variables
# output_ar1 <- mle_ar1(betas.rrphylo=betas.rrphylo, tree=tree)
# output_ar2 <- mle_ar2(betas.rrphylo=betas.rrphylo, tree=tree)
# output_arma11 <- mle_arma11(betas.rrphylo=betas.rrphylo, tree=tree)
# output_arma21 <- mle_arma21(betas.rrphylo=betas.rrphylo, tree=tree)
# output_arma22 <- mle_arma22(betas.rrphylo=betas.rrphylo, tree=tree)
# 
# # Calculate AIC for each model
# aic_ar1 <- calculate_aic(output_ar1)
# aic_ar2 <- calculate_aic(output_ar2)
# aic_arma11 <- calculate_aic(output_arma11)
# aic_arma21 <- calculate_aic(output_arma21)
# aic_arma22 <- calculate_aic(output_arma22)
# 
# # Print AIC values
# aic_values <- c(aic_ar1, aic_ar2, aic_arma11, aic_arma21, aic_arma22)
# names(aic_values) <- c("AR1", "AR2", "ARMA11", "ARMA21", "ARMA22")
# aic_values

# par(mfrow = c(1, 3))
# plot(tree,main="Raw Tree")
# nodelabels()
# tiplabels()
# 
# ## split the tree from the root to two clades
# clades <- getMainCladesFromRoot(tree)
# clades[[1]]
# clades[[2]]
# 
# for (i in seq_along(clades)) {
#   plot(clades[[i]], main = paste("Clade", i))
#   nodelabels()
#   tiplabels()
# }


#tree2subtrees(tree)
# 
# clades <- getMainCladesFromRoot(tree)
# 
# ## get the first tree
# subtree1<-clades[[1]]
# subtree1$node.label
# # Extract the numbers from the original node labels
# node_numbers <- gsub("Node_", "", subtree1$node.label)
# # Rename node labels
# subtree1$node.label <- paste("t", node_numbers, sep = "")
# # Print the new node labels
# print(subtree1$node.label)
# subtree1$node.label<-subtree1$node.label
# subtree1$node.label
# subtree1$tip.label
# subtree1$edge
# y1<-y[names(y) %in% subtree1$tip.label]
# y1
# ## get the rate branch estimates from RRphylo
# betas.rrphylo.1 <- compute_betas_rrphylo(subtree1,y1)$betas.rrphylo
# 
# #rownames(betas.rrphylo.2)[1:(Ntip(subtree1)-1)]<-get_subtree_numbers(subtree2)
# names(betas.rrphylo.1)<-c( (length(subtree1$tip.label)+1): (length(subtree1$tip.label)+  length((subtree1$node.label))), 1:length(subtree1$tip.label))
# 
# #names(betas.rrphylo.1)<-c(subtree1$node.label,                          subtree1$tip.label)
# betas.rrphylo.1
# 
# # Assuming your data is in a vector named 'data'
# names(betas.rrphylo.1) <- gsub("t", "", names(betas.rrphylo.1))
# 
# # Print the new names
# print(names(betas.rrphylo.1))
# 
# 
# ################################
# ###### get the second tree #####
# ################################
# subtree2<-clades[[2]]
# y2<-y[names(y) %in% subtree2$tip.label]
# 
# subtree2$tip.label
# subtree2$node.label
# # Extract the numbers from the original node labels
# node_numbers <- gsub("Node_", "", subtree2$node.label)
# # Rename node labels
# subtree2$node.label <- paste("t", node_numbers, sep = "")
# subtree2$node.label
# 
# # Print the new node labels
# print(subtree2$node.label)
# y2<-y[names(y) %in% subtree2$tip.label]
# y2
# ## get the rate branch estimates from RRphylo
# betas.rrphylo.2 <- compute_betas_rrphylo(subtree2,y2)$betas.rrphylo
# betas.rrphylo.2
# #rownames(betas.rrphylo.2)[1:(Ntip(subtree1)-1)]<-get_subtree_numbers(subtree2)
# 
# names(betas.rrphylo.2)<-c( (length(subtree2$tip.label)+1): (length(subtree2$tip.label)+  length((subtree2$node.label))), 1:length(subtree2$tip.label))
# betas.rrphylo.2
# print(names(betas.rrphylo.2))
# 
# ##################
# ##################
# param.mle <- unlist(ar1.mle(betas.rrphylo = betas.rrphylo, tree = tree))
# names(param.mle) <- c("phi", "sigsq")
# param.mle
# ar1.negloglike(param.mle, betas.rrphylo=betas.rrphylo, tree=tree)
# 
# 
# ####
# 
# mle_ar1(betas.rrphylo=betas.rrphylo.1, tree=subtree1)
# ar1.mle(betas.rrphylo=betas.rrphylo.1,tree=subtree1)
# param.mle.1<-unlist(ar1.mle(betas.rrphylo=betas.rrphylo.1,tree=subtree1))
# names(param.mle.1)<-c("phi","sigsq")
# param.mle.1
# ar1.negloglike(param.mle.1, betas.rrphylo=betas.rrphylo.1, tree=subtree1)
# 
# rownames(betas.rrphylo.2)<-c( (length(subtree2$tip.label)+1):(length(subtree2$tip.label)+length(subtree2$node.label)), 1:length(subtree2$tip.label))
# betas.rrphylo.2
# param.mle.2<-unlist(ar1.mle(betas.rrphylo=betas.rrphylo.2,tree=subtree2))
# names(param.mle.2)<-c("phi","sigsq")
# param.mle.2
# ar1.negloglike(param.mle.2, betas.rrphylo=betas.rrphylo.2, tree=subtree1)
# 
# 
# length(betas.rrphylo)
# length(betas.rrphylo.1)
# length(betas.rrphylo.2)
# 
# 
# ## need to wrap above into a function asking user to enter tree, and y the it would return the mle, and the likelihood
# # do the testing homogeity of the two parameter H0: phi1-phi2 = 0
# # Compute the negative log-likelihoods
# LL_global <- ar1.negloglike(param.mle, betas.rrphylo=betas.rrphylo, tree=tree)
# LL_sub1 <- ar1.negloglike(param.mle.1, betas.rrphylo=betas.rrphylo.1, tree=subtree1)
# LL_sub2 <- ar1.negloglike(param.mle.2, betas.rrphylo=betas.rrphylo.2, tree=subtree2)
# LL_global
# LL_sub1
# LL_sub2
# 
# # Compute the test statistic
# test_statistic <- 2 * ((LL_sub1 + LL_sub2)-LL_global)
# names(test_statistic)<-c("test.stat")
# test_statistic
# 
# # Perform the chi-square test
# p_value <- 1 - pchisq(test_statistic, df=1)
# names(p_value)<-c("p_value")
# p_value # zero seems make sense as the rate is from the y simulated from 

# treecase<-c("treeAR1", "subtree1AR1","subtree2AR1")
# 
# mles<-testrate2subtrees.output$mles[[1]]
# mles.1<-testrate2subtrees.output$mles[[2]]
# mles.2<-testrate2subtrees.output$mles[[3]]
# ##########################################
# 
# numparam<-c(2,2,2)
# mlesoutphi<-c(mles['phi'],mles.1['phi'],mles.2['phi'])
# mlesoutsigsq<-c(mles['sigsq'],mles.1['sigsq'],mles.2['sigsq'])
# 
# likelihoods<-testrate2subtrees.output$likelihoods[[1]]
# likelihoods1<-testrate2subtrees.output$likelihoods[[2]]
# likelihoods2<-testrate2subtrees.output$likelihoods[[3]]
# 
# likelihoodsout<-c(likelihoods,likelihoods1,likelihoods2)
# 
# teststat<-c(testrate2subtrees.output$test_statistic,NA,NA)
# deg<-c(2,NA,NA)
# pval<-c(format(testrate2subtrees.output$p_value, scientific = TRUE),NA,NA)
# 
# tworatedf<-data.frame(numparam=numparam,mlesoutphi=mlesoutphi,mlesoutsigsq=mlesoutsigsq,likelihoodsout=likelihoodsout,teststat=teststat,deg=deg,pval=pval)
# rownames(tworatedf)<-treecase
# tworatedf
# things to do 
# how to simulate collreated AR(1) and od test
# Can merge two trees for two rates model sims and then test for the power it is fun to do…
# how to do two phi cases. 

# can discuss simulation scheme with yr, so she can do it by herself
# Type I error: reject the one phi case when it is true