library(mvtnorm)
library(denoiseR)
source("~/GMM_synthet_PP/GMM_algos.R")

L <- 2
N <- 10000
p <- 2

# cluster weights
alpha0 <- 1:2
alpha0 <- alpha0 / sum(alpha0)
alpha0



n <- 100
N <- n*ngroup
# cluster means
mu <- rmvnorm(n = L,
              mean = rep(0, p),
              sigma = 5 * diag(p))

# covariance matrix
Sigma <- diag(p) # think later how to set Sigma
Sigma <- Sigma + 0.1*LRsim(p, p, 1, 1)$mu
Sigma <- Sigma + t(Sigma)
Sigma <- Sigma/norm(Sigma)

# Z variables (unobserved)
Z <- rmultinom(N, 1, alpha0)

# Y variables (observed)
Y <- sapply(1:N, function(i) {
  cluster <- which(Z[, i] == 1)
  rmvnorm(1, mu[cluster,], Sigma)
})

groups <- as.factor(rep(1:n, each = ngroup))
Nc <- aggregate(rep(1, N), list(groups), sum)[, 2]
Ylist <- list()
for (c in 1:n) {
  Ylist[[c]] <- t(Y[, which(groups == c)])
}

for(prob in c(0.75)){
  
  tt <-
    FedEM_gmm(
      Ylist = Ylist,
      groups = groups,
      L = 2,
      maxiter = 2*1e3,
      thresh = 1e-20,
      gamma = 1e-2,
      nbatch = 20,
      Sigma_init = Sigma,
      bits=bits
    )
  fileFedEM <- paste('~/GMM_synthet_PP/results/FedEM_traj_normH_prob', 100*prob, '.txt', sep="")
  write.table(tt$normH, file=fileFedEM)
  fileFedEM_h <- paste('~/GMM_synthet_PP/results/FedEM_traj_normh_prob', 100*prob, '.txt', sep="")
  write.table(tt$normh, file=fileFedEM_h)
  
  
  ttSpid <-
    FedSpiderEM_gmm(
      Ylist = Ylist,
      groups = groups,
      L = 2,
      kout = 200,
      kin = 20,
      thresh = 1e-20,
      gamma = 1e-2,
      Sigma_init = Sigma,
      nbatch = 5,
      bits=bits
    )
  fileFedSpidEM <- paste('~/GMM_synthet_PP/results/FedSpidEM_traj_normH_prob', 100*prob, '.txt', sep="")
  write.table(unlist(ttSpid$normH), file=fileFedSpidEM)
  fileFedSpidEM_h <- paste('~/GMM_synthet_PP/results/FedSpidEM_traj_normh_prob', 100*prob, '.txt', sep="")
  write.table(unlist(ttSpid$normh), file=fileFedSpidEM_h)
}

