# This file contains the code for the (small) simulation study reported in
# the Psychometrika paper:
#
# "A Note on the Connection Between Trek Rules and Separable Nonlinear Least
#  Squares in Linear Structural Equation Models"
#
# Authors: Maximilian S. Ernst, Aaron Peikert, Andreas M. Brandmaier,
#          and Yves Rosseel
#
# This script requires:
# - a recent version of lavaan (>0.6-10) that supports the RAM representation
# - the file SNLLS_algorithm.R which provides an initial implementation of
#   the SNLLS algorithm for SEM

library("lavaan")
source("SNLLS_algorithm.R")

# Define the population model
pop.model <- '
    # factor loadings
    Y =~ 1*y1 + 0.8*y2 + 0.6*y3
    X =~ 1*x1 + 0.8*x2 + 0.6*x3

    # regression part
    Y ~ 0.25*X
'

# Model specification with simple starting values
model <- '
    # factor loadings
    Y =~ y1 + start(1)*y2 + start(1)*y3
    X =~ x1 + start(1)*x2 + start(1)*x3

    # regression part
    Y ~ start(0)*X

    # factor variances
    Y ~~ start(1)*Y
    X ~~ start(1)*X

    # observed residual variances
    y1 ~~ start(1)*y1
    y2 ~~ start(1)*y2
    y3 ~~ start(1)*y3
    x1 ~~ start(1)*x1
    x2 ~~ start(1)*x2
    x3 ~~ start(1)*x3
'

# Sample sizes
SAMPLE.SIZES <- c(10, 15, 20, 25, 30, 40, 50, 60, 70, 80, 90, 100)
NSIZES <- length(SAMPLE.SIZES)

# Numer of replications
REP <- 1000

# Container to store GLS results (converged/niterations)
GLS <- matrix(as.numeric(NA), nrow = REP*NSIZES, ncol = 2L)
colnames(GLS) <- c("converged", "nterations")

# Container to store SNLLS results (converged/niterations)
SNLLS <- matrix(as.numeric(NA), nrow = REP*NSIZES, ncol = 2L)
colnames(SNLLS) <- c("converged", "nterations")

# Start simulation
for(ss in 1:length(SAMPLE.SIZES)) {

    # set sample size
    N <- SAMPLE.SIZES[ss]

    # run REP replications for this sample size
    for(i in seq_len(REP)) {

        # be verbose
        cat("sample size = ", sprintf("%4d", N), " rep = ", sprintf("%4d", i),
            "\n")

        # set seed for this replication/sample-size
        this.seed <- 50000 + i + (ss - 1)*REP
        set.seed(this.seed)

        # row number in container
        idx <- i + (ss - 1)*REP

        # generate data
        Data <- simulateData(pop.model, sample.nobs = N, empirical = FALSE)

        # regular GLS
        fit <- try(suppressWarnings(sem(model, data = Data, optim.attempts = 1L,
                       baseline = FALSE, h1 = FALSE, check.post = FALSE,
                       se = "none", estimator = "GLS", representation = "RAM")),
               silent = TRUE)
        if(!inherits(fit, "try-error")) {
            if(lavInspect(fit, "converged")) {
                GLS[idx, 1] <- 1L
                GLS[idx, 2] <- lavInspect(fit, "iterations")
            } else {
                GLS[idx, 1] <- 0L
            }
        }

        # SNLLS
        fit.dummy <- sem(model, data = Data, estimator = "GLS",
                         do.fit = FALSE, representation = "RAM")
        lavmodel <- fit.dummy@Model
        S <- fit.dummy@SampleStats@cov[[1]]
        S.inv <- fit.dummy@SampleStats@icov[[1]]
        MLIST <- try(suppressWarnings(lav_optim_snlss(lavmodel =
                     fit.dummy@Model, S = S, S.inv = S.inv, verbose = FALSE)),
                     silent = TRUE)
        if(!inherits(MLIST, "try-error")) {
            if(MLIST$converged) {
                SNLLS[idx, 1] <- 1L
                SNLLS[idx, 2] <- MLIST$iterations
            } else {
                SNLLS[idx, 1] <- 0L
            }
        }

    } # i
} # ss

# Create table with number of non-converged replications
BAD.GLS   <- integer(NSIZES)
BAD.SNLLS <- integer(NSIZES)
for(i in seq_len(NSIZES)) {
    ss.idx <- (i-1)*REP + seq_len(REP)

    # converged - GLS
    idx <- ss.idx[GLS[ss.idx, 1] == 1L]
    BAD.GLS[i] <- REP - length(idx)

    # converged - SNLLS
    idx <- ss.idx[SNLLS[ss.idx, 1] == 1L]
    BAD.SNLLS[i] <- REP - length(idx)
}
Table.convergence <- as.data.frame(cbind(SAMPLE.SIZES, BAD.GLS, BAD.SNLLS))

# Create table with number of iterations (from nlminb)
NITER.GLS   <- integer(NSIZES)
NITER.SNLLS <- integer(NSIZES)
for(i in seq_len(NSIZES)) {
    ss.idx <- (i-1)*REP + seq_len(REP)

    # converged - GLS
    idx <- ss.idx[GLS[ss.idx, 1] == 1L]
    NITER.GLS[i] <- median(GLS[idx, "nterations"])

    # converged - SNLLS
    idx <- ss.idx[SNLLS[ss.idx, 1] == 1L]
    NITER.SNLLS[i] <- median(SNLLS[idx, "nterations"])
}
Table.iterations <- as.data.frame(cbind(SAMPLE.SIZES, NITER.GLS, NITER.SNLLS))

# Print tables
print(Table.convergence)
print(round(Table.iterations))