###########################
## Load required package ##
###########################

library(BCDAG)


###############
## Section 3 ##
###############

## Sample from DAG-Wishart distribution

q <- 4
DAG <- matrix(c(0,1,1,0,0,0,0,1,0,0,0,1,0,0,0,0), nrow = q)
DAG

set.seed(123)
outDL <- rDAGWishart(n = 1, DAG = DAG, a = q, U = diag(1, q))

options("scipen" = 100, digits = 4)

outDL$D
outDL$L


###############
## Section 4 ##
###############

## Compute causal effects from DAG-model parameters

causaleffect(targets = c(3,4), response = 1, L = outDL$L, D = outDL$D)


###############
## Section 5 ##
###############

## Comparison between exact and approximate proposal distribution

## For FIGURE 2 (q = 10)

N_sim = 5000

ratio_O_q10 <- rep(NA, N_sim)

dag         <- matrix(0, 10, 10)
O_dag       <- get_neighboringDAGs(dag)
O_dag_size  <- dim(O_dag)[3]

set.seed(123)

for(s in 1:N_sim){
  
  ## find all DAG neighbors and compute the size of the O set
  
  dag_star        <- O_dag[,,sample(1:O_dag_size, 1)]
  O_dag_star      <- get_neighboringDAGs(dag_star)
  O_dag_star_size <- dim(O_dag_star)[3]
  
  ## compute the ratio
  
  ratio_O_q10[s] <- O_dag_size/O_dag_star_size
  
  ## update the dag
  
  dag        <- dag_star
  O_dag      <- O_dag_star
  O_dag_size <- O_dag_star_size
  
}

## For FIGURE 2 (q = 20)

ratio_O_q20 <- rep(NA, N_sim)

dag         <- matrix(0, 20, 20)
O_dag       <- get_neighboringDAGs(dag)
O_dag_size  <- dim(O_dag)[3]

set.seed(123)

for(s in 1:N_sim){
  
  ## find all DAG neighbors and compute the size of the O set
  
  dag_star        <- O_dag[,,sample(1:O_dag_size, 1)]
  O_dag_star      <- get_neighboringDAGs(dag_star)
  O_dag_star_size <- dim(O_dag_star)[3]
  
  ## compute the ratio
  
  ratio_O_q20[s] <- O_dag_size/O_dag_star_size
  
  ## Update the dag
  
  dag        <- dag_star
  O_dag      <- O_dag_star
  O_dag_size <- O_dag_star_size
  
}

## For FIGURE 2 (q = 40)

ratio_O_q40 <- rep(NA, N_sim)

dag         <- matrix(0, 40, 40)
O_dag       <- get_neighboringDAGs(dag)
O_dag_size  <- dim(O_dag)[3]

set.seed(123)

for(s in 1:N_sim){
  
  ## find all DAG neighbors and compute the size of the O set
  
  dag_star        <- O_dag[,,sample(1:O_dag_size, 1)]
  O_dag_star      <- get_neighboringDAGs(dag_star)
  O_dag_star_size <- dim(O_dag_star)[3]
  
  ## compute the ratio
  
  ratio_O_q40[s] <- O_dag_size/O_dag_star_size
  
  ## Update the dag
  
  dag        <- dag_star
  O_dag      <- O_dag_star
  O_dag_size <- O_dag_star_size
  
}

df_ratio <- data.frame(ratio = c(ratio_O_q10, ratio_O_q20, ratio_O_q40),
                    q = c(rep("10", N_sim), rep("20", N_sim), rep("40", N_sim)),
                    t = c(1:N_sim, 1:N_sim, 1:N_sim))

## FIGURE 2

library(ggplot2)

ggplot(df_ratio[df_ratio$q == 10,], aes(x = t, y = ratio)) +
  geom_point(size = 0.9, shape = 16, col = "gray40") +
  labs(x = "t", y = "ratio  (q = 10)") +
  theme_bw() +
  theme(plot.margin = unit(c(0.5,0.5,0.4,0.4), "cm")) +
  theme(axis.title.y = element_text(margin = margin(
                                             t = 0, r = 10, b = 0, l = 0))) +
  theme(axis.title.x = element_text(margin = margin(
                                             t = 5, r = 0, b = 0, l = 0))) +
  ylim(0.74,1.26) +
  theme(legend.position = "none")

ggplot(df_ratio[df_ratio$q == 20,], aes(x = t, y = ratio)) +
  geom_point(size = 0.9, shape = 16, col = "gray50") +
  labs(x = "t", y = "ratio  (q = 20)") +
  theme_bw() +
  theme(plot.margin = unit(c(0.5,0.5,0.4,0.4), "cm")) +
  theme(axis.title.y = element_text(margin = margin(
                                             t = 0, r = 10, b = 0, l = 0))) +
  theme(axis.title.x = element_text(margin = margin(
                                             t = 5, r = 0, b = 0, l = 0))) +
  ylim(0.74,1.26) +
  theme(legend.position = "none")

ggplot(df_ratio[df_ratio$q == 40,], aes(x = t, y = ratio)) +
  geom_point(size = 0.9, shape = 16, col = "gray50") +
  labs(x = "t", y = "ratio  (q = 40)") +
  theme_bw() +
  theme(plot.margin = unit(c(0.5,0.5,0.4,0.4), "cm")) +
  theme(axis.title.y = element_text(margin = margin(
                                             t = 0, r = 10, b = 0, l = 0))) +
  theme(axis.title.x = element_text(margin = margin(
                                             t = 5, r = 0, b = 0, l = 0))) +
  ylim(0.74,1.26) +
  theme(legend.position = "none")


###############
## Section 6 ##
###############

## Apply the MCMC to simulated data

q <- 8
set.seed(123)
DAG <- rDAG(q = q, w = 0.2)
DAG

## FIGURE 3

#to install Rgraphviz (R session 4.5), please uncomment the code:
#if (!require("BiocManager", quietly = TRUE))
#  install.packages("BiocManager")
#BiocManager::install("Rgraphviz")

graphics::par(pty = "s")
Rgraphviz::plot(as_graphNEL(DAG), main = "")

L <- matrix(runif(n = q*q, min = 0.1, max = 1), q, q)*DAG
diag(L) <- 1
D <- diag(1, q)

Omega <- L%*%solve(D)%*%t(L)

set.seed(123)                                       
X <- mvtnorm::rmvnorm(n = 200, sigma = solve(Omega))

out_mcmc <- learn_DAG(S = 5000, burn = 1000, data = X, a = q, U = diag(1,q),
                      w = 0.1, fast = TRUE, save.memory = FALSE, 
                      collapse = FALSE, verbose = FALSE)

## Recover posterior probabilities of edge inclusion, MAP and MPM DAG estimates

get_edgeprobs(learnDAG_output = out_mcmc)
get_MAPdag(learnDAG_output = out_mcmc)
get_MPMdag(learnDAG_output = out_mcmc)

## Derive the posterior distribution of selected causal effect-coefficients 
## given the MCMC output

joint_causal <- get_causaleffect(learnDAG_output = out_mcmc,
                                 targets = c(5,6,7), response = 1)

head(joint_causal$causaleffects)

## FIGURE 4

plot(joint_causal)

joint_causal$post_mean
joint_causal$post_ci
joint_causal$Probs

confint(object = joint_causal, parm = c(5,6,7), level = 0.95)


###############
## Section 6 ##
###############

## Run MCMC on leukemia data and perform causal effect estimation

data(leukemia)
names(leukemia)

X <- leukemia
q <- ncol(X); n <- nrow(X)

set.seed(123)
out_mcmc <- learn_DAG(S = 60000, burn = 5000, data = X, a = q,
                      U = diag(1,q)/n, w = 0.5, fast = TRUE, 
                      collapse = FALSE, verbose = FALSE)    ###


## FIGURE 5 and 6

get_diagnostics(learnDAG_output = out_mcmc, nodes = 13:18, ask = FALSE)

post_edge_probs <- get_edgeprobs(learnDAG_output = out_mcmc)

MPM_dag <- get_MPMdag(out_mcmc)

## FIGURE 7, 8 and 9

plot(out_mcmc, ask = FALSE)

causal_all <- sapply(1:q, function(j) get_causaleffect(
                          out_mcmc, targets = j, response = 1)$causaleffects)


## FIGURE 10

colnames(causal_all) <- 1:ncol(X)
par(mar = c(6, 5, 3, 3))
boxplot(causal_all[,-1], ylab = "causal effect", 
        outline = F, cex.lab = 1, cex.axis = 1, las = 2)

causal_out_12_13 <- get_causaleffect(learnDAG_output = out_mcmc, 
                                     targets = c(12,13), response = 1)

causal_out_12_13$post_mean
causal_out_12_13$post_ci
causal_out_12_13$Probs

confint(object = causal_out_12_13, level = 0.95)


sessionInfo()