library(foreach)
library(parallel)
library(iterators)
library(doParallel)
#library(mvtnorm)
library(MASS)
detectCores() # the number of CPU cores

ccf_k <- function(x, K, p)
{
  
  lm0 <- 0
  
  lm3 <- matrix(0,p,n)
  
  lm4 <- matrix(0,p,n)
  
  lm5 <- rep(0,p)
  
  for (j in 1:p)
  {
    lm0 = lm0 + log(K[j,j])
    
    for (k in 1:n)
    {
      for (i in 1:p)
      {
        if(i != j)
        {
          lm3[j,k] = lm3[j,k] + K[i,j]*x[k,i]
        }
        
      }
      lm4[j,k] =  K[j,j]*(x[k,j] + 1/K[j,j]*lm3[j,k])^2
      
    }
    lm5[j] = lm5[j] + sum(lm4[j,])
    
  }
  
  
  f_k <- -1/2*n*lm0 + 1/2*sum(lm5)
  
  return(f_k)
  
}

#ccf_k <- ccf_k(x, K, p) # -l_c(\theta) negative composite log-likelihood
#ccf_k

## grisearch function is used to do grid search for different turning parameters

gridsearch <- function(n,gridpoints,p,sigma)
  
{
  mu <- rep(0, p)
  x <- mvrnorm(n,mu,sigma) # observations
  le <- length(gridpoints)
  
  compbic_p <- Inf
  compbic <- Inf
  turning_b <- rep(0,4)
  tau_t = lambda_1 = lambda_2 = lambda_3 = 0
  
  
  for (i in 1:le)
  {
    tau_t = gridpoints[i]
    
    for (j in 1:le)
    {
      lambda_1 = lambda_2 = lambda_3 = gridpoints[j]
      K <- ccselection(x, p, tau_t, lambda_1, lambda_2, lambda_3,idd)[[1]]
      #print(K)  
      compbic <- 2*ccf_k(x,  MBVB(K,tau_t)[[2]], p) + MBVB(K,tau_t)[[3]]*log(n)
      
      if(compbic < compbic_p)
      {
        compbic_p = compbic
        turning_b = c(tau_t,lambda_1,lambda_2,lambda_3)
        #print(compbic_p)
        #print(turning_b)
      }
    }
    
    
    
  }
  
  #compbic_p
  #turning_b
  
  tau_t = turning_b[1]
  lambda_1  =  turning_b[2]
  lambda_2 = turning_b[3]
  lambda_3 = turning_b[4] 
  
  best <- ccselection(x, p, tau_t, lambda_1, lambda_2, lambda_3,idd)
  #best[[1]]
  
  color <- MBVB(best[[1]],tau_t)
  #color
  
  return(list(color[[1]],color[[2]]))
}

ccselection <- function(x, p, tau_t, lambda_1, lambda_2, lambda_3,idd)
{
  
  K = solve(cov(x))
  #K = diag(1,p,p)
  
  ig_d <- diag(rep(0, p), nrow = p, ncol = p)  # Gd, the initial value of diagnoal parameters \theta
  
  
  ig_0 <- matrix(c(0), p*(p-1)/2, p*(p-1)/2) # G0, the initial value of off-diagnoal parameters \theta
  
  
  
  for (i in 1:p)
  {
    ig_d[i,i] = K[i,i]
  }
  
  for (i in 1:(p-1))
  {
    for (j in (i+1):p)
    {
      ig_d[i,j] = K[i,i] - K[j,j]
    }
  }
  
  for (i in 1:(p*(p-1)/2))
  {
    ig_0[i,i] = K[idd[i, 1],idd[i, 2]]
  }
  
  for (i in 1:(p*(p-1)/2-1))
  {
    for (j in (i+1):(p*(p-1)/2))
    {
      ig_0[i,j] = K[idd[i, 1],idd[i, 2]] - K[idd[j, 1],idd[j, 2]]
    }
  }
  
  ig_0
  
  ig_d
  
  
  ## compute f_k
  
  cf_k <- function(ig_d, ig_0, K, tau_t)
  {
    
    lm1 <- 0
    
    for(i in 1:p)
      for(j in i:p)
      {
        if(i != j)
        {
          if(abs(ig_d[i,j]) < tau_t)
          { lm1 <- lm1 + lambda_1 * abs(ig_d[i,j])/tau_t
          }
          else
          {lm1 <- lm1 + lambda_1}
        }
      }
    
    lm2 <- 0
    
    for(i in 1:dim(ig_0)[1])
      for(j in i:dim(ig_0)[1])
      {
        if(i == j)
        { 
          if(abs(ig_0[i,j]) < tau_t)
          {
            lm2 <- lm2 + lambda_2 * abs(ig_0[i,j])/tau_t
          }
          else
          {
            lm2 = lm2 + lambda_2
          }
        }
        if(i != j)
        { 
          if(abs(ig_0[i,j]) < tau_t)
          {
            lm2 <- lm2 + lambda_3 * abs(ig_0[i,j])/tau_t
          }
          else
          {
            lm2 <- lm2 + lambda_3
          }
        }
      }
    
    lm0 <- 0
    
    lm3 <- matrix(0,p,n)
    
    lm4 <- matrix(0,p,n)
    
    lm5 <- rep(0,p)
    
    for (j in 1:p)
    {
      lm0 = lm0 + log(K[j,j])
      
      for (k in 1:n)
      {
        for (i in 1:p)
        {
          if(i != j)
          {
            lm3[j,k] = lm3[j,k] + K[i,j]*x[k,i]
          }
          
        }
        lm4[j,k] =  K[j,j]*(x[k,j] + 1/K[j,j]*lm3[j,k])^2
        
      }
      lm5[j] = lm5[j] + sum(lm4[j,])
      
    }
    
    
    f_k <- -1/2*lm0 + 1/(2*n)*sum(lm5)+ lm1 + lm2
    
    return(f_k)
    
  }
  
  f_k <- cf_k(ig_d, ig_0, K, tau_t) # need to minimize
  
  # f_k
  
  s <- (t(x)%*%x)/n
  
  # st is the soft thresholding operator
  
  st <- function(a_1, a_2)
  {
    if( abs(a_1)-a_2 > 0 )
    {
      if(a_1 > 0)
      {
        st_r <- abs(a_1) - a_2
      }
      if(a_1 < 0)
      {
        st_r <- a_2 - abs(a_1)
      }
      if(a_1 == 0) st_r <- 0
    }
    if(abs(a_1)- a_2 <= 0)
    {
      st_r <- 0
    }
    return(st_r)
  }
  
  
  # st(-3,2)
  
  ### First Iteration
  
  ## repeat{}
  
  for(kkk in 1:500)
  {
    
    ig_d_m <- ig_d
    
    ig_0_m <- ig_0
    
    K_m <- K
    
    VE_d <- matrix( c(0), nrow = p, ncol = p)
    
    for( i in 1:(p-1))
      for(j in (i+1):p)
      {
        if(abs(ig_d[i, i] - ig_d[j, j]) < tau_t)
          VE_d[i, j] <- 2
      }
    
    VE_0 <- matrix(c(0), nrow = p*(p-1)/2, ncol = p*(p-1)/2)
    
    for(i in 1:(p*(p-1)/2-1))
      for(j in (i+1):(p*(p-1)/2))
      {
        if(abs(ig_0[i, i] - ig_0[j, j]) < tau_t)
          VE_0[i, j] <- 2
      }
    
    for(i in 1:(p*(p-1)/2))
    {
      if(abs(ig_0[i, i]) < tau_t)
      {VE_0[i, i] <- 2}
    }
    
    VE_d
    VE_0
    
    a <- matrix(c(1), nrow = p, ncol = p) # for update a_jj'
    a[!upper.tri(a)] <- 0
    
    b <- matrix(rep(1, p*p), nrow = p, ncol = p) # for update b_jj'
    b[!upper.tri(b)] <- 0
    
    c <- matrix(c(1), nrow = p*(p-1)/2, ncol = p*(p-1)/2) # for update c_jj'
    c[!upper.tri(c)] <- 0
    
    d <- matrix(rep(1, p^2*(p-1)^2/4), nrow = p*(p-1)/2) # for update d_jj'
    d[!upper.tri(d)] <- 0
    
    tt <- 0
    
    ### 
    repeat{
      
      tt = tt + 1
      
      
      ig_d_1 <- ig_d
      
      ig_0_1 <- ig_0
      
      a_1 <- a
      
      b_1 <- b
      
      c_1 <- c
      
      d_1 <- d
      
      
      ### (1) update non-zero's of K[i, i] using (3)
      
      for(i in 1:p) # i <- 1
      {
        
        ## VE_d[i, j] == 0 if |ig_d[i,j]| > tau_t
        ## VE_d[i, j] == 1 if             = 0
        ##               2 if  0<  < tau_t
        
        trm3 <- 0 
        
        for (k in 1:p)
        {
          if(k != i )
          {
            for ( l in 1:p)
            {
              if(l != i)
              {
                trm3 = trm3 + K[k,i]*K[l,i]*t(x[,k])%*%x[,l]
              }
            } 
          }
          
        }
        
        trm3
        trm3 = c(trm3)
        
        trm1 <- 0; trm2 <- 0
        
        for (j in 1:p)
        {
          if(j < i)
          {
            if(VE_d[j, i] == 2)
            {
              trm1 <- trm1 + b[j,i]
              trm2 <- trm2 - a[j,i] - b[j,i]*(K[j,j] - ig_d[j,i])
            }
          }
          if(j > i)
          {
            if(VE_d[i, j] == 2)
            {
              trm1 <- trm1 + b[i,j]
              trm2 <- trm2 + a[i,j] - b[i,j]*(K[j,j] + ig_d[i,j])
            }
          }
        }
        
        trm2 = trm2 + 1/(2*n)*t(x[,i])%*%x[,i]
        trm2 = c(trm2)
        
        
        
        
        fun1 <- function(x)
        {
          y <- -1/(2*n)*trm3*x^(-2) - 1/2*x^(-1) + trm1*x + trm2
        }
        
        
        All <- uniroot.all( fun1, c(0, 10000))
        All
        
        if(length(All) == 1) { K[i, i] <- All}
        
        if(length(All) > 1)
        {
          f_k_c <- c(0)
          
          for(m in 1:length(All))
          {
            K[i, i] <- All[m]
            
            ig_d[i, i] <- All[m]
            ## compute f_k
            
            f_k_c[m] <- cf_k(ig_d, ig_0, K)
          }
          
          K[i, i] <- All[which(f_k_c == min(f_k_c) )]
        }
        
        
        
        ig_d[i, i] <- K[i, i]
        
        
      }
      
      K
      ig_d
      
      ### (2) update K[i,j] (i!=j)
      D = CC = X1 = X2 = SA = SB <- rep(0, p*(p-1)/2)
      
      
      for (j in 1: (p*(p-1)/2))
      {
        
        ###################
        for (k in 1: (p*(p-1)/2))
        {
          if(j < k)
          {
            if(VE_0[j,k] == 2)
            {
              D[j] = D[j] + d[j,k]
              CC[j] = CC[j] - c[j,k] + d[j,k]*(ig_0[k,k] + ig_0[j,k])
            }
          }
          
          if(j > k)
          {
            if(VE_0[k,j] == 2)
            {
              D[j] = D[j] + d[k,j]
              CC[j] = CC[j] + c[k,j] + d[k,j]*(ig_0[k,k] - ig_0[k,j])
            }
          }   
        }
        
        #####################
        
        
        for (k in 1:p)
        {
          if(k != idd[j,][1] & k != idd[j,][2])
          {
            X1[j] = X1[j] + t(x[,idd[j,][2]])%*%x[,k]*K[k,idd[j,][1]]
            X2[j] = X2[j] + t(x[,idd[j,][1]])%*%x[,k]*K[k,idd[j,][2]]
          }
        }
        
        
        
        
        ######################
        
        
        SA[j] = 1/n*K[idd[j,][1],idd[j,][1]]^(-1)*t(x[,idd[j,][2]])%*%x[,idd[j,][2]] + 
          1/n*K[idd[j,][2],idd[j,][2]]^(-1)*t(x[,idd[j,][1]])%*%x[,idd[j,][1]] + D[j]
        
        SB[j] = -1/n*t(x[,idd[j,][1]])%*%x[,idd[j,][2]] - 1/n*K[idd[j,][1],idd[j,][1]]^(-1)*X1[j] - 1/n*t(x[,idd[j,][2]])%*%x[,idd[j,][1]] -
          1/n*K[idd[j,][2],idd[j,][2]]^(-1)*X2[j] + CC[j]
        
        
        #####################
        
        
        if(VE_0[idd[j,][1],idd[j,][2]] == 2)
        {
          K[idd[j,][1],idd[j,][2]] = K[idd[j,][2],idd[j,][1]] = st(SB[j]/SA[j], lambda_2/tau_t/SA[j])
        }
        else
        {
          K[idd[j,][1],idd[j,][2]] =  K[idd[j,][2],idd[j,][1]] = SB[j]/SA[j]
        }
        
        
        
        ig_0[j,j] = K[idd[j,][1],idd[j,][2]]
      }
      
      K
      ig_0
      
      ################################################################################################
      
      
      ### (3) update ig_d[j, j'], j< j' in E_d    0 <   < tau_t
      
      
      for(i in 1:(p-1))
      { 
        
        for(j in (i+1):p)
        {
          if(VE_d[i, j] == 2)
          {
            a_11 <- (a[i,j] + b[i,j]*(ig_d[i,i] - ig_d[j,j]))/b[i,j]
            
            a_22 <- lambda_1/(tau_t*b[i,j])
            
            ig_d[i, j] = ig_d[j, i] <- st(a_11, a_22)
          }
        }
      }
      
      ig_d
      
      
      
      #### (4) update ig_0[j, j'], j < j] in E_0  0 <  < tau_t
      
      for(i in 1:(dim(idd)[1]-1))
        for(j in (i+1):dim(idd)[1])
        {
          if(VE_0[i,j] == 2)
          {
            b_1 <- (c[i,j] + d[i,j]*(ig_0[i,i] - ig_0[j,j]))/d[i,j]
            
            b_2 <- lambda_3/(tau_t*d[i,j])
            
            ig_0[i,j] = ig_0[j,i] <- st(b_1, b_2)
          }
        }
      
      ig_0
      
      ### second layer repeat
      
      if( max(abs(ig_d - ig_d_1) ) < e & max(abs(ig_0 - ig_0_1)) < e )
      {break}
      
      for(i in 1:(p-1))
        for(j in (i+1):p)
        {
          a[i,j] <- a[i,j] + b[i,j]*(ig_d[i,i]-ig_d[j,j] - ig_d[i,j])
          
          b[i,j] <- rho*b[i,j]
        }
      
      
      for(i in 1:(dim(idd)[1]-1))
        for(j in (i+1):dim(idd)[1])
        {
          c[i,j] <- c[i,j] + d[i,j]*(ig_0[i,i] - ig_0[j,j] - ig_0[i,j])
          
          d[i,j] <- rho*d[i,j]
        }
      
      a
      b
      c
      d
      K
      ig_d
      ig_0
      
    }
    
    # K
    # sigma
    
    if( max(abs(ig_d - ig_d_m) ) < e & max(abs(ig_0 - ig_0_m)) < e )
    {break}
   
    kkk = kkk + 1
    
  }
  
  #proc.time() - time1
  
  kkk
  K
  #K_0
  
  return(list(K, VE_0, VE_d))
}


# function MBVB is used to obtain the best colored model and the estimated K corresponding to the best colored model

MBVB <- function(K,tau_t)
{
  
  modelB <- matrix(1,p,p)
  EVB <- K
  qq <- 2
  su <- c()
  inde <- c()
  su[qq] = 0
  inde[qq] = 1
  
  # the beginning of vertertices
  
  for (i in 1:(p-1))
  {
    flag = F
    if(modelB[i,i] == 1)
    {
      for (j in (i+1):p)
      {
        if (abs(K[j,j]-K[i,i]) < tau_t)
        {
          modelB[j,j] = qq
          inde[qq] = inde[qq] + 1
          su[qq] = su[qq] + K[j,j] 
          flag = T
        }
      }
      if(flag == T)
      {
        modelB[i,i] = qq
        su[qq] = su[qq] + K[i,i]
        qq = qq + 1
        su[qq] = 0
        inde[qq] = 1
      }
    }
  }
  
  #the beginning for calculate the average of vertices
  
  for (i in 1:p)
  {
    if(qq > 2)
    {
      for (j in 2:(qq-1))
      {
        if(modelB[i,i] == j)
        {
          EVB[i,i] = su[j]/inde[j]
        }
      }
    }
  }
  
  vq <- qq
  
  modelB
  
  # the beginning for edges with 0 constrains
  
  for (i in 1:(p-1))
  {
    for (j in (i+1):p)
    {
      if(abs(K[i,j]) < tau_t)
      {
        modelB[i,j] = modelB[j,i] = 0
        EVB[i,j] = EVB[j,i] = 0
      }
    }
  }
  
  # the beginning for edges without edge constrains
  
  for (i in 1:(p-2))
  {
    for (j in (i+1):p)
    {
      flag = F
      if(modelB[i,j] == 1)
      {
        ####################### 
        
        for (ii in (i+1):(p-1))
        {
          for (jj in (ii+1):p)
          {
            if(abs(K[i,j]-K[ii,jj]) < tau_t)
            {
              modelB[ii,jj] = modelB[jj,ii] = qq
              inde[qq] = inde[qq] + 1
              su[qq] = su[qq] + K[ii,jj] 
              flag = T
            }
          } 
        }
        
        if(j != p)
        {
          for (kk in (j+1):p)
          {
            if(abs(K[i,j]-K[i,kk]) < tau_t)
            {
              modelB[i,kk] = modelB[kk,i] = qq
              inde[qq] = inde[qq] + 1
              su[qq] = su[qq] + K[i,kk] 
              flag = T
            }
          }
        }
        ########################     
        
        
        if(flag == T)
        {
          modelB[i,j] = modelB[j,i] = qq
          su[qq] = su[qq] + K[i,j]
          qq = qq + 1
          su[qq] = 0
          inde[qq] = 1
        }
      }
      
    }
  }
  
  # the beginning for calculating the summation for the edges without 0 constraints
  
  for (i in 1:(p-1))
  {
    for (j in (i+1):p)
    {
      if(qq > vq)
      {
        for (k in vq:(qq-1))
        {
          if(modelB[i,j] == k)
          {
            EVB[i,j] = EVB[j,i] = su[k]/inde[k]
          }
        }
      }
    }
  }
  
  modelB # the best model
  EVB # the estimated K corresponding to the best model
  
  df <- 0
  for(i in 1:p)
  {
    for(j in i:p)
    {
      if(modelB[i,j] == 1)
      {df = df + 1}
    }
  }
  
  df = df + qq - 2
  
  df
  
  return(list(modelB,EVB,df))
  
}


indicator <- function(c) ifelse(c <= tau_t, 1, 0)

rho <- 1.1 # rho for updating a_jj', b_jj', c_jj', d_jj'

e <- 10^(-2) # error for iterating ig_d and ig_0
n <- 1000 # number of observations

######################## model setting

p <- 10 #p=30 done

K_0 <- matrix(0,p,p)
model <- matrix(0,p,p)

for(i in 1: (p/2))
{
  K_0[2*i-1,2*i-1] = 1
  model[2*i-1,2*i-1] = 2
}
for(i in 1: (p/2))
{
  K_0[2*i,2*i] = 1.5
  model[2*i,2*i] = 3
}
for(i in 1: (p/2))
{
  K_0[2*i-1,2*i] = 0.5
  K_0[2*i,2*i-1] = 0.5
  model[2*i-1,2*i] = 4
  model[2*i,2*i-1] = 4
}
for(i in 1: ((p-2)/2))
{
  K_0[2*i,2*i+1] = 0.3
  K_0[2*i+1,2*i] = 0.3
  model[2*i,2*i+1] = 5
  model[2*i+1,2*i] = 5
}
K_0[p, 1] = 0.3
K_0[1, p] = 0.3
model[p, 1] = 5
model[1, p] = 5
K_0
model

sigma <- solve(K_0)

eigen(sigma)$values
eigen(K_0)$values

model

id <- NULL

for(i in 1:(p-1))
  for(j in (i+1):p)
  {
    id <- c(id, i, j)
  }

idd <- matrix(as.vector(id), ncol = 2, byrow = T)

#idd

gridpoints <- c(0.01, 0.05, 0.1, 0.5, 1)

cl <- makeCluster(3)
registerDoParallel(cl)
getDoParWorkers()


ptm <- proc.time()
t <- 3
A <- foreach(i=1:t, .combine='cbind') %dopar% {
  library(rootSolve)
  library(MASS)
  gridsearch(n,gridpoints,p,sigma)
}


stopImplicitCluster()

times <- proc.time()-ptm
times
A[1,]
A[2,]