# An illustration of the mr_mash_simple implementation applied to a
# small, simulated data set.
suppressMessages(library(MBSP))
# source("../../R/misc.R")
# source("../../R/mr_mash_simple.R")

# SCRIPT PARAMETERS
# -----------------
# Number of samples (n) and number of predictors (p).
n <- 1200
p <- 400

# Residual covariance matrix.
V <- rbind(c(1.0,0.2),
           c(0.2,0.4))
r <- nrow(V)

# True effects used to simulate the data.
B <- rbind(c(-2.0, -1.5),
           c( 1.0,  1.0),
           matrix(0,p - 2,r))

# Covariances in the mixture-of-normals prior on the regression
# coefficients.
S0 <- list(k1 = rbind(c(3,0),
                      c(0,3)),
           k2 = rbind(c(4,2),
                      c(2,4)),
           k3 = rbind(c(6,3.5),
                      c(3.5,4)),
           k4 = rbind(c(5,0),
                      c(0,0)),
           k5 = rbind(c(0,0),
                      c(0,0)))

# The mixture weights in the mixture-of-normals prior on the
# regression coefficients.
w0 <- c(0.1,0.2,0.2,0.1,0.4)

# SIMULATE DATA
# -------------
set.seed(1)
X <- matrix(rnorm(n*p),n,p)
X <- scale(X,scale = FALSE)

# Simulate Y ~ MN(X*B,I,V). Note that matrix.normal from the MBSP
# package appears to be much faster than rmatrixnorm from the
# MixMatrix package.
Y <- matrix.normal(X %*% B,diag(n),V)

# FIT MR-MASH MODEL
# -----------------
# Assign some missing values in Y
i  <- sample(n)
i1 <- i[seq(1,n/10)]
i2 <- i[seq(n/10+1,n/5)]
Y[i1,1] <- NA
Y[i2,2] <- NA

# FIT MR-MASH MODEL, ALLOWING FOR MISSING Ys
# ------------------------------------------
# Run 20 co-ordinate ascent updates.
B0 <- matrix(0,p,r)
t1 <- proc.time()
fit <- mr_mash_simple(X,Y,V,S0,w0,B0,numiter = 100,tol = 1e-6,
                      update_w0 = TRUE,update_V = TRUE,
                      verbose = TRUE)
t2 <- proc.time()
print(t2 - t1)


# Compare the posterior mean estimates of the regression coefficients
# against the coefficients used to simulate the data.
plot(B,fit$B,pch = 20,xlab = "true",ylab = "estimated", main = "Effects")
abline(a = 0,b = 1,col = "skyblue",lty = "dotted")

# Compare the estimate of the residual covariance
# against the true residual covariance used to simulate the data.
plot(V,fit$V,pch = 20,xlab = "true",ylab = "estimated", main = "Residual covariance")
abline(a = 0,b = 1,col = "skyblue",lty = "dotted")

