model{

# model has ARMA(1,0) structure on the latent level with two states S and one-vs. three factors
# state switch from S=1 to S=2 is predicted by BAI sum score and WAI 3-factors from state S=1
# cross-lagged effects are included and BAI is also DV

# P(S_t=s|S_(t-1)=2), i.e. backswitch probability
PSb[1] <- 0.01
PSb[2] <- 0.99


########################################## level 3 ##########################################
# multivariate normal distribution random effects therapists in S=1 and S=2
for(i in 1:N2){#therapists
  v0.S1[i,1:4] ~ dmnorm(mu0[1:4],psizeta3.S1[1:4,1:4])
  v0.S2[i,1:2] ~ dmnorm(mu0[1:2],psizeta3.S2[1:2,1:2])
}

for(i in 1:N1){#patients

########################################## level 2 ##########################################
# multivariate normal distribution random effects patients in S=1 and S=2
  
  u0.S1[i,1:4] ~ dmnorm(mu0[1:4],psizeta2.S1[1:4,1:4])
  u0.S2[i,1:2] ~ dmnorm(mu0[1:2],psizeta2.S2[1:2,1:2])
  
########################################## level 1 ##########################################

  ##################
  # Markov Switching Model
  ##################
  
  for(t in 2:Nt){
    # logistic model for P(S_t=s|S_(t-1)=1)
    logit(eta2[i,t]) <- b2[1] + 
                        b2[2]*eta1.S1[i,t-1,1]+
                        b2[3]*eta1.S1[i,t-1,2]+
                        b2[4]*eta1.S1[i,t-1,3]+
                        b2[5]*eta1.S1[i,t-1,4]
    
    # transition matrix
    PSa[i,t,1] <- eta2[i,t]
    PSa[i,t,2] <- 1-PSa[i,t,1]
    
    PS1[i,t,1:2] <- ifelse(S[i,t-1]==1,PSa[i,t,1:2],PSb[1:2])
    
    # state membership
    S[i,t] ~ dcat(PS1[i,t,1:2]) 
  }
  
  # constraint for t=1 (all in S_1=1)
  S[i,1] <- 1
  
  ####################################
  # latent factor model ("DSEM" part within each state): ARMA(1,0)
  # In S=1: eta1.S1[,,1:3] are WAI subscales, eta[,,4] is BAI or PSWQ
  # In S=2: eta1.S2[,,1]   is WAI global factor, eta[,,2] is BAI or PSWQ
  ####################################
  
  for(t in 1:Nt){
    eta1.S1[i,t,1:4] ~ dmnorm(mueta1.S1[i,t,1:4],psizeta1.S1[1:4,1:4])
    eta1.S2[i,t,1:2] ~ dmnorm(mueta1.S2[i,t,1:2],psizeta1.S2[1:2,1:2])
  }
  
  ##################
  # means groups S==1,S==2
  # mean is random effect with patient (u) and therapist (v) effect
  # random terms are state-specific 
  ##################
  
  ##################
  # S==1
  ##################
  # t==1
  for(j in 1:4){
    mueta1.S1[i,1,j] <- b0.S1[j] + u0.S1[i,j] + v0.S1[x1[i],j] 
  }
    
  # t==2...T
  for(t in 2:Nt){
    mueta1.S1[i,t,1] <- b0.S1[1] + u0.S1[i,1] + v0.S1[x1[i],1] + phi.S1[1]*eta1.S1[i,t-1,1] + b1.S1[1,1]*eta1.S1[i,t-1,4] 
    mueta1.S1[i,t,2] <- b0.S1[2] + u0.S1[i,2] + v0.S1[x1[i],2] + phi.S1[2]*eta1.S1[i,t-1,2] + b1.S1[2,1]*eta1.S1[i,t-1,4] 
    mueta1.S1[i,t,3] <- b0.S1[3] + u0.S1[i,3] + v0.S1[x1[i],3] + phi.S1[3]*eta1.S1[i,t-1,3] + b1.S1[3,1]*eta1.S1[i,t-1,4] 
    mueta1.S1[i,t,4] <- b0.S1[4] + u0.S1[i,4] + v0.S1[x1[i],4] + phi.S1[4]*eta1.S1[i,t-1,4] + b1.S1[1,2]*eta1.S1[i,t-1,1] + b1.S1[2,2]*eta1.S1[i,t-1,2] + b1.S1[3,2]*eta1.S1[i,t-1,3]
  }
  
  ##################
  # S==2
  ##################
  # t==1
  for(j in 1:2){
    mueta1.S2[i,1,j] <- b0.S2[j] + u0.S2[i,j] + v0.S2[x1[i],j]
  }
    
  # t==2...T
  for(t in 2:Nt){
    mueta1.S2[i,t,1] <- b0.S2[1] + u0.S2[i,1] + v0.S2[x1[i],1] + phi.S2[1]*eta1.S2[i,t-1,1] + b1.S2[1]*eta1.S2[i,t-1,2]
    mueta1.S2[i,t,2] <- b0.S2[2] + u0.S2[i,2] + v0.S2[x1[i],2] + phi.S2[2]*eta1.S2[i,t-1,2] + b1.S2[2]*eta1.S2[i,t-1,1]
  }
  
  
  ##################
  # measurement model: no additional time-structures
  ##################
  # distribution is state-specific (mean structure given the factor distribution/loading pattern, and residual variances/precisions)
  for(t in 1:Nt){
    for(j in 1:12){y1[i,t,j] ~ dnorm(muy1[i,t,j,S[i,t]],psiy1[j,S[i,t]])}
    
    ########################################
    # S==1
    ########################################
    # scaling indicator for each factor
    muy1[i,t,1,1] <- eta1.S1[i,t,1]
    muy1[i,t,4,1] <- eta1.S1[i,t,2]
    muy1[i,t,7,1] <- eta1.S1[i,t,3]
    muy1[i,t,10,1]<- eta1.S1[i,t,4]
    
    # rest 
    for(j in 2:3){  muy1[i,t,j,1] <- ly1.S1[j-1]*eta1.S1[i,t,1]}
    for(j in 5:6){  muy1[i,t,j,1] <- ly1.S1[j-2]*eta1.S1[i,t,2]}
    for(j in 8:9){  muy1[i,t,j,1] <- ly1.S1[j-3]*eta1.S1[i,t,3]}
    for(j in 11:12){muy1[i,t,j,1] <- ly1.S1[j-4]*eta1.S1[i,t,4]}
    
    ########################################
    # S==2
    ########################################
    # scaling indicator for each factor
    muy1[i,t,1,2]  <- eta1.S2[i,t,1]
    muy1[i,t,10,2] <- eta1.S2[i,t,2]
    
    # rest 
    for(j in 2:9){  muy1[i,t,j,2] <- ly1.S2[j-1]*eta1.S2[i,t,1]}
    for(j in 11:12){muy1[i,t,j,2] <- ly1.S2[j-2]*eta1.S2[i,t,2]}
  
  }
  
}# end N1 loop

########################################## Priors ##########################################

# zero mean for random effects
for(j in 1:4){mu0[j] <- 0}

########################################
# S==1
########################################
for(j in 1:4){
  # mean in S=1, constrained because items were rescaled to have zero mean at t=1 
  b0.S1[j]  <- 0
  # AR coefficient
  phi.S1[j] ~ dnorm(0,1)
}

# cross lagged effects
for(j in 1:3){for(k in 1:2){b1.S1[j,k]  ~ dnorm(0,1)}}

########################################
# S==2
########################################
for(j in 1:2){
  # mean in S=2
  b0.S2[j]  ~ dnorm(0,1)
  # AR coefficient
  phi.S2[j] ~ dnorm(0,1)
  # cross lagged effects
  b1.S2[j]  ~ dnorm(0,1)
}

  
########################################
# measurement model
# note: due to rescaling all items have zero means at time 1 and hence zero intercepts
########################################
# factor loadings in S=1 and S=2
for(j in 1:8){ ly1.S1[j] ~ dnorm(1,1)} 
for(j in 1:10){ly1.S2[j] ~ dnorm(1,1)} 


########################################
# precisions measurement model
########################################
for(j in 1:12){for(k in 1:2){psiy1[j,k] ~ dgamma(9,4)}}

########################################
# L1&L2&L3 state specific precisions (wishart)
########################################
psizeta1.S1[1:4,1:4] ~ dwish(PHinv[1:4,1:4],4)
psizeta2.S1[1:4,1:4] ~ dwish(PHinv[1:4,1:4],4)
psizeta3.S1[1:4,1:4] ~ dwish(PHinv[1:4,1:4],4)

psizeta1.S2[1:2,1:2] ~ dwish(PHinv[1:2,1:2],2)
psizeta2.S2[1:2,1:2] ~ dwish(PHinv[1:2,1:2],2)
psizeta3.S2[1:2,1:2] ~ dwish(PHinv[1:2,1:2],2)

########################################
# Markov switching model
########################################
for(j in 1:5){b2[j]  ~ dnorm(0,.5)}



########################################## transformations ##########################################

##############################################
# (co)variances
##############################################
for(k in 1:2){for(j in 1:12){sigmay1[j,k] <- 1/psiy1[j,k]}}

sigmazeta1.S1[1:4,1:4] <- inverse(psizeta1.S1[1:4,1:4])
sigmazeta2.S1[1:4,1:4] <- inverse(psizeta2.S1[1:4,1:4])
sigmazeta3.S1[1:4,1:4] <- inverse(psizeta3.S1[1:4,1:4])

sigmazeta1.S2[1:2,1:2] <- inverse(psizeta1.S2[1:2,1:2])
sigmazeta2.S2[1:2,1:2] <- inverse(psizeta2.S2[1:2,1:2])
sigmazeta3.S2[1:2,1:2] <- inverse(psizeta3.S2[1:2,1:2])

# some descriptives for states (St indicates how often persons were in S=2; Send is the final state)
for(i in 1:N1){
  St[i] <- Nt+1-mean(S[i,1:(Nt)]-1)*Nt
  Send[i] <- S[i,Nt]
}

#########################################################
# loglikelihood for wai/loo (loglik1 is relevant parameter to be extracted in R-loo)
#########################################################
for(i in 1:N1){
  for(t in 1:Nt){
    for(j in 1:12){
      loglik0[i,t,j] = logdensity.norm(y1[i,t,j],muy1[i,t,j,S[i,t]],psiy1[j,S[i,t]]);
    }
    loglik01[i,t] = sum(loglik0[i,t,1:12]);
  }
  loglik1[i] = sum(loglik01[i,1:Nt]);
}
  
loglik = sum(loglik1[1:N1]);

    


###########################################################
# some descriptives for illustration: average trajectories for each state
###########################################################
for(t in 1:Nt){
  # no of persons per class at each t
  nc2[t] <- sum(S[,t]-1)
  nc1[t] <- N1-sum(S[,t]-1)
  
  ########################################
  # ysum for persons in each class at each t
  ########################################
  # S==1
  for(j in 1:4){
    etasum.S1[1,t,j]   <- ifelse(S[1,t]==1,eta1.S1[1,t,j],0)
    etasumsq.S1[1,t,j] <- ifelse(S[1,t]==1,prod(eta1.S1[1,t,j],eta1.S1[1,t,j]),0)
    
    for(i in 2:N1){
      etasum.S1[i,t,j]   <- ifelse(S[i,t]==1,eta1.S1[i,t,j]                     +etasum.S1[i-1,t,j]  ,etasum.S1[i-1,t,j])
      etasumsq.S1[i,t,j] <- ifelse(S[i,t]==1,prod(eta1.S1[i,t,j],eta1.S1[i,t,j])+etasumsq.S1[i-1,t,j],etasumsq.S1[i-1,t,j])
    }
  }
  
  # S==2
  for(j in 1:2){
    etasum.S2[1,t,j]   <- ifelse(S[1,t]==2,eta1.S2[1,t,j],0)
    etasumsq.S2[1,t,j] <- ifelse(S[1,t]==2,prod(eta1.S2[1,t,j],eta1.S2[1,t,j]),0)
    
    for(i in 2:N1){
      etasum.S2[i,t,j]   <- ifelse(S[i,t]==2,eta1.S2[i,t,j]                     +etasum.S2[i-1,t,j]  ,etasum.S2[i-1,t,j])
      etasumsq.S2[i,t,j] <- ifelse(S[i,t]==2,prod(eta1.S2[i,t,j],eta1.S2[i,t,j])+etasumsq.S2[i-1,t,j],etasumsq.S2[i-1,t,j])
    }
  }


}

# inverse of persons per class with a tiny correction to ensure 
# that it is well-defined (otherwise division by zero)
inv.nc1 <- 1/(nc1+0.001)
inv.nc2 <- 1/(nc2+0.001)

for(t in 1:Nt){
  ########################################
  # ymean and yvar for persons in each class at each t
  ########################################
  
  # S==1
  for(j in 1:4){
    etamean.S1[t,j] <- etasum.S1[N1,t,j]*inv.nc1[t]
    etavar.S1[t,j]  <- (etasumsq.S1[N1,t,j]-prod(etasum.S1[N1,t,j],etasum.S1[N1,t,j])*inv.nc1[t])*inv.nc1[t]
  }
  
  # S==2
  for(j in 1:2){
    etamean.S2[t,j] <- etasum.S2[N1,t,j]*inv.nc2[t]
    etavar.S2[t,j]  <- (etasumsq.S2[N1,t,j]-prod(etasum.S2[N1,t,j],etasum.S2[N1,t,j])*inv.nc2[t])*inv.nc2[t]
  }
  
}


###########################################################
# differences in the means between classes at each time point
# posterior distribution can be used to describe zero or non-zero overlap
###########################################################
for(t in 2:Nt){
  delta.eta[t,1] <- etamean.S1[t,1] - etamean.S2[t,1]
  delta.eta[t,2] <- etamean.S1[t,2] - etamean.S2[t,1]
  delta.eta[t,3] <- etamean.S1[t,3] - etamean.S2[t,1]
  delta.eta[t,4] <- etamean.S1[t,4] - etamean.S2[t,2]
}



}