### Set Up ----
rm(list = ls())

# devtools::install_github("manuelarnold/semtree", force=T)
library(semtree)
library(lavaan)
library(future)



# Factor model (lavaan) - Homogeneous sample ----
## Generate data
m <- '
f =~ 1*y1 + 1*y2 + 1*y3
f ~~ 1*f
y1 ~~ 0.5*y1
y2 ~~ 0.5*y2
y3 ~~ 0.5*y3
'

m_samp <- '
f =~ y1 + y2 + y3
f ~~ 1*f
y1 ~~ y1
y2 ~~ y2
y3 ~~ y3
'

N <- 1000
ntrees <- 10
set.seed(10)
data <- simulateData(model = m, sample.nobs = N)

## Generate covariates
noise_c1 <- rnorm(N)
noise_o1 <- rbinom(N, 4, 0.5) + 1; noise_o1 <- factor(noise_o1, ordered=T)
noise_m1 <- round(runif(N, 1, 3)); noise_m1 <- factor(noise_m1, ordered=F)
noise_d1 <- rbinom(N, 1, 0.5); noise_d1 <- factor(noise_d1, ordered=F)

covariates <- data.frame(noise_c1,
                         noise_o1,
                         noise_m1,
                         noise_d1)
covnames <- names(covariates)

data_complete <- cbind(data, covariates)

## Shuffle the data
shuffle <- sample(1:N, size=N, replace=F)
data_complete <- data_complete[shuffle,]

## Fit
fit <- lavaan(model = m_samp, data = data_complete)
# summary(fit)

## Controls
tree_ctrl <- semtree.control(method = "score", min.bucket = 100, verbose = T, refit=TRUE, bonferroni = T)
forest_ctrl <- semforest.control(num.trees = 5, control = tree_ctrl, mtry = 2)

# tree <- semtree(model = fit, data = data_complete, control = tree_ctrl, predictors = covnames)


## Run forest
ncores <- availableCores() - 2
plan(multisession, workers = ncores)
forest <- semforest(model = fit, 
                    data = data_complete, 
                    control = forest_ctrl,
                    predictors = covnames)
plan(sequential)