library(irtoys)

###
# Please cite the following paper when using the code
# Wang,c., Chen, P., & Jiang, S. (2019) Item Calibration Methods with Multiple Subscale Multistage Testing.
# Journal of Educational Measurement.

###############Single-group MML/EM ########################################

library(irtoys)

##group is a index vector with the same length as row in repsonse matrix, 
##elements (1,...G) indicate group membership
ipest.sg <- function(resp,Q=21,eps =1e-3,max.iter=50)
{
  N <- nrow(resp)
  J <- ncol(resp)
  
  D <- 1.7
  #One set of quadrature points and weights based on standard normal
  GH <- normal.qu(Q)
  X <- GH$quad.points
  A <- GH$quad.weights
  #starting values
  b <- rep(0,J)
  a <- rep(1,J)
  #intermediate quantities
  ngq <- double(Q)
  rgq <- double(Q)
  pstar <- matrix(,J,Q)
  pi <- double(N)
  niq <- liq <- matrix(double(N*Q),N,Q)
  pjq <-ljq <-  matrix(double(J*Q),J)
  
  df.a <- df.b <-  1
  iter <- 0
  
  while(max(df.a)>eps & max(df.b)>eps  & iter<max.iter)
  {
    iter <- iter+1
    aold <- a
    bold <- b
    
    for (i in 1:N){
      Xg <- X
      Ag <- A
      
      for (j in 1:J){
        if(is.na(resp[i,j]))
        {ljq[j,] <- 1}
        
        else{
          temp.b <- b[j]
          temp.a <- a[j]
          pstar[j,] <- 1/(1+exp(-temp.a*D*(Xg-temp.b)))
          pjq[j,] <-  pstar[j,]
          ljq[j,] <- pjq[j,] ^resp[i,j]*(1-pjq[j,]) ^(1-resp[i,j])}
      }
      
      liq[i,] <- temp.post<- apply(ljq,2,prod)*Ag
      pi[i] <- sum(temp.post)
    } 
    
    niq <- liq/pi
    ngq <- colSums(niq)
    
    for (j in 1:J){
      temp.b <- b[j]
      temp.a <- a[j]
      temp.resp <- resp[,j]
      
      ngq <- colSums(niq[!is.na(temp.resp),])
      riq <- temp.resp*niq
      rgq <- colSums(riq,na.rm = T) 
      
      lik <- function(ip){
        a <- ip[1]
        b <- ip[2]
        p <- 1/(1+exp(-a*D*(X-b)))
        p[p==1] <- 1-1e-7
        p[p==0] <- 1e-7
        -sum(rgq*log(p)+(ngq-rgq)*log(1-p))
      }
      
      gr <- function(ip){
        a <- ip[1]
        b <- ip[2]
        p <- 1/(1+exp(-a*D*(X-b)))
        c(
          -sum((rgq-ngq*p)*D*(X-b)),
          sum((rgq-ngq*p)*a*D)
        )
      }
      temp <- optim(c(temp.a,temp.b),lik,gr,method = "L-BFGS-B", lower =c (0,-Inf))$par
      temp.a <- temp[1]
      temp.b <- temp[2]
      
      
      a[j] <- temp.a
      b[j] <- temp.b
    }   
    
    df.a <- abs(aold-a)
    df.b <- abs(bold-b)
  }
  
  return(list(est=cbind(a,b),iter=iter))
}

#end of function

###############Multiple group MML/EM with truncated normal: M-MML-T ##################

##group is a index vector with the same length as row in repsonse matrix, elements (1,...G) indicate
##group membership, a and b are cut offs, Q # of quadrature, eps convergence criterion,
##and max.iter maximal # of iteration for EM
ipest.mg <- function(resp,group,a=qnorm(1/3),b=qnorm(2/3),Q=21,eps =1e-3,max.iter=25)
{
#sample size and test length
  N <- nrow(resp)
  J <- ncol(resp)

  D <- 1.7
#number of groups, X the quadrature points, A the weights
  G <- length(unique(group))
  A <- X <- matrix(,G,Q)
#three sets of X and A for the three groups, with cut offs provided in function argument
  qu1 <- normal.qu(n=Q,upper =a)
  X[1,] <- qu1$quad.points
  A[1,] <- qu1$quad.weights

  qu2 <- normal.qu(n=Q,upper =b,lower = a)
  X[2,] <- qu2$quad.points
  A[2,] <- qu2$quad.weights

  qu3 <- normal.qu(n=Q,lower = b)
  X[3,] <- qu3$quad.points
  A[3,] <- qu3$quad.weights

#starting values
  b <- rep(0,J)
  a <- rep(1,J)
#some intermediate quantities
  ngq <- matrix(double(G*Q),G)
  rgq <- matrix(double(G*Q),G)
  pstar <- matrix(,J,Q)
  pi <- double(N)
  niq <- liq <- matrix(double(N*Q),N,Q)
  pjq <-ljq <-  matrix(double(J*Q),J)
  # FI <- matrix(double(4),2,2)
  
  df.a <- df.b <- 1
  iter <- 0
#start of EM
  while(max(df.a)>eps & max(df.b)>eps &  iter<max.iter)
  {
    iter <- iter+1
    aold <- a
    bold <- b
#E step, first calculate the expected number of subjects in each quadrature for each group
      for (i in 1:N){
        # temp.J <- sum(!is.na(resp[i,]))
        g <- group[i]
        Xg <- X[g,]
        Ag <- A[g,]
  
        for (j in 1:J){
#ignore missing value when calculating likelihood
           if(is.na(resp[i,j]))
          {ljq[j,] <- 1}
          
          else{
          temp.b <- b[j]
          temp.a <- a[j]
          pstar[j,] <- 1/(1+exp(-temp.a*D*(Xg-temp.b)))
          pjq[j,] <- pstar[j,]
          ljq[j,] <- pjq[j,] ^resp[i,j]*(1-pjq[j,]) ^(1-resp[i,j])}
        }
        
        liq[i,] <- temp.post<- apply(ljq,2,prod)*Ag
        pi[i] <- sum(temp.post)
      } 
    
    niq <- liq/pi
    for (g in 1:G){
      ngq[g,] <- apply(niq[group==g,],2,sum)
    }
#here each row of ngq should sum to N/3
    
    for (j in 1:J){
#continue E step, calculate rgq for each item, for items not in the routing block, only one 
#out of three has nonzero values
      temp.b <- b[j]
      temp.a <- a[j]
      # temp.c <- c[j]
      temp.resp <- resp[,j]

      riq <- temp.resp*niq
      for (g in 1:G){
        rgq[g,] <- apply(riq[group==g,],2,sum,na.rm=T)
      }
#M step 
#likelihood function to be minimized
        lik <- function(ip){
          a <- ip[1]
          b <- ip[2]
          p <- 1/(1+exp(-a*D*(X-b)))
          p[p==1] <- 1-1e-7
          p[p==0] <- 1e-7
          ll <- -sum((rgq*log(p)+(ngq-rgq)*log(1-p))*(rgq>0))
          ll
        }
#gradient function for a and b        
        gr <- function(ip){
          a <- ip[1]
          b <- ip[2]
          p <- 1/(1+exp(-a*D*(X-b)))
          c(
            -sum((rgq-ngq*p)*(X-b)*D*(rgq>0)),
            sum((rgq-ngq*p)*a*D*(rgq>0))
          )
        }
#optimization, retrieve a and b estimate       
        temp <- optim(c(temp.a,temp.b),lik,gr,method = "L-BFGS-B", lower =c (0,-Inf))$par

          temp.a <- temp[1]
          temp.b <- temp[2]

      a[j] <- temp.a
      b[j] <- temp.b
    }   
#calculate the amout of change to compare with termination criterion    
    df.a <- abs(aold-a)
    df.b <- abs(bold-b)
  }
  
  return(list(est=cbind(a,b),iter=iter))
}

#end of function


###############Multiple group MML/EM with normal theta per group: M-MML-N ##################

ipest.mg2 <- function(resp,group,mu=0,sig=1,Q=21,eps =1e-3,max.iter=50)
{
  
  N <- nrow(resp)
  J <- ncol(resp)
  D <- 1.7
  #starting values for the first and third group population parameters
  mu1 <- -1
  mu3 <- 1
  #var
  sig1 <- 1
  sig3 <- 1
  #three sets of X and A for the three groups, to be updated in each cycle of EM
  G <- length(unique(group))
  A <- X <- matrix(,G,Q)
  #starting values  
  b <- rep(0,J)
  a <- rep(1,J)
  #intermediate quantities
  ngq <- matrix(double(G*Q),G)
  rgq <- matrix(double(G*Q),G)
  pstar <- matrix(,J,Q)
  pi <- double(N)
  niq <- liq <- matrix(double(N*Q),N,Q)
  pjq <-ljq <-  matrix(double(J*Q),J)
  
  df.a <- df.b <- 1
  iter <- 0
  #start of EM cycle
  while(max(df.a)>eps & max(df.b)>eps &  iter<max.iter)
  {
    iter <- iter+1
    aold <- a
    bold <- b
    #set up three groups' quadrature and weights, all three groups based on a normal 
    #distribution with population parameters in
    #first and third group estimated, middle group provided in function argument
    qu1 <- normal.qu(n=Q,mu=mu1,sigma = sqrt(sig1),scaling = 1)
    X[1,] <- qu1$quad.points
    A[1,] <- qu1$quad.weights
    
    qu2 <- normal.qu(n=Q,mu=mu,sigma=sig,scaling = 1)
    X[2,] <- qu2$quad.points
    A[2,] <- qu2$quad.weights
    
    qu3 <- normal.qu(n=Q,mu=mu3,sigma = sqrt(sig3),scaling = 1)
    X[3,] <- qu3$quad.points
    A[3,] <- qu3$quad.weights
    #calculate ngq, expected number of subjects in each quadrature for each group    
    for (i in 1:N){
      #read in group membership and corresponding quadrature and weights
      g <- group[i]
      Xg <- X[g,]
      Ag <- A[g,]
      
      for (j in 1:J){
        if(is.na(resp[i,j]))
        {ljq[j,] <- 1}
        
        else{
          temp.b <- b[j]
          temp.a <- a[j]
          pstar[j,] <- 1/(1+exp(-temp.a*D*(Xg-temp.b)))
          pjq[j,] <- pstar[j,]
          ljq[j,] <- pjq[j,] ^resp[i,j]*(1-pjq[j,]) ^(1-resp[i,j])}
      }
      liq[i,] <- temp.post<- apply(ljq,2,prod)*Ag
      pi[i] <- sum(temp.post)
    } 
    
    niq <- liq/pi
    for (g in 1:G){
      ngq[g,] <- apply(niq[group==g,],2,sum)
    }
    #continue E-step, calculate how many people got item j correct in each quadrature for 
    #each group
    for (j in 1:J){
      temp.b <- b[j]
      temp.a <- a[j]
      temp.resp <- resp[,j]
      
      riq <- temp.resp*niq
      for (g in 1:G){
        rgq[g,] <- apply(riq[group==g,],2,sum,na.rm=T)
      }
      
      lik <- function(ip){
        a <- ip[1]
        b <- ip[2]
        p <- 1/(1+exp(-a*D*(X-b)))
        p[p==1] <- 1-1e-7
        p[p==0] <- 1e-7
        -sum((rgq*log(p)+(ngq-rgq)*log(1-p))*(rgq>0))
      }
      
      gr <- function(ip){
        a <- ip[1]
        b <- ip[2]
        p <- 1/(1+exp(-a*D*(X-b)))
        c(
          -sum((rgq-ngq*p)*(X-b)*D*(rgq>0)),
          sum((rgq-ngq*p)*a*D*(rgq>0))
        )
      }
      
      temp <- optim(c(temp.a,temp.b),lik,gr,method = "L-BFGS-B", lower =c (0,-Inf))$par
      
      temp.a <- temp[1]
      temp.b <- temp[2]
      
      a[j] <- temp.a
      b[j] <- temp.b
    }   
    #after all item parameters are updated, calculate population parameters    
    Ng <- table(group)
    Xiq <- X[rep(1:nrow(X),times=Ng),]
    muiq <- Xiq*niq
    Xi <- apply(muiq,1,sum)
    
    mu1 <- mean(Xi[group==1])
    mu3 <- mean(Xi[group==3])
    
    sigiq <- apply((Xiq-Xi)^2*niq,1,sum)
    sig1 <- mean(sigiq[group==1]+(Xi[group==1]-mu1)^2)
    sig3 <- mean(sigiq[group==3]+(Xi[group==3]-mu3)^2)
    
    df.a <- abs(aold-a)
    df.b <- abs(bold-b)
  }
  
  return(list(est=cbind(a,b),pop=c(mu1,mu3,sig1,sig3),iter=iter))
}

#end of function