library(OpenMx)

## Set the number of individuals and of repeated measures and seed
n <- 500; p <- 10

## Set "true" values to parameters
###################################
### Set true values of parameters of TICs
sd.x1 <- 1; sd.x2 <- 1; rho12 <- 0.3

### Set true values of parameters of latent variables
eta0.mean <- 100; eta0.var <- 25
eta1.mean <- -5; eta1.var <- 1
rho <- 0.3

beta <- matrix(c(1.8, 1.8,
                 0.4, 0.4), nrow = 2, byrow = T)/sqrt(2)

sd <- 1

### Mean vectors of parameters
mean0 <- c(eta0.mean, eta1.mean, 0, 0)
### Var-cov matrix of parameters
#### The psi matrix
eta0eta1 <- rho * sqrt(eta0.var * eta1.var) 
eta <- matrix(c(eta0.var, eta0eta1, 
                eta0eta1, eta1.var), nrow = 2, ncol = 2)

#### The Phi matrix
var.x1 <- sd.x1^2; var.x2 <- sd.x2^2; x1x2 <- rho12 * sd.x1 * sd.x2

phi <- matrix(c(var.x1, x1x2, x1x2, var.x2), nrow = 2, ncol = 2)
psi <- eta - beta %*% phi %*% t(beta)
#### Calculate the model implied variance-covariance matrix
sigma.yy <- beta %*% phi %*% t(beta) + psi
sigma.xy <-  beta %*% phi
sigma0 <- rbind(cbind(sigma.yy, sigma.xy), cbind(t(sigma.xy), phi))
init <- c(mean0[1:2], psi[row(psi) >= col(psi)], c(beta),
          mean0[3:4], var.x1, x1x2, var.x2)

### Generate individual intercept, slope1, slope2 and knot
eta0.eta1.x <- MASS::mvrnorm(n, mu = mean0, Sigma = sigma0)

### Long-format data framework
long_dat <- expand.grid(id = 1:n, time = 0:(p - 1))
long_dat <- long_dat[order(long_dat$id, long_dat$time), ]

### Calculate the repeated measurement for each individual at eath time-point
score <- c(NA)
for (i in 1:nrow(long_dat)){
  df <- long_dat[i, ]
  score[i] <- eta0.eta1.x[df$id, 1] + 
    df$time * eta0.eta1.x[df$id, 2] + rnorm(1, mean = 0, sd = sd)
}
long_dat$measurement <- score

### Convert long format to wide format
dat <- reshape2::dcast(long_dat, id ~ time, value.var = "measurement")
colnames(dat) <- c("id", paste0("Y", 0:9))

dat$x1 <- eta0.eta1.x[, 3]
dat$x2 <- eta0.eta1.x[, 4]

model_mx <- mxModel("Trajectories with TICs", type = "RAM",
                    manifestVars = c(paste0("Y", 0:9), "x1", "x2"), latentVars = c("eta0", "eta1"),
                    mxData(observed = dat[, c(-1, -(2 * p + 2))], type = "raw"),
                    #### Define factor loadings from latent variables to manifests
                    mxPath(from = "eta0", to = paste0("Y", 0:9), arrows = 1, free = F, values = 1),
                    mxPath(from = "eta1", to = paste0("Y", 0:9), arrows = 1, free = F, values = 0:9),
                    #### Define the variances of residuals
                    mxPath(from = paste0("Y", 0:9), to = paste0("Y", 0:9), arrows = 2, free = T, values = sd^2,
                           labels = "residuals"),
                    #### Define means of latent variables/intercepts of regression coefficients
                    mxPath(from = "one", to = c("eta0", "eta1"), arrows = 1, free = T, values = init[1:2],
                           labels = c("mueta0", "mueta1")),
                    #### Define var-cov matrix of latent variables
                    mxPath(from = c("eta0", "eta1"), to = c("eta0", "eta1"), arrows = 2,
                           connect = "unique.pairs", free = T,
                           values = init[3:5],
                           labels = c("psi00", "psi01", "psi11")),
                    #### Include time-invariant covariates
                    ##### Means
                    mxPath(from = "one", to = c("x1", "x2"), arrows = 1, free = T, 
                           values = c(0, 0), labels = c("mux1", "mux2")),
                    mxPath(from = c("x1", "x2"), to = c("x1", "x2"), connect = "unique.pairs", 
                           arrows = 2, free = T,
                           values = c(sd.x1^2, rho12 * sd.x1 * sd.x2, sd.x2^2), 
                           labels = c("phi11", "phi12", "phi22")),
                    ##### Regression coefficients
                    mxPath(from = "x1", to = c("eta0", "eta1"), arrows = 1, free = T, 
                           values = init[6:7], labels = paste0("beta1", 0:1)),
                    mxPath(from = "x2", to = c("eta0", "eta1"), arrows = 1, free = T, 
                           values = init[8:9], labels = paste0("beta2", 0:1)))
model0 <- mxTryHard(model_mx, extraTries = 9, 
                    initialGradientIterations = 20, OKstatuscodes = 0)
output <- summary(model0)
estimates <- output$parameters[, 1:6]
estimates$true <- init[c(6:9, 12:14, 3:5, 10:11, 1:3)]
write.table(dat, file = "growth.dat", row.names = F, col.names = F)