#----------Install R packages (if not already installed)----------#
## from CRAN
#  install.packages("hibayes")
## or from GitHub
#  devtools::install_github("YinLiLin/hibayes")

library("hibayes")

#----------run individual level Bayesian regression model----------#
pheno_file_path <- system.file("extdata", "demo.phe", package = "hibayes")
pheno <- read.table(pheno_file_path, header = TRUE)
dim(pheno)
head(pheno, 4)

bfile_path <- system.file("extdata", "demo", package = "hibayes")
bin <- read_plink(bfile = bfile_path, mode = "A", threads = 4)
fam <- bin[["fam"]]
geno <- bin[["geno"]]
map <- bin[["map"]]

## or redirect the memory-mapping files into a new path:
# bin <- read_plink(bfile = bfile_path, out = "./demo")

## fast load the memory-mapping files into memory at the next time:
# geno <- attach.big.matrix("./demo.desc")

dim(geno)
geno[1:4, 1:5]

geno.id <- fam[, 2]
head(geno.id)

head(map, 4)

# model fitting
fitCpi <- ibrm(T1 ~ season + bwt + (1 | loc) + (1 | dam), data = pheno, 
               M = geno, M.id = geno.id, method = "BayesCpi", printfreq = 100,  
               Pi = c(0.98, 0.02), niter = 50000, nburn = 40000, thin = 1,
               seed = 666666, map = map, windsize = 1e6, verbose = TRUE)

sumfit <- summary(fitCpi)
sumfit
names(sumfit)

str(sumfit)

# fixed effects, variances of environmental and genetic random effects
sumfit$beta
sumfit$VER
sumfit$VGR

head(sumfit$alpha, 4)

# plot some main parameters
par(mfrow = c(2, 2), mar = c(4, 4, 1, 1))
plot(pheno$T1, fitCpi$g$gebv[match(pheno$id, fitCpi$g$id)], mgp = c(2.5, 1, 0),
     xlab = "A. Phenotype", ylab = "GEBV")
plot(density(fitCpi$alpha, adjust = 5), lwd = 2, mgp = c(2.5, 1, 0), 
     xlim = c(-1, 1), xlab = "B. Estimated marker effects", main = "")
plot(fitCpi$MCMCsamples$Vg[1, ], type = "l", mgp = c(2.5, 1, 0),
     xlab = "C. MCMC iteration after burn in", ylab = "Genetic variance")
abline(h = fitCpi$Vg[1], col = "red", lwd = 2)
plot(fitCpi$MCMCsamples$Ve[1, ], type = "l", mgp = c(2.5, 1, 0),
     xlab = "D. MCMC iteration after burn in", ylab = "Residual variance")
abline(h = fitCpi$Ve[1], col = "red")

# get the estimated marker effects and the genomic breeding values:
SNPeffect <- fitCpi[["alpha"]]
gebv <- fitCpi[["g"]]

# visualize the marker effects
CMplot(cbind(map[, 1:3], SNPeffect), type = "h", plot.type = "m",
       LOG10 = FALSE, ylab = "SNP effect", file.output = FALSE)

# fit different methods and visualize the density distribution of marker effects
fitRR <- ibrm(T1 ~ season + bwt + (1 | loc) + (1 | dam), data = pheno, 
              M = geno, M.id = geno.id, method = "BayesRR", thin = 1, 
              niter = 50000, nburn = 40000, verbose = FALSE)
fitA <- ibrm(T1 ~ season + bwt + (1 | loc) + (1 | dam), data = pheno, 
             M = geno, M.id = geno.id, method = "BayesA", thin = 1, 
             niter = 50000, nburn = 40000, verbose = FALSE)
fitB <- ibrm(T1 ~ season + bwt + (1 | loc) + (1 | dam), data = pheno, 
             M = geno, M.id = geno.id, method = "BayesB", thin = 1, 
             niter = 50000, nburn = 40000, verbose = FALSE)
fitR <- ibrm(T1 ~ season + bwt + (1 | loc) + (1 | dam), data = pheno, 
             M = geno, M.id = geno.id, method = "BayesR", thin = 1, 
             niter = 50000, nburn = 40000, verbose = FALSE)

plot(density(fitB[["alpha"]], adjust = 5), xlim = c(-0.5, 0.5), lwd = 2,
     col = "#f9c00c", xlab = "SNP effect", main = "")
lines(density(fitCpi[["alpha"]], adjust = 5), col = "#00b9f1", lwd = 2)
lines(density(fitA[["alpha"]], adjust = 5), col = "#7200da", lwd = 2)
lines(density(fitR[["alpha"]], adjust = 5), col = "#f9320c", lwd = 2)
lines(density(fitRR[["alpha"]], adjust = 5), col = "#75D701", lwd = 2)
legend("topleft",
       legend = c("BayesRR", "BayesA", "BayesB", "BayesCpi", "BayesR"),
       col = c("#75D701", "#7200da", "#f9c00c", "#00b9f1", "#f9320c"),
       lty = 1, lwd = 2)

# GWAS results and its visualization:
gwas <- fitCpi[["gwas"]]
head(gwas, 4)
highlight <- gwas[(1 - gwas[, "WPPA"]) < 0.5, 1]
CMplot(data.frame(gwas[, c(1, 2, 4)], wppa = 1 - gwas[, "WPPA"]), type = "h", 
       plot.type = "m", LOG10 = TRUE, threshold = 0.5, ylim = c(0, 0.6),
       ylab = expression(-log[10](1 - italic(WPPA))), highlight = highlight,
       highlight.col = NULL, highlight.text = highlight, file.output = FALSE)

index5 <- map[, 2] == 5
chr5 <- cbind(map[index5, 1:3], pip = (1 - fitCpi[["pip"]])[index5])
highlight <- chr5[chr5[, 4] < 0.5, 1]
CMplot(chr5, plot.type = "m", width = 9, height = 5, threshold = 0.1,
       ylab = expression(-log[10](1 - italic(PIP))), LOG10 = TRUE,
       ylim = c(0, 0.7), amplify = FALSE, highlight = highlight,
       highlight.col = NULL, highlight.text = highlight, file.output = FALSE)

#----------run summary level Bayesian regression model----------#
sumstat_path <- system.file("extdata", "demo.ma", package = "hibayes")
sumstat <- read.table(sumstat_path, header = TRUE)
head(sumstat, 4)

bfile_path <- system.file("extdata", "demo", package = "hibayes")
bin <- read_plink(bfile_path)
geno <- bin[["geno"]]
map <- bin[["map"]]

# compute LD variance and covariance matrix between markers
ldm1 <- ldmat(geno, threads = 4)
ldm2 <- ldmat(geno, chisq = 5, threads = 4)
ldm3 <- ldmat(geno, map, ldchr = FALSE, threads = 4)
ldm4 <- ldmat(geno, map, ldchr = FALSE, chisq = 5, threads = 4)

p1 <- image(as(ldm1, "dgCMatrix"), colorkey = FALSE, xlab = "", ylab = "",
            sub = "ldm1", cex.sub = 2, col = "red")
p2 <- image(ldm2, colorkey = FALSE, xlab = "", ylab = "", sub = "ldm2",
            cex.sub = 2, col = "red")
p3 <- image(ldm3, colorkey = FALSE, xlab = "", ylab = "", sub = "ldm3",
            cex.sub = 2, col = "red")
p4 <- image(ldm4, colorkey = FALSE, xlab = "", ylab = "", sub = "ldm4",
            cex.sub = 2, col = "red")
library("gridExtra")
plot(grid.arrange(p1, p2, p3, p4, ncol = 2))

sumstat <- sumstat[match(map[, 1], sumstat[, 1]), ]

# model fitting
fitCpi <- sbrm(sumstat = sumstat, ldm = ldm1, method = "BayesCpi",
               Pi = c(0.95, 0.05), niter = 20000, nburn = 12000, seed = 666666,
               map = map, windsize = 1e6)

names(summary(fitCpi))
summary(fitCpi)

#----------run single-step Bayesian regression model----------#
pheno_file_path <- system.file("extdata", "demo.phe", package = "hibayes")
pheno <- read.table(pheno_file_path, header = TRUE)

bfile_path <- system.file("extdata", "demo", package = "hibayes")
bin <- read_plink(bfile = bfile_path, mode = "A", threads = 4)
fam <- bin[["fam"]]
geno.id <- fam[, 2]
geno <- bin[["geno"]]
map <- bin[["map"]]

ped_file_path <- system.file("extdata", "demo.ped", package = "hibayes")
ped <- read.table(ped_file_path, header = TRUE)
dim(ped)
head(ped, 4)

fitR <- ssbrm(T1 ~ sex + bwt + (1 | dam), data = pheno, M = geno,
              M.id = geno.id, pedigree = ped, method = "BayesR", niter = 20000,
              nburn = 12000, printfreq = 100, Pi = c(0.95, 0.02, 0.02, 0.01),
              fold = c(0, 0.0001, 0.001, 0.01), seed = 666666, map = map, windsize = 1e6)

summary(fitR)