library(OpenMx)
mxOption(key="Default optimizer", value="NPSOL")

# ----------------------------------------------------------------------------------------------------------------------
# PREPARE DATA (SIMULATION)

# Set parameters
set.seed(123456789)
n <- 10000   # n/2 per zygosity
VA1 <- 0.4   # VE1 = 1-VA1
cut <- 0.8   # z-value

# Prepares data frame
dt <- data.frame(matrix(NA,nrow=n,ncol=3,dimnames = list(NULL,c("x1","x2","rel"))))

# Simulates genetic components
A1_mz <- MASS::mvrnorm(n/2,c(0,0), matrix(c(1,1,
                                            1,1),2,2,T))
A1_dz <- MASS::mvrnorm(n/2,c(0,0), matrix(c( 1,.5,
                                            .5, 1),2,2,T))
# Weighs genetic component and adds environmental component.
dt[1:(n/2)   , 1:2]   <- sqrt(VA1)*A1_mz + rnorm(n,0,sqrt(1-VA1))
dt[(n/2+1):n , 1:2]   <- sqrt(VA1)*A1_dz + rnorm(n,0,sqrt(1-VA1))

# Adds relatedness coefficients to data frame
dt[,"rel"] <- c(rep(1,n/2),rep(0.5,n/2))

# Dichotomises the data
dt[,1:2] <- ifelse(dt[,1:2] > cut,1,0)

# Select Variables for Analysis
selVars      <- c('x1','x2')    # list of variables namess

# Converts to factor variables
dt[,selVars] <- mxFactor(dt[,selVars], levels = c(0,1))


# ----------------------------------------------------------------------------------------------------------------------
# PREPARE MODEL

#####################################################
####          Creates Data Objects              ####
#####################################################
dataMZ    <- mxData( observed=subset(dt, select = selVars, subset = rel == 1), type="raw" )
dataDZ    <- mxData( observed=subset(dt, select = selVars, subset = rel == 0.5), type="raw" )

#####################################################
####         Creates Expectation Objects         ####
#####################################################
expCovMZ      <- mxMatrix(name="expCovMZ", type="Symm", nrow=2, ncol=2, free=c(F,T,F), values=c(1,0.5,1), label=c("var1", "MZ_r", "var1") ) 
expCovDZ      <- mxMatrix(name="expCovDZ", type="Symm", nrow=2, ncol=2, free=c(F,T,F), values=c(1,0.5,1), label=c("var1", "DZ_r", "var1") )

         
threMZ      <- mxMatrix( type="Full", nrow=1, ncol=2, free=TRUE, values=0, labels=c("MZ_th1","MZ_th2"), name="threMZ")
threDZ      <- mxMatrix( type="Full", nrow=1, ncol=2, free=TRUE, values=0, labels=c("DZ_th1","DZ_th2"), name="threDZ")

meanG     <- mxMatrix( type="Full", nrow=1, ncol=2, free=FALSE, values=0, labels="mean", name="meanG" )

expMZ     <- mxExpectationNormal( covariance = "expCovMZ", 
                                  means      = "meanG", 
                                  thresholds = "threMZ", 
                                  dimnames   = selVars )
expDZ     <- mxExpectationNormal( covariance = "expCovDZ", 
                                  means      = "meanG", 
                                  thresholds = "threDZ", 
                                  dimnames   = selVars )


#######################################################
####               FINAL PREPARATIONS              ####
#######################################################


fitFun     <- mxFitFunctionWLS("WLS")

modelMZ <- mxModel(threMZ, meanG, expCovMZ, expMZ, dataMZ, fitFun, name = "MZ")
modelDZ <- mxModel(threDZ, meanG, expCovDZ, expDZ, dataDZ, fitFun, name = "DZ")

multi     <- mxFitFunctionMultigroup( c("MZ","DZ") )



#######################################################
####               Run first model                 ####
#######################################################

model1   <- mxModel(meanG, expCovMZ, expCovDZ,  
                    modelMZ, modelDZ, multi, 
                    name="SAT" )

fit1    <- mxRun(model1, intervals = T)
#summary(fit1)


#######################################################
####      Prepare and run constrained models       ####
#######################################################


# Equal threshold across twin order
model2 <- mxModel(name = "SAT_thres1", fit1)
model2 <- omxSetParameters(model2, label=c("MZ_th1","MZ_th2"),values=0, newlabels = "MZ_th" )
model2 <- omxSetParameters(model2, label=c("DZ_th1","DZ_th2"),values=0, newlabels = "DZ_th" )

fit2    <- mxRun(model2, intervals = T)
#summary(fit2)


# Equal threshold across twin type
model3 <- mxModel(name = "SAT_thres2", model2)
model3 <- omxSetParameters(model3, label=c("MZ_th","DZ_th"),values=0, newlabels = "th" )

fit3    <- mxRun(model3, intervals = T)
#summary(fit3)

# Compares models
mxCompare(fit1,list(fit2,fit3))
mxCompare(fit2,fit3)