
## EM for the estimate the covariance matrix

## function to calculate the parameters
par.est<-function(pM, pS){
  sX2<-pS[1,1]
  sM2<-pS[2,2]
  sY2<-pS[3,3]
  sXM<-pS[1,2]
  sXY<-pS[1,3]
  sMY<-pS[2,3]
  mX <-pM[1]
  mM <-pM[2]
  mY <-pM[3]

  pa<-sXM/sX2
  pb<-(sMY*sX2-sXM*sXY)/(sX2*sM2-sXM^2)
  pc<-(sXY*sM2-sXM*sMY)/(sX2*sM2-sXM^2)
  ps1<-sX2
  ps2<-sM2-sXM^2/sX2
  ps3<-(sX2*sM2*sY2-sX2*sMY^2-sM2*sXY^2-sY2*sXM^2+2*sXM*sXY*sMY)/(sX2*sM2-sXM^2)

  pm1<-mX
  pm2<-mM - pa*mX
  pm3<-mY - pb*mM - pc*mX



  return(c(pm2,pm3,pa,pb,pc,ps2,ps3))
}

## the M-step
M_step<-function(sM, sS, N){
## sM: first moment vector
## sS: second moment

  pM <- sM/N

  pS <- sS/(N-1) - N/(N-1)*pM%*%t(pM)
  

  return(list(pM=pM, pS=pS))

}

## the E-step
E_step<-function(dset, pM, pS){

X<-dset[,1]
M<-dset[,2]
Y<-dset[,3]

n<-dim(dset)[1]

  sX2<-rep(0,n)
  sM2<-rep(0,n)
  sY2<-rep(0,n)
  sXM<-rep(0,n)
  sXY<-rep(0,n)
  sMY<-rep(0,n)
  mX <-rep(0,n)
  mM <-rep(0,n)
  mY <-rep(0,n)

for (i in 1:n){

   ## process missing data in X
   if (!is.na(X[i])){
      ## X is available
      if (!is.na(M[i])){
         ## M is available
         if (!is.na(Y[i])){
            ## X, M, Y is available
            mX[i]<-X[i]
            mM[i]<-M[i]
            mY[i]<-Y[i]
            sX2[i]<-X[i]^2
            sY2[i]<-Y[i]^2
            sM2[i]<-M[i]^2
            sXM[i]<-X[i]*M[i]
            sXY[i]<-X[i]*Y[i]
            sMY[i]<-M[i]*Y[i]
            
         }else{
            ## X & M observed but Y is missing
            mX[i]<-X[i]
            mM[i]<-M[i]
            mY[i]<-pM[3] + pS[3,1:2]%*%solve(pS[1:2,1:2])%*%(c(X[i]-pM[1], M[i]-pM[2]))
            sX2[i]<-mX[i]*mX[i]
            sM2[i]<-mM[i]*mM[i]
            sY2[i]<-mY[i]*mY[i]+pS[3,3]-pS[3,1:2]%*%solve(pS[1:2,1:2])%*%(pS[1:2,3])
            sXM[i]<-mX[i]*mM[i]
            sXY[i]<-mX[i]*mY[i]
            sMY[i]<-mM[i]*mY[i]
         }
      }else{
         ## M is missing
         if (!is.na(Y[i])){
            ## X, Y is available but M is missing
            mX[i]<-X[i]
            mM[i]<-pM[2] + pS[2,c(1,3)]%*%solve(pS[c(1,3),c(1,3)])%*%(c(X[i]-pM[1],Y[i]-pM[3]))
            mY[i]<-Y[i]
            sX2[i]<-mX[i]*mX[i]
            sM2[i]<-mM[i]*mM[i]+pS[2,2]-pS[2,c(1,3)]%*%solve(pS[c(1,3),c(1,3)])%*%(pS[c(1,3),2])
            sY2[i]<-mY[i]*mY[i]
            sXM[i]<-mX[i]*mM[i]
            sXY[i]<-mX[i]*mY[i]
            sMY[i]<-mM[i]*mY[i]
            
         }else{
            ## X is observed but M and Y are missing
            mX[i]<-X[i]
            mM[i]<-pM[2] + pS[1,2]/pS[1,1]*(X[i]-pM[1])
            mY[i]<-pM[3] + pS[1,3]/pS[1,1]*(X[i]-pM[1])
            sX2[i]<-mX[i]*mX[i]
            sM2[i]<-mM[i]*mM[i] + pS[2,2]-pS[1,2]^2/pS[1,1]
            sY2[i]<-mY[i]*mY[i] + pS[3,3]-pS[1,3]^2/pS[1,1]
            sXM[i]<-mX[i]*mM[i]
            sXY[i]<-mX[i]*mY[i]
            sMY[i]<-mM[i]*mY[i] + pS[2,3]-pS[1,2]*pS[1,3]/pS[1,1] 
         }
      }
   }else{
      ## X is missing
      if (!is.na(M[i])){
         ## M is available
         if (!is.na(Y[i])){
            ## X is missing, M & Y are available
            mX[i]<-pM[1] + pS[1,2:3]%*%solve(pS[2:3,2:3])%*%(c(M[i]-pM[2], Y[i]-pM[3]))
            mM[i]<-M[i]
            mY[i]<-Y[i]
            sX2[i]<-mX[i]*mX[i]+pS[1,1]-pS[1,2:3]%*%solve(pS[2:3,2:3])%*%pS[2:3,1]
            sM2[i]<-mM[i]*mM[i]
            sY2[i]<-mY[i]*mY[i]
            sXM[i]<-mX[i]*mM[i]
            sXY[i]<-mX[i]*mY[i]
            sMY[i]<-mM[i]*mY[i]
         }else{
            ## Y and X is missing but M is available
            mX[i]<-pM[1]+pS[1,2]/pS[2,2]*(M[i]-pM[2])
            mM[i]<-M[i]
            mY[i]<-pM[3]+pS[2,3]/pS[2,2]*(M[i]-pM[2])
            sX2[i]<-mX[i]*mX[i]+pS[1,1]-pS[1,2]^2/pS[2,2]
            sM2[i]<-mM[i]*mM[i]
            sY2[i]<-mY[i]*mY[i]+pS[3,3]-pS[2,3]^2/pS[2,2]
            sXM[i]<-mX[i]*mM[i]
            sXY[i]<-mX[i]*mY[i]+pS[1,3]-pS[1,2]*pS[2,3]/pS[2,2]
            sMY[i]<-mM[i]*mY[i]            
         }
      }else{
         ## X, M is missing
            mX[i]<-pM[1]+pS[1,3]/pS[3,3]*(Y[i]-pM[3])
            mM[i]<-pM[2]+pS[2,3]/pS[3,3]*(Y[i]-pM[3])
            mY[i]<-Y[i]
            sX2[i]<-mX[i]*mX[i]+pS[1,1]-pS[1,3]^2/pS[3,3]
            sM2[i]<-mM[i]*mM[i]+pS[2,2]-pS[2,3]^2/pS[3,3]
            sY2[i]<-mY[i]*mY[i]
            sXM[i]<-mX[i]*mM[i]+pS[1,2]-pS[1,3]*pS[2,3]/pS[3,3]
            sXY[i]<-mX[i]*mY[i]
            sMY[i]<-mM[i]*mY[i]
      }

   }
}
     sM<-c(sum(mX),sum(mM),sum(mY))
     sS<-array(c( sum(sX2),sum(sXM),sum(sXY),sum(sXM),sum(sM2),sum(sMY),sum(sXY),sum(sMY),sum(sY2)),dim=c(3,3))
     return(list(sM=sM, sS=sS))
}

## Matrice to store the results
R<-100
pairdel<-array(NA,dim=c(R,7))
em<-array(NA,dim=c(R,7))
em2<-array(NA,dim=c(R,7))
em3<-array(NA,dim=c(R,7))
true<-array(NA,dim=c(R,7))
listdel<-array(NA,dim=c(R,7))
ml<-array(NA,dim=c(R,7))

N<-100
for (j in 1:R){
## dset generation

X<-rnorm(N)
M<-.5+.5*X + sqrt(.1)*rnorm(N)
Y<-1+1*M + .1*X + sqrt(.3)*rnorm(N)
dset<-cbind(X,M,Y)

sampM<-c(sum(X),sum(M),sum(Y))
sampS<-array(c(sum(X^2),sum(X*M),sum(X*Y),sum(X*M),sum(M^2),sum(M*Y),sum(X*Y),sum(M*Y),sum(Y^2)),c(3,3))

## results from complete data analysis
para<-M_step(sampM, sampS, N)
true[j,]<-par.est(para$pM, para$pS)

## For the missing data analysis
for (i in 1:N){
#  if (runif(1)<.1) X[i]<-NA
  if (runif(1)<.1) M[i]<-NA
  if (runif(1)<.1) Y[i]<-NA
}

dset<-cbind(X,M,Y)

e<-1

while (e>.0000001){
 E.res<-E_step(dset,para$pM, para$pS)
 M.res<-M_step(E.res$sM, E.res$sS, N)

 e<-sum(abs(M.res$pM-para$pM))+sum(abs(M.res$pS-para$pS))

 para<-M.res
}

em[j,]<-par.est(para$pM,para$pS)

temp1 <- prelim.norm(dset)  
temp2 <- em.norm(temp1)  
temp3 <- getparam.norm(temp1,temp2)
em3[j,]<-par.est(temp3$mu, temp3$sigma)

#test<-prelim.norm(dset)
#res<-em.norm(test)
#getparam.norm(test,res)

## ML method
## save the data
write.table(dset, "data.dat", na='.',row.names=F, col.names=F)
system('c:\\programs\\mplus\\mplus.exe mle.inp',show.output.on.console = F)
tempres<-scan('est.txt',quiet=T)
ml[j,]<-tempres[c(2,1,6,4,5,7,8)]

## the list wise delete method
isna<-is.na(dset)
sumisna<-apply(isna,1,sum)
listdata<-dset[sumisna==0,]
sampN[j]<-dim(listdata)[1]

temp<-cov(listdata,use='complete.obs')
par<-par.est(c(mean(listdata[,1]),mean(listdata[,2]),mean(listdata[,3])),temp)
listdel[j,]<-par

## pairwise delete
temp<-cov(dset,use='pairwise.complete.obs')
par<-par.est(c(mean(dset[,1],na.rm=T),mean(dset[,2],na.rm=T),mean(dset[,3],na.rm=T)),temp)
pairdel[j,]<-par

## EM for path model
temp<-cov(dset,use='pairwise.complete.obs')
par<-est(c(temp[1,1],temp[2,2],temp[3,3],temp[1,2],temp[1,3],temp[2,3],mean(dset[,1],na.rm=T),mean(dset[,2],na.rm=T),mean(dset[,3],na.rm=T)))

e<-1

while (e>.00001){
  SS<-EMe(dset, par)
  para<-est(SS)
  e<-sum(abs(para-par))
  par<-para
 # print(par,digits=10)
}

em2[j,]<-para[c(8,9,1,2,3,5,6)]
}

apply(true,2,mean)
apply(pairdel,2,mean)
apply(em,2,mean)
apply(em2,2,mean)
apply(em3,2,mean)
apply(listdel,2,mean)
apply(ml,2,mean)