set.seed(1)
Sigma = cbind(c(1,0.7,0.2), c(0.7,1,0.4), c(0.2,0.4,1))
U0 = matrix(0,3,3)
U1 = matrix(0,3,3); U1[1,1] = 1
U2 = diag(3); U2[3,3] = 0
U3 = matrix(1,3,3)
data = generate_data(n=4000, p=3, V=Sigma, Utrue = list(U0=U0, U1=U1,U2=U2,U3=U3))

m.data = mash_set_data(data$Bhat, data$Shat)
m.1by1 = mash_1by1(m.data)
strong = get_significant_results(m.1by1)

U.pca = cov_pca(m.data, 3, subset = strong)
U.ed = cov_ed(m.data, U.pca, subset = strong)
U.c = cov_canonical(m.data)

Vhat.mle = estimateV(m.data, c(U.c, U.ed), init_rho = c(-0.5,0,0.5), tol=1e-4, optmethod = 'mle')
Vhat.mle$V

Vhat.em = estimateV(m.data, c(U.c, U.ed), init_rho = c(-0.5,0,0.5), tol = 1e-4, optmethod = 'em2')
Vhat.em$V

table = data.frame(rbind(Vhat.mle$ttime, Vhat.em$ttime), row.names = c('mle', 'em'))
colnames(table) = c('12','13','23')
table %>% kable() %>% kable_styling()

Vhat.emV = estimateV(m.data, c(U.c, U.ed), init_V = list(diag(ncol(m.data$Bhat)), clusterGeneration::rcorrmatrix(3), clusterGeneration::rcorrmatrix(3)), tol = 1e-4, optmethod = 'emV')
Vhat.emV$V

Vhat.tru = estimate_null_correlation(m.data)
Vhat.tru

V.mash = cor((data$Bhat/data$Shat)[-strong,])
V.mash

set.seed(100)
n=4000; p = 5
U0 = matrix(0,p,p)
U1 = U0; U1[1,1] = 1
U2 = U0; U2[c(1:2), c(1:2)] = 1
U3 = matrix(1, p,p)
Utrue = list(U0 = U0, U1 = U1, U2 = U2, U3 = U3)
for(t in 1:10){
  Vtrue = clusterGeneration::rcorrmatrix(p)
  data = generate_data(n, p, Vtrue, Utrue)
  # mash cov
  m.data = mash_set_data(Bhat = data$Bhat, Shat = data$Shat)
  m.1by1 = mash_1by1(m.data)
  strong = get_significant_results(m.1by1)

  U.pca = cov_pca(m.data, 3, subset = strong)
  U.ed = cov_ed(m.data, U.pca, subset = strong)
  U.c = cov_canonical(m.data)
  Vhat.mle <- estimateV(m.data, c(U.c, U.ed), init_rho = c(-0.5,0,0.5), tol=1e-4, optmethod = 'mle')
  Vhat.em <- estimateV(m.data, c(U.c, U.ed), init_rho = c(-0.5,0,0.5), tol=1e-4, optmethod = 'em2')
  Vhat.emV <- estimateV(m.data, c(U.c, U.ed), init_V = list(diag(ncol(m.data$Bhat)), clusterGeneration::rcorrmatrix(p), clusterGeneration::rcorrmatrix(p)),tol=1e-4, optmethod = 'emV')
  saveRDS(list(V.true = Vtrue, V.mle = Vhat.mle, V.em = Vhat.em, V.emV = Vhat.emV, data = data, strong=strong),
          paste0('../output/MASH.result.',t,'.rds'))
}

files = dir("../output/AddEMV/"); files = files[grep("MASH.result",files)]
times = length(files)
result = vector(mode="list",length = times)
for(i in 1:times) {
  result[[i]] = readRDS(paste("../output/AddEMV/", files[[i]], sep=""))
}

mle.pd = numeric(times)
em.pd = numeric(times)
for(i in 1:times){
  m.data = mash_set_data(result[[i]]$data$Bhat, result[[i]]$data$Shat)
  
  result[[i]]$V.trun = estimate_null_correlation(m.data, apply_lower_bound = FALSE)
  m.1by1 = mash_1by1(m.data)
  strong = get_significant_results(m.1by1)
  result[[i]]$V.1by1 = cor(m.data$Bhat[-strong,])
  U.c = cov_canonical(m.data)
  U.pca = cov_pca(m.data, 3, subset = strong)
  U.ed = cov_ed(m.data, U.pca, subset = strong)

  m.data.true = mash_set_data(Bhat = m.data$Bhat, Shat = m.data$Shat, V = result[[i]]$V.true)
  m.model.true = mash(m.data.true, c(U.c,U.ed), verbose = FALSE)
  
  m.data.trunc = mash_set_data(Bhat = m.data$Bhat, Shat = m.data$Shat, V = result[[i]]$V.trun)
  m.model.trunc = mash(m.data.trunc, c(U.c,U.ed), verbose = FALSE)
  
  m.data.1by1 = mash_set_data(Bhat = m.data$Bhat, Shat = m.data$Shat, V = result[[i]]$V.1by1)
  m.model.1by1 = mash(m.data.1by1, c(U.c,U.ed), verbose = FALSE)
  
  m.data.emV = mash_set_data(Bhat = m.data$Bhat, Shat = m.data$Shat, V = result[[i]]$V.emV$V)
  m.model.emV = mash(m.data.emV, c(U.c,U.ed), verbose = FALSE)
  
  # MLE
  
  m.model.mle = m.model.mle.F = m.model.mle.2 = list()
  R <- tryCatch(chol(result[[i]]$V.mle$V),error = function (e) FALSE)
  if(is.matrix(R)){
    mle.pd[i] = 1
    m.data.mle = mash_set_data(Bhat = m.data$Bhat, Shat = m.data$Shat, V = result[[i]]$V.mle$V)
    m.model.mle = mash(m.data.mle, c(U.c,U.ed), verbose = FALSE)
  }else{
    V.mle.near.F = as.matrix(Matrix::nearPD(result[[i]]$V.mle$V, conv.norm.type = 'F', keepDiag = TRUE)$mat)
    V.mle.near.2 = as.matrix(Matrix::nearPD(result[[i]]$V.mle$V, conv.norm.type = '2', keepDiag = TRUE)$mat)
    
    result[[i]]$V.mle.F = V.mle.near.F
    result[[i]]$V.mle.2 = V.mle.near.2
    
    # mashmodel
    m.data.mle.F = mash_set_data(Bhat = m.data$Bhat, Shat = m.data$Shat, V = V.mle.near.F)
    m.model.mle.F = mash(m.data.mle.F, c(U.c,U.ed), verbose = FALSE)
    
    m.data.mle.2 = mash_set_data(Bhat = m.data$Bhat, Shat = m.data$Shat, V = V.mle.near.2)
    m.model.mle.2 = mash(m.data.mle.2, c(U.c,U.ed), verbose = FALSE)
  }
  
  # EM
  m.model.em = m.model.em.F = m.model.em.2 = list()
  R <- tryCatch(chol(result[[i]]$V.em$V),error = function (e) FALSE)
  if(is.matrix(R)){
    em.pd[i] = 1
    m.data.em = mash_set_data(Bhat = m.data$Bhat, Shat = m.data$Shat, V = result[[i]]$V.em$V)
    m.model.em = mash(m.data.em, c(U.c,U.ed), verbose = FALSE)
  }else{
    V.em.near.F = as.matrix(Matrix::nearPD(result[[i]]$V.em$V, conv.norm.type = 'F', keepDiag = TRUE)$mat)
    V.em.near.2 = as.matrix(Matrix::nearPD(result[[i]]$V.em$V, conv.norm.type = '2', keepDiag = TRUE)$mat)
    
    result[[i]]$V.em.F = V.em.near.F
    result[[i]]$V.em.2 = V.em.near.2
    
    # mashmodel
    m.data.em.F = mash_set_data(Bhat = m.data$Bhat, Shat = m.data$Shat, V = V.em.near.F)
    m.model.em.F = mash(m.data.em.F, c(U.c,U.ed), verbose = FALSE)
      
    m.data.em.2 = mash_set_data(Bhat = m.data$Bhat, Shat = m.data$Shat, V = V.em.near.2)
    m.model.em.2 = mash(m.data.em.2, c(U.c,U.ed), verbose = FALSE)
  }
  
  result[[i]]$m.model = list(m.model.true = m.model.true, m.model.trunc = m.model.trunc, 
                             m.model.1by1 = m.model.1by1, m.model.emV = m.model.emV,
                             m.model.mle = m.model.mle,
                             m.model.mle.F = m.model.mle.F, m.model.mle.2 = m.model.mle.2, 
                             m.model.em = m.model.em,
                             m.model.em.F = m.model.em.F, m.model.em.2 = m.model.em.2)
}

sessionInfo()