library("makemyprior")

#### Code to produce graphs in Section 3 ####

formula <- y ~ -1 + mc(a)
data <- list(
  a = rep(1:10, 10),
  y = rep(0, 100)
)

prior1 <- make_prior(
  formula, data, prior = list(tree = "(a); (eps)"))
prior2 <- make_prior(
  formula, data,
  prior = list(tree = "s1 = (a, eps)",
               w = list(s1 = list(prior = "pc0", param = 0.25))))
prior3 <- make_prior(
  formula, data,
  prior = list(tree = "s1 = (a, eps)",
               w = list(s1 = list(prior = "pc1", param = 0.75))))
prior4 <- make_prior(
  formula, data,
  prior = list(tree = "s1 = (a, eps)",
               w = list(s1 = list(prior = "pcM", param = c(0.25, 0.85)))))

plot_marginal_prior(seq(0, 5, 0.050), prior1, "sigma^2[a]", sd = T) # Figure 2a
plot_marginal_prior(seq(0, 1, 0.001), prior2, "w[a/a_eps]") # Figure 2b
plot_marginal_prior(seq(0, 1, 0.001), prior3, "w[a/a_eps]") # Figure 2c
plot_marginal_prior(seq(0, 1, 0.001), prior4, "w[a/a_eps]") # Figure 2d

#### Example model for Section 5 ####

formula <- y ~ x + mc(a) + mc(b)

p <- 10
m <- 10
n <- m*p

set.seed(1)
data <- list(a = rep(1:p, each = m),
             b = rep(1:m, times = p),
             x = runif(n))
data$y <- data$x + rnorm(p, 0, 0.5)[data$a] +
  rnorm(m, 0, 0.3)[data$b] + rnorm(n, 0, 1)

prior <- make_prior(formula, data, family = "gaussian",
                    intercept_prior = c(0, 1000),
                    covariate_prior = list(x = c(0, 100)))

# Do not open GUI unless interactive mode
if (interactive()){
  new_prior <- makemyprior_gui(prior)
  summary(new_prior)
}

prior <- make_prior(
  formula, data,
  prior = list(
    tree = "s1 = (a, b); s2 = (s1, eps)",
    w = list(s1 = list(prior = "pcM", param = c(0.7, 0.5)),
             s2 = list(prior = "pc0", param = 0.25)),
    V = list(s2 = list(prior = "pc0", param = c(3, 0.05)))
  ),
  covariate_prior = list(x = c(0, 100)))

prior

#### Genomic model for wheat breeding from Section 6.1 ####

formula <- y ~
  mc(a, model = "generic0", Cmatrix = Q_a, constr = TRUE) +
  mc(d, model = "generic0", Cmatrix = Q_d, constr = TRUE) +
  mc(x, model = "generic0", Cmatrix = Q_x, constr = TRUE)

wheat_data_scaled <- wheat_data
wheat_data_scaled$Q_a <- scale_precmat(wheat_data$Q_a)
wheat_data_scaled$Q_d <- scale_precmat(wheat_data$Q_d)
wheat_data_scaled$Q_x <- scale_precmat(wheat_data$Q_x)

prior <- make_prior(formula, wheat_data_scaled, prior = list(
  tree = "s1 = (d, x); s2 = (a, s1); s3 = (s2, eps)",
  w = list(s1 = list(prior = "pcM", param = c(0.67, 0.8)),
           s2 = list(prior = "pcM", param = c(0.85, 0.8)),
           s3 = list(prior = "pc0", param = 0.25))))

# the first time you do inference with Stan, we recommend to run:
# compile_stan(save = T)

posterior <- inference_stan(prior, iter = 15000, warmup = 5000,
                            seed = 1, init = "0", chains = 1)

plot_posterior_stan(posterior, param = "prior", prior = TRUE) # Figure 6

#### Latin square experiment from Section 6.1 ####

formula <- y ~ lin + mc(row) + mc(col) + mc(iid) +
  mc(rw2, model = "rw2", constr = T, lin_constr = TRUE)

prior <- make_prior(
  formula, latin_data,
  prior = list(tree = "s1 = (rw2, iid);
                               s2 = (row, col, s1); s3 = (s2, eps)",
               w = list(s1 = list(prior = "pc0", param = 0.25),
                        s2 = list(prior = "dirichlet"),
                        s3 = list(prior = "pc0", param = 0.25))))

posterior <- inference_stan(prior, iter = 15000, warmup = 5000,
                            seed = 1, init = "0", chains = 1,
                            control = list(adapt_delta = 0.9))

plot_posterior_stan(posterior, param = "prior", prior = TRUE) # Figure 7

#### Neonatal mortality from Section 6.2 ####

set.seed(1)
find_pc_prior_param(lower = 0.1, upper = 10, prob = 0.9, N = 2e5)

graph_path <- paste0(path.package("makemyprior"), "/neonatal.graph")

formula <- y ~ urban + mc(nu) + mc(v) +
  mc(u, model = "besag", graph = graph_path, scale.model = TRUE)

prior <- make_prior(
  formula, neonatal_data, family = "binomial",
  prior = list(tree = "s1 = (u, v); s2 = (s1, nu)",
               w = list(s1 = list(prior = "pc0", param = 0.25),
                        s2 = list(prior = "pc1", param = 0.75)),
               V = list(s2 = list(prior = "pc",
                                  param = c(3.35, 0.05)))))

posterior <- inference_stan(prior, iter = 15000, warmup = 5000,
                            seed = 1, init = "0", chains = 1,
                            control = list(adapt_delta = 0.95))

plot_posterior_fixed(posterior) # Figure 8a
plot_posterior_stan(posterior, param = "prior", prior = TRUE) # Figure 8b
posterior


u <- extract_posterior_effect(posterior, "u")


prior_samps <- inference_stan(prior, use_likelihood = F, print_prior = F,
                              iter = 15000, warmup = 5000,
                              seed = 1, init = "0", chains = 1)

plot_several_posterior_stan(
  list(Prior = prior_samps, Posterior = posterior), "stdev") # Figure 9