#' PriorThetaY list
#' @description generate prior value for parameter Theta and Y.
#' @importFrom gtools rdirichlet
#' @importFrom mvtnorm rmvnorm
#' @import stats
#' @param m the number of cluster
#' @param n sample size
#' @param p number of covariates
#' @param muBar parameter
#' @param hparam hyperparameters
#' @param qVec the vector of the number of factors in each clusters
#' @param ZOneDim ZOneDim
#' @param constraint constraint
#' @noRd
generate_prior_theta_y <- function(m,
                                n,
                                p,
                                muBar,
                                hparam,
                                qVec,
                                ZOneDim,
                                constraint) {
  ggamma <- hparam@ggamma
  delta <- hparam@delta
  bbeta <- hparam@bbeta
  alpha1 <- hparam@alpha1
  alpha2 <- hparam@alpha2

  # prior tao
  if (m == 1) {
    tao <- rbeta(n = 1, shape1 = 1, shape2 = m)
  } else {
    tao <- gtools::rdirichlet(n = 1, alpha = rep(ggamma, m))
  }

  # prior psy
  psy <- generate_prior_psi(p, m, delta, bbeta, constraint)

  # prior M
  M <- list()
  for (i in 1:m) {
    M[[i]] <- (mvtnorm::rmvnorm(n = 1, mean = muBar, sigma = 1 / alpha1 * psy[[i]]))
  }

  # prior lambda
  lambda <- generate_prior_lambda(p, m, alpha2, qVec, psy, constraint)


  # Zmat = get_z_mat_r(ZOneDim, m, n)
  Zmat <- get_z_mat_cpp(ZOneDim, m, n)

  # post Y
  Y <- list()
  for (k in 1:m) {
    Y[[k]] <- matrix(NA, qVec[k], n)
    for (i in 1:n) {
      Y[[k]][, i] <- mvtnorm::rmvnorm(1, mean = rep(0, qVec[k]), sigma = diag(qVec[k]))
    }
  }

  new("ThetaYList",
    tao = tao,
    psy = psy,
    M = M,
    lambda = lambda,
    Y = Y
  )
}

#' (internal)
#' @noRd
#' @noRd
evaluate_prior <- function(m,
                          p,
                          muBar,
                          hparam,
                          thetaYList,
                          ZOneDim,
                          qVec,
                          constraint,
                          clusInd) {
  ggamma <- hparam@ggamma
  alpha1 <- hparam@alpha1
  alpha2 <- hparam@alpha2
  bbeta <- hparam@bbeta
  delta <- hparam@delta

  loopm <- which(clusInd == 1)

  ## 2.2: Y
  # Yval = 0
  # for(k in loopm){
  #   Yval = Yval + sum(dnorm(thetaYList@Y[[k]], log = T))
  # }


  ## 2.3: Z
  # adjustTao = thetaYList@tao/sum(thetaYList@tao, na.rm = T)

  Zval <- sum(log(thetaYList@tao[ZOneDim]), na.rm = TRUE)
  # Zval = log(sum(adjustTao[ZOneDim], na.rm = T))
  ## 2.6: tao
  # taoVal = log(gtools::ddirichlet(x = adjustTao[!is.na(adjustTao)], alpha = rep(ggamma, m)))
  # print(thetaYList@tao[clusInd == 1])
  # print(rep(ggamma, m))

  taoVal <- log(gtools::ddirichlet(x = thetaYList@tao[clusInd == 1], alpha = rep(ggamma, m)))

  ## test
  # Zval = 0;taoVal = 0
  ##


  ## 2.7: M
  Mval <- 0
  for (k in loopm) {
    # cat("k =", k , "===> \n")
    Mval <- Mval + mvtnorm::dmvnorm(c(thetaYList@M[[k]]), mean = c(muBar), sigma = 1 / alpha1 * thetaYList@psy[[k]], log = T)
  }

  ## 2.8: lambda
  lambdaVal <- evaluate_prior_lambda(
    p, m, alpha2, qVec, thetaYList@psy, thetaYList@lambda, constraint,
    clusInd
  )

  ## 2.9: psy
  psyVal <- evaluate_prior_psi(thetaYList@psy, p, m, delta, bbeta, constraint, clusInd)


  # print(Zval)
  # print(taoVal)
  # print(Mval)
  # print(lambdaVal)
  # print(psyVal)
  # totalVal = sum(Yval + Zval + taoVal + Mval + lambdaVal + psyVal)


  # print(c( Zval,taoVal,Mval,lambdaVal,psyVal))
  totalVal <- sum(Zval + taoVal + Mval + lambdaVal + psyVal)
  return(totalVal)
}
#' generate_prior_psi
#'
#' @description generate prior value for parameter Psi
#' @import stats
#' @param p the number of features
#' @param m the number of clusters
#' @param delta hyperparameters
#' @param bbeta hyperparameters
#' @param constraint the pgmm constraint, a vector of length three with binary entry. For example, c(1,1,1) means the fully constraint model
#'
#' @noRd
generate_prior_psi <- function(p,
                             m,
                             delta,
                             bbeta,
                             constraint) {
  psy <- list()

  if (constraint[2] == T & constraint[3] == T) {
    for (i in 1:m) {
      if (i == 1) {
        psyValue <- 1 / rgamma(1, shape = delta, rate = bbeta)
        psy[[i]] <- diag(rep(psyValue, p))
      } else {
        psy[[i]] <- psy[[1]]
      }
    }
  } else if (constraint[2] == T & constraint[3] == F) {
    for (i in 1:m) {
      if (i == 1) {
        psy[[i]] <- diag(1 / rgamma(p, shape = delta, rate = bbeta))
      } else {
        psy[[i]] <- psy[[1]]
      }
    }
  } else if (constraint[2] == F & constraint[3] == T) {
    for (i in 1:m) {
      psyValue <- 1 / rgamma(1, shape = delta, rate = bbeta)
      psy[[i]] <- diag(rep(psyValue, p))
    }
  } else {
    for (i in 1:m) {
      psy[[i]] <- diag(1 / rgamma(p, shape = delta, rate = bbeta))
    }
  }
  return(psy)
}

#' (internal)
#' @noRd
#' @noRd
evaluate_prior_psi <- function(psy,
                             p,
                             m,
                             delta,
                             bbeta,
                             constraint,
                             clusInd) {
  evaluate_prior_psi_cpp(psy, p, m, delta, bbeta, constraint, clusInd)
}


#' generate_prior_lambda
#'
#' @description evaluate prior value for parameter Lambda
#' @importFrom mvtnorm rmvnorm
#' @param p the number of features
#' @param m the number of clusters
#' @param alpha2 hyper parameter
#' @param qVec parameter
#' @param psy parameter
#' @param constraint parameter
#'
#' @noRd
generate_prior_lambda <- function(p,
                                m,
                                alpha2,
                                qVec,
                                psy,
                                constraint) {
  lambda <- list()
  if (constraint[1] == T & constraint[2] == T) {
    for (k in 1:m) {
      if (k == 1) {
        qk <- qVec[k]
        lambda[[k]] <- matrix(0, p, qk)
        for (j in 1:qk) {
          lambda[[k]][, j] <- mvtnorm::rmvnorm(1, rep(0, p), 1 / alpha2 * psy[[k]])
        }
      } else {
        lambda[[k]] <- lambda[[1]]
      }
    }
  } else if (constraint[1] == T & constraint[2] == F) {
    psyAve <- matrix(0, nrow = p, ncol = p)
    for (k in 1:m) {
      psyAve <- psyAve + solve(psy[[k]])
    }
    psyAve <- solve(1 / m * psyAve)

    for (k in 1:m) {
      if (k == 1) {
        qk <- qVec[k]
        lambda[[k]] <- matrix(0, p, qk)
        for (j in 1:qk) {
          lambda[[k]][, j] <- mvtnorm::rmvnorm(1, rep(0, p), 1 / alpha2 * psyAve)
        }
      } else {
        lambda[[k]] <- lambda[[1]]
      }
    }
  } else {
    for (k in 1:m) {
      qk <- qVec[k]
      lambda[[k]] <- matrix(0, p, qk)
      for (j in 1:qk) {
        lambda[[k]][, j] <- mvtnorm::rmvnorm(1, rep(0, p), 1 / alpha2 * psy[[k]])
      }
    }
  }
  return(lambda)
}

#' (internal)
#' @noRd
#' @noRd
evaluate_prior_lambda <- function(p,
                                m,
                                alpha2,
                                qVec,
                                psy,
                                lambda,
                                constraint,
                                clusInd) {
  evaluate_prior_lambda_cpp(p, m, alpha2, qVec, psy, lambda, constraint, clusInd)
}
