library(trunmnt)
library(mvtnorm)
library(tmvtnorm) 
library(MomTrunc)

set.seed(123)
sigma2e <- 1 ; sigma2u <- 2          # Error variance and random effect variance
n  <- 5 ; k <- 3
u  <- rnorm(k, sd = sqrt(sigma2u))
Z  <- model.matrix(~ factor(rep(1:k, each = n)) - 1) # Design matrix for random
D  <- diag(sigma2u, k)                               # effects
S  <- Z %*% D %*% t(Z) + diag(sigma2e, n*k)
mu <- rep(seq(-1,1, length.out = n), k)              # Mean vector mu
y  <- mu + Z %*% u + rnorm(n*k, sd = sqrt(sigma2e))
w  <- ifelse(y >= 0, 1, 0)
a  <- rep(-Inf, n*k)
b  <- rep( Inf, n*k)    
a[w == 1] <- 0
b[w == 0] <- 0

obj <- mtrunmnt(mu[1:5], a[1:5], b[1:5], Sigmae = 1, D = 2)
# obj <- mtrunmnt(mu[1:5], a[1:5], b[1:5], Sigma = S[1:5,1:5])
probntrun(obj)
c(prodmnt(obj, c(4,0,0,0,0)), prodmnt(obj, c(3,1,0,0,0)),
  prodmnt(obj, c(2,2,0,0,0)), prodmnt(obj, c(2,1,1,0,0)),
  prodmnt(obj, c(1,1,1,1,0)))
meanvar(obj)


# Test q = 1
obj <- mtrunmnt(mu[1:5], a[1:5], b[1:5], Sigmae = 1, D = 2)
trun <- meanvar(obj)
tmv  <- mtmvnorm(mu[1:5], S[1:5, 1:5], lower = a[1:5], upper = b[1:5])
Mom  <- meanvarTMD(a[1:5], b[1:5], mu[1:5], S[1:5, 1:5], dist="normal")


# Test q = 2

obj <- mtrunmnt(mu[1:10], a[1:10], b[1:10], Sigma = S[1:10,1:10])
trun <- meanvar(obj)
tmv  <- mtmvnorm(mu[1:10], S[1:10, 1:10], lower = a[1:10], upper = b[1:10])
Mom  <- meanvarTMD(a[1:10], b[1:10], mu[1:10], S[1:10, 1:10], dist="normal")

# Test q = 3

obj <- mtrunmnt(mu, a, b, Sigma = S)
trun <- meanvar(obj)
tmv  <- mtmvnorm(mu, S, lower = a, upper = b)
Mom  <- meanvarTMD(a, b, mu, S, dist="normal")

