#########################################################
### R programs for simulation study in Section 5 of the paper
### Defining and Estimating Intervention Effects for Groups That Will
### Develop an Auxiliary Outcome
### by Marshall Joffe, Dylan Small and Chi-yuan Hsu
#########################################################

expit=function(x){
  exp(x)/(1+exp(x));
}

#####################################
# Estimation for a Model that includes a binary covariate
# Set covariate X to have a constant effect over all A levels and
# principal strata
#####################################

# Setting I(A)
# Simulation study
sims=1; # Number of simulations
# Vectors to store estimates from simulation study
mean.immune.a1.est.1=rep(0,sims);
mean.immune.a0.est.1=rep(0,sims);
mean.treatprotected.a1.est.1=rep(0,sims);
mean.treatprotected.a0.est.1=rep(0,sims);
mean.treatharmful.a1.est.1=rep(0,sims);
mean.treatharmful.a0.est.1=rep(0,sims);
mean.doomed.a1.est.1=rep(0,sims);
mean.doomed.a0.est.1=rep(0,sims);
sd.immune.a1.est.1=rep(0,sims);
sd.immune.a0.est.1=rep(0,sims);
sd.treatprotected.a1.est.1=rep(0,sims);
sd.treatprotected.a0.est.1=rep(0,sims);
sd.treatharmful.a1.est.1=rep(0,sims);
sd.treatharmful.a0.est.1=rep(0,sims);
sd.doomed.a1.est.1=rep(0,sims);
sd.doomed.a0.est.1=rep(0,sims);
prob.immune.est.1=rep(0,sims);
prob.treatprotected.est.1=rep(0,sims);
prob.treatharmful.est.1=rep(0,sims);
prob.doomed.est.1=rep(0,sims);
xeffect.est.1=rep(0,sims);
psi0.est=rep(0,sims); # Structural nested model coefficients
psi1.est=rep(0,sims); 
betahat.a=rep(0,sims);
betahat.mua=rep(0,sims);
lldistvec.1=rep(0,sims);
no.notconv=0; # Number of times EM algorithm does not converge
itt.immune.treatprotected.est.2=rep(0,sims);
mean.immune.a0.est.2=rep(0,sims);
mean.treatprotected.a0.est.2=rep(0,sims);
itt.treatharmful.doomed.est.2=rep(0,sims);
mean.treatharmful.a0.est.2=rep(0,sims);
mean.doomed.a0.est.2=rep(0,sims);
sd.immune.a1.est.2=rep(0,sims);
sd.immune.a0.est.2=rep(0,sims);
sd.treatprotected.a1.est.2=rep(0,sims);
sd.treatprotected.a0.est.2=rep(0,sims);
sd.treatharmful.a1.est.2=rep(0,sims);
sd.treatharmful.a0.est.2=rep(0,sims);
sd.doomed.a1.est.2=rep(0,sims);
sd.doomed.a0.est.2=rep(0,sims);
xeffect.est.2=rep(0,sims);
lldistvec.2=rep(0,sims);

k=1;
while(k<=sims){
# Simulation from model with a binary covariate
# Probabilities of principal strata
prob.immune=.25;
prob.treatprotected=.4;
prob.treatharmful=.05;
prob.doomed=.3;

# Means and standard deviations of each principal strata under aggressive treatment
# for hypertension and under not aggressive treatment for hypertension
mean.immune.a1=2; # Mean of immune subjects when aggressively treated for hypertension
sd.immune.a1=1; # Variance of immune subjects when aggressively treated
mean.immune.a0=1; # Mean of immune subjects when not agressively treated
sd.immune.a0=1; # Variance of immune subjects when not aggressively treated
mean.treatprotected.a1=2.5;
sd.treatprotected.a1=1;
mean.treatprotected.a0=1.5;
sd.treatprotected.a0=1;
mean.treatharmful.a1=1.25;
sd.treatharmful.a1=1;
mean.treatharmful.a0=.75;
sd.treatharmful.a0=1;
mean.doomed.a1=1.75;
sd.doomed.a1=1;
mean.doomed.a0=1.25;
sd.doomed.a0=.25;
xeffect=.5; # Addition to mean for binary effect

# Probabilities of x for different principal strata
probx1.immune=.5;
probx1.treatprotected=.75;
probx1.treatharmful=.25;
probx1.doomed=.5;
totalprobx=probx1.immune*prob.immune+probx1.treatprotected*prob.treatprotected+probx1.treatharmful*prob.treatharmful+probx1.doomed*prob.doomed;

n=500; # total sample size
# Assume that half of subjects are randomized into aggressive treatment
# of hypertension and half into nonaggressive treatment
aggressive.treatment=c(rep(1,n/2),rep(0,n/2));
principalstrata=rmultinom(n,1,c(prob.immune,prob.treatprotected,prob.treatharmful,prob.doomed));
# Dummy variables for principal strata
principalstrata.immune=principalstrata[1,];
principalstrata.treatprotected=principalstrata[2,];
principalstrata.treatharmful=principalstrata[3,];
principalstrata.doomed=principalstrata[4,];
# Covariate X
x=principalstrata.immune*rbinom(n,1,probx1.immune)+principalstrata.treatprotected*rbinom(n,1,probx1.treatprotected)+principalstrata.treatharmful*rbinom(n,1,probx1.treatharmful)+principalstrata.doomed*rbinom(n,1,probx1.doomed);

# ESRD dummy
esrd=principalstrata.doomed+principalstrata.treatprotected*(1-aggressive.treatment)+principalstrata.treatharmful*aggressive.treatment;
# Mean and standard deviation of outcome
mean.outcome=mean.immune.a1*principalstrata.immune*aggressive.treatment+mean.immune.a0*principalstrata.immune*(1-aggressive.treatment)+mean.treatprotected.a1*principalstrata.treatprotected*aggressive.treatment+mean.treatprotected.a0*principalstrata.treatprotected*(1-aggressive.treatment)+mean.treatharmful.a1*principalstrata.treatharmful*aggressive.treatment+mean.treatharmful.a0*principalstrata.treatharmful*(1-aggressive.treatment)+mean.doomed.a1*principalstrata.doomed*aggressive.treatment+mean.doomed.a0*principalstrata.doomed*(1-aggressive.treatment)+xeffect*x;

sd.outcome=sd.immune.a1*principalstrata.immune*aggressive.treatment+sd.immune.a0*principalstrata.immune*(1-aggressive.treatment)+sd.treatprotected.a1*principalstrata.treatprotected*aggressive.treatment+sd.treatprotected.a0*principalstrata.treatprotected*(1-aggressive.treatment)+sd.treatharmful.a1*principalstrata.treatharmful*aggressive.treatment+sd.treatharmful.a0*principalstrata.treatharmful*(1-aggressive.treatment)+sd.doomed.a1*principalstrata.doomed*aggressive.treatment+sd.doomed.a0*principalstrata.doomed*(1-aggressive.treatment);

# Generate outcomes from normal distribution
y=rnorm(n,mean=mean.outcome,sd=sd.outcome);
# Generate outcomes from gamma distribution
# scalegamma=sd.outcome^2/mean.outcome;
# shapegamma=mean.outcome/scalegamma;
# y=rgamma(n,shape=shapegamma,scale=scalegamma);

# EM algorithm

# Select starting values
# For now, start at truth
prob.immune.curr=prob.immune;
prob.treatprotected.curr=prob.treatprotected;
prob.treatharmful.curr=prob.treatharmful;
prob.doomed.curr=prob.doomed;
prob.immune.x1.curr=(prob.immune*probx1.immune)/totalprobx;
prob.treatprotected.x1.curr=(prob.treatprotected*probx1.treatprotected)/totalprobx;
prob.treatharmful.x1.curr=(prob.treatharmful*probx1.treatharmful)/totalprobx;
prob.doomed.x1.curr=(prob.doomed*probx1.doomed)/totalprobx;
prob.immune.x0.curr=(prob.immune*(1-probx1.immune))/(1-totalprobx);
prob.treatprotected.x0.curr=(prob.treatprotected*(1-probx1.treatprotected))/(1-totalprobx);
prob.treatharmful.x0.curr=(prob.treatharmful*(1-probx1.treatharmful))/(1-totalprobx);
prob.doomed.x0.curr=(prob.doomed*(1-probx1.doomed))/(1-totalprobx);
mean.immune.a1.curr=mean.immune.a1;
sd.immune.a1.curr=sd.immune.a1; 
mean.immune.a0.curr=mean.immune.a0; 
sd.immune.a0.curr=sd.immune.a0;
mean.treatprotected.a1.curr=mean.treatprotected.a1;
sd.treatprotected.a1.curr=sd.treatprotected.a1;
mean.treatprotected.a0.curr=mean.treatprotected.a0;
sd.treatprotected.a0.curr=sd.treatprotected.a0;
mean.treatharmful.a1.curr=mean.treatharmful.a1;
sd.treatharmful.a1.curr=sd.treatharmful.a1;
mean.treatharmful.a0.curr=mean.treatharmful.a0;
sd.treatharmful.a0.curr=sd.treatharmful.a0;
mean.doomed.a1.curr=mean.doomed.a1;
sd.doomed.a1.curr=sd.doomed.a1;
mean.doomed.a0.curr=mean.doomed.a0;
sd.doomed.a0.curr=sd.doomed.a0;
xeffect.curr=xeffect;

maxiters=500; # Maximum iterations in EM algorithm
iter=1;
lldist=1; # Will measure difference in log likelihood between iterations
toler=.001; # Tolerance for difference in log likelihood between iterations
# for stopping EM algorithm
conv=1;

while(lldist>toler && iter<maxiters){
  # Compute log likelihood with current parameters
ll.old.p1a=sum((1-aggressive.treatment)*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*esrd*x)*log(prob.doomed.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*esrd*x)*log(prob.doomed.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x1.curr)/(prob.immune.x1.curr+prob.treatharmful.x1.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+xeffect.curr,sd=sd.treatharmful.a0.curr))));
  ll.old.p1b=sum((1-aggressive.treatment)*esrd*x*log(((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+xeffect.curr,sd=sd.treatprotected.a0.curr))));
ll.old.p2=sum(aggressive.treatment*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.immune.a1.curr+xeffect.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.immune.x1.curr)*dnorm(y,mean=mean.treatprotected.a1.curr+xeffect.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*x*((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.doomed.a1.curr+xeffect.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x1.curr)/(prob.treatharmful.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatharmful.a1.curr+xeffect.curr,sd=sd.treatharmful.a1.curr)));
ll.old.p3a=sum((1-aggressive.treatment)*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x0.curr)/(prob.immune.x0.curr+prob.treatharmful.x0.curr)*dnorm(y,mean=mean.treatharmful.a0.curr,sd=sd.treatharmful.a0.curr))))
ll.old.p3b=sum((1-aggressive.treatment)*esrd*(1-x)*log(((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatprotected.a0.curr,sd=sd.treatprotected.a0.curr))));
ll.old.p4=sum(aggressive.treatment*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.immune.a1.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.immune.x0.curr)*dnorm(y,mean=mean.treatprotected.a1.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*(1-x)*((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.doomed.a1.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x0.curr)/(prob.treatharmful.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatharmful.a1.curr,sd=sd.treatharmful.a1.curr)));
ll.old=ll.old.p1a+ll.old.p1b+ll.old.p2+ll.old.p3a+ll.old.p3b+ll.old.p4;

# Expectation step
# Compute probabilities for each subject's principal strata
prob.immune.subject=rep(0,n);
prob.treatprotected.subject=rep(0,n);
prob.treatharmful.subject=rep(0,n);
prob.doomed.subject=rep(0,n);
for(j in 1:n){
  if(aggressive.treatment[j]==0 && esrd[j]==0 && x[j]==1){
    prob.immune.subject[j]=prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)/(prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)+prob.treatharmful.x1.curr*dnorm(y[j],mean=mean.treatharmful.a0.curr+xeffect.curr,sd=sd.treatharmful.a0.curr));
    prob.treatharmful.subject[j]=1-prob.immune.subject[j];
  }
  if(aggressive.treatment[j]==0 && esrd[j]==0 && x[j]==0){
  prob.immune.subject[j]=prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)/(prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)+prob.treatharmful.x0.curr*dnorm(y[j],mean=mean.treatharmful.a0.curr,sd=sd.treatharmful.a0.curr));
    prob.treatharmful.subject[j]=1-prob.immune.subject[j];
  }  
  if(aggressive.treatment[j]==0 && esrd[j]==1 && x[j]==1){
    prob.doomed.subject[j]=prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)/(prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)+prob.treatprotected.x1.curr*dnorm(y[j],mean=mean.treatprotected.a0.curr+xeffect.curr,sd=sd.treatprotected.a0.curr));
    prob.treatprotected.subject[j]=1-prob.doomed.subject[j];
  }
  if(aggressive.treatment[j]==0 && esrd[j]==1 && x[j]==0){
    prob.doomed.subject[j]=prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)/(prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)+prob.treatprotected.x0.curr*dnorm(y[j],mean=mean.treatprotected.a0.curr,sd=sd.treatprotected.a0.curr));
    prob.treatprotected.subject[j]=1-prob.doomed.subject[j];
  }
  if(aggressive.treatment[j]==1 && esrd[j]==0 && x[j]==1){
    prob.immune.subject[j]=prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a1.curr+xeffect.curr,sd=sd.immune.a1.curr)/(prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a1.curr+xeffect.curr,sd=sd.immune.a1.curr)+prob.treatprotected.x1.curr*dnorm(y[j],mean=mean.treatprotected.a1.curr+xeffect.curr,sd=sd.treatprotected.a1.curr));
    prob.treatprotected.subject[j]=1-prob.immune.subject[j];
  }
  if(aggressive.treatment[j]==1 && esrd[j]==0 && x[j]==0){
    prob.immune.subject[j]=prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a1.curr,sd=sd.immune.a1.curr)/(prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a1.curr,sd=sd.immune.a1.curr)+prob.treatprotected.x0.curr*dnorm(y[j],mean=mean.treatprotected.a1.curr,sd=sd.treatprotected.a1.curr));
    prob.treatprotected.subject[j]=1-prob.immune.subject[j];
  }  
  if(aggressive.treatment[j]==1 && esrd[j]==1 && x[j]==1){
    prob.doomed.subject[j]=prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a1.curr+xeffect.curr,sd=sd.doomed.a1.curr)/(prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a1.curr+xeffect.curr,sd=sd.doomed.a1.curr)+prob.treatharmful.x1.curr*dnorm(y[j],mean=mean.treatharmful.a1.curr+xeffect.curr,sd=sd.treatharmful.a1.curr));
    prob.treatharmful.subject[j]=1-prob.doomed.subject[j];
  }
  if(aggressive.treatment[j]==1 && esrd[j]==1 && x[j]==0){
    prob.doomed.subject[j]=prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a1.curr,sd=sd.doomed.a1.curr)/(prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a1.curr,sd=sd.doomed.a1.curr)+prob.treatharmful.x0.curr*dnorm(y[j],mean=mean.treatharmful.a1.curr,sd=sd.treatharmful.a1.curr));
    prob.treatharmful.subject[j]=1-prob.doomed.subject[j];
  }  
}
 
# Maximization Step (Use Conditional Maximization, i.e., ECM algorithm)
# First maximize expected likelihood over xeffect, principal strata means
# and then maximize likelihood over standard deviations in each principal
# stratra
# Set up weighted regression
weighty=rep(0,2*n);
weightx=rep(0,2*n);
weight.immune.a1=rep(0,2*n);
weight.immune.a0=rep(0,2*n);
weight.treatprotected.a1=rep(0,2*n);
weight.treatprotected.a0=rep(0,2*n);
weight.treatharmful.a1=rep(0,2*n);
weight.treatharmful.a0=rep(0,2*n);
weight.doomed.a1=rep(0,2*n);
weight.doomed.a0=rep(0,2*n);

for(j in 1:n){
  if(aggressive.treatment[j]==0 && esrd[j]==0){
  tempweight1=(prob.immune.subject[j]/(prob.immune.subject[j]+prob.treatharmful.subject[j]))^.5/sd.immune.a0.curr;
  weighty[((j-1)*2+1)]=y[j]*tempweight1;
  weightx[((j-1)*2+1)]=x[j]*tempweight1;
  weight.immune.a0[((j-1)*2+1)]=tempweight1;
  tempweight2=(prob.treatharmful.subject[j]/(prob.immune.subject[j]+prob.treatharmful.subject[j]))^.5/sd.treatharmful.a0.curr;
  weighty[(j*2)]=y[j]*tempweight2;
  weightx[(j*2)]=x[j]*tempweight2;
  weight.treatharmful.a0[(j*2)]=tempweight2;
}
  if(aggressive.treatment[j]==0 && esrd[j]==1){
  tempweight1=(prob.doomed.subject[j]/(prob.doomed.subject[j]+prob.treatprotected.subject[j]))^.5/sd.doomed.a0.curr;
  weighty[((j-1)*2+1)]=y[j]*tempweight1;
  weightx[((j-1)*2+1)]=x[j]*tempweight1;
  weight.doomed.a0[((j-1)*2+1)]=tempweight1;
  tempweight2=(prob.treatprotected.subject[j]/(prob.doomed.subject[j]+prob.treatprotected.subject[j]))^.5/sd.treatprotected.a0.curr;
  weighty[(j*2)]=y[j]*tempweight2;
  weightx[(j*2)]=x[j]*tempweight2;
  weight.treatprotected.a0[(j*2)]=tempweight2;
}
  if(aggressive.treatment[j]==1 && esrd[j]==0){
    tempweight1=(prob.immune.subject[j]/(prob.immune.subject[j]+prob.treatprotected.subject[j]))^.5/sd.immune.a1.curr;
    weighty[((j-1)*2+1)]=y[j]*tempweight1;
    weightx[((j-1)*2+1)]=x[j]*tempweight1;
    weight.immune.a1[((j-1)*2+1)]=tempweight1;
    tempweight2=(prob.treatprotected.subject[j]/(prob.immune.subject[j]+prob.treatprotected.subject[j]))^.5/sd.treatprotected.a1.curr;
    weighty[(j*2)]=y[j]*tempweight2;
    weightx[(j*2)]=x[j]*tempweight2;
    weight.treatprotected.a1[(j*2)]=tempweight2;
  }
  if(aggressive.treatment[j]==1 && esrd[j]==1){
    tempweight1=(prob.doomed.subject[j]/(prob.doomed.subject[j]+prob.treatharmful.subject[j]))^.5/sd.doomed.a1.curr;
    weighty[((j-1)*2+1)]=y[j]*tempweight1;
    weightx[((j-1)*2+1)]=x[j]*tempweight1;
    weight.doomed.a1[((j-1)*2+1)]=tempweight1;
    tempweight2=(prob.treatharmful.subject[j]/(prob.treatharmful.subject[j]+prob.doomed.subject[j]))^.5/sd.treatharmful.a1.curr;
    weighty[(j*2)]=y[j]*tempweight2;
    weightx[(j*2)]=x[j]*tempweight2;
    weight.treatharmful.a1[(j*2)]=tempweight2;
  }
}
  
weightreg=lm(weighty~weightx+weight.immune.a0+weight.immune.a1+weight.treatharmful.a0+weight.treatharmful.a1+weight.doomed.a0+weight.doomed.a1+weight.treatprotected.a0+weight.treatprotected.a1-1);
coef.weightreg=coef(weightreg);
  xeffect.curr=coef.weightreg[1];
  mean.immune.a0.curr=coef.weightreg[2];
  mean.immune.a1.curr=coef.weightreg[3];
  mean.treatharmful.a0.curr=coef.weightreg[4];
  mean.treatharmful.a1.curr=coef.weightreg[5];
  mean.doomed.a0.curr=coef.weightreg[6];
  mean.doomed.a1.curr=coef.weightreg[7];
  mean.treatprotected.a0.curr=coef.weightreg[8];
  mean.treatprotected.a1.curr=coef.weightreg[9];

sd.immune.a1.curr=sqrt(sum(prob.immune.subject*aggressive.treatment*(y-mean.immune.a1.curr-x*xeffect.curr)^2)/sum(prob.immune.subject*aggressive.treatment))
sd.immune.a0.curr=sqrt(sum(prob.immune.subject*(1-aggressive.treatment)*(y-mean.immune.a0.curr-x*xeffect.curr)^2)/sum(prob.immune.subject*(1-aggressive.treatment)));
sd.treatprotected.a1.curr=sqrt(sum(prob.treatprotected.subject*aggressive.treatment*(y-mean.treatprotected.a1.curr-x*xeffect.curr)^2)/sum(prob.treatprotected.subject*aggressive.treatment));
sd.treatprotected.a0.curr=sqrt(sum(prob.treatprotected.subject*(1-aggressive.treatment)*(y-mean.treatprotected.a0.curr-x*xeffect.curr)^2)/sum(prob.treatprotected.subject*(1-aggressive.treatment)));
sd.doomed.a1.curr=sqrt(sum(prob.doomed.subject*aggressive.treatment*(y-mean.doomed.a1.curr-x*xeffect.curr)^2)/sum(prob.doomed.subject*aggressive.treatment));
sd.doomed.a0.curr=sqrt(sum(prob.doomed.subject*(1-aggressive.treatment)*(y-mean.doomed.a0.curr-x*xeffect.curr)^2)/sum(prob.doomed.subject*(1-aggressive.treatment)));
sd.treatharmful.a1.curr=sqrt(sum(prob.treatharmful.subject*aggressive.treatment*(y-mean.treatharmful.a1.curr-x*xeffect.curr)^2)/sum(prob.treatharmful.subject*aggressive.treatment));
sd.treatharmful.a0.curr=sqrt(sum(prob.treatharmful.subject*(1-aggressive.treatment)*(y-mean.treatharmful.a0.curr-x*xeffect.curr)^2)/sum(prob.treatharmful.subject*(1-aggressive.treatment)));
  
prob.immune.x1.curr=sum(prob.immune.subject*x)/sum(x);
prob.treatprotected.x1.curr=sum(prob.treatprotected.subject*x)/sum(x);
prob.treatharmful.x1.curr=sum(prob.treatharmful.subject*x)/sum(x);
prob.doomed.x1.curr=sum(prob.doomed.subject*x)/sum(x);
prob.immune.x0.curr=sum(prob.immune.subject*(1-x))/(n-sum(x));
prob.treatprotected.x0.curr=sum(prob.treatprotected.subject*(1-x))/(n-sum(x));
prob.treatharmful.x0.curr=sum(prob.treatharmful.subject*(1-x))/(n-sum(x));
prob.doomed.x0.curr=sum(prob.doomed.subject*(1-x))/(n-sum(x));

# Compute log likelihood with updated parameters
  ll.new.p1a=sum((1-aggressive.treatment)*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*esrd*x)*log(prob.doomed.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*esrd*x)*log(prob.doomed.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x1.curr)/(prob.immune.x1.curr+prob.treatharmful.x1.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+xeffect.curr,sd=sd.treatharmful.a0.curr))));
  ll.new.p1b=sum((1-aggressive.treatment)*esrd*x*log(((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+xeffect.curr,sd=sd.treatprotected.a0.curr))));
ll.new.p2=sum(aggressive.treatment*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.immune.a1.curr+xeffect.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.immune.x1.curr)*dnorm(y,mean=mean.treatprotected.a1.curr+xeffect.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*x*((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.doomed.a1.curr+xeffect.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x1.curr)/(prob.treatharmful.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatharmful.a1.curr+xeffect.curr,sd=sd.treatharmful.a1.curr)));
ll.new.p3a=sum((1-aggressive.treatment)*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x0.curr)/(prob.immune.x0.curr+prob.treatharmful.x0.curr)*dnorm(y,mean=mean.treatharmful.a0.curr,sd=sd.treatharmful.a0.curr))))
ll.new.p3b=sum((1-aggressive.treatment)*esrd*(1-x)*log(((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatprotected.a0.curr,sd=sd.treatprotected.a0.curr))));
ll.new.p4=sum(aggressive.treatment*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.immune.a1.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.immune.x0.curr)*dnorm(y,mean=mean.treatprotected.a1.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*(1-x)*((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.doomed.a1.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x0.curr)/(prob.treatharmful.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatharmful.a1.curr,sd=sd.treatharmful.a1.curr)));
ll.new=ll.new.p1a+ll.new.p1b+ll.new.p2+ll.new.p3a+ll.new.p3b+ll.new.p4;

# If log likelihood has become infinite, mark that estimate has not converged
# and restart simulation
# If log likelihood has become infinite, mark that estimate has not converged
# and restart simulation
if(is.na(ll.new)==FALSE){
# If any of the standard deviations has become zero, mark that estimate
# has not converged and restart simulation
if(sd.immune.a1.curr<.00001 | sd.immune.a0.curr<.00001 | sd.treatprotected.a1.curr<.00001 | sd.treatprotected.a0.curr<.00001 | sd.doomed.a1.curr<.00001 | sd.doomed.a0.curr<.00001 | sd.treatharmful.a1.curr<.00001 | sd.treatharmful.a0.curr<.00001){
  no.notconv=no.notconv+1
  conv=0;
  
}
if(sd.immune.a1.curr>=.00001 & sd.immune.a0.curr>=.00001 & sd.treatprotected.a1.curr>=.00001 & sd.treatprotected.a0.curr>=.00001 & sd.doomed.a1.curr>=.00001 & sd.doomed.a0.curr>=.00001 & sd.treatharmful.a1.curr>=.00001 & sd.treatharmful.a0.curr>=.00001){
lldist=ll.new-ll.old;
iter=iter+1;
}
}

if(is.na(ll.new)==TRUE){
  no.notconv=no.notconv+1;
  conv=0;
  lldist=0;
}
}

mean.immune.a1.est.1[k]=mean.immune.a1.curr;
mean.immune.a0.est.1[k]=mean.immune.a0.curr;
mean.treatprotected.a1.est.1[k]=mean.treatprotected.a1.curr;
mean.treatprotected.a0.est.1[k]=mean.treatprotected.a0.curr;
mean.treatharmful.a1.est.1[k]=mean.treatharmful.a1.curr;
mean.treatharmful.a0.est.1[k]=mean.treatharmful.a0.curr;
mean.doomed.a1.est.1[k]=mean.doomed.a1.curr;
mean.doomed.a0.est.1[k]=mean.doomed.a0.curr;
sd.immune.a1.est.1[k]=sd.immune.a1.curr;
sd.immune.a0.est.1[k]=sd.immune.a0.curr;
sd.treatprotected.a1.est.1[k]=sd.treatprotected.a1.curr;
sd.treatprotected.a0.est.1[k]=sd.treatprotected.a0.curr;
sd.treatharmful.a1.est.1[k]=sd.treatharmful.a1.curr;
sd.treatharmful.a0.est.1[k]=sd.treatharmful.a0.curr;
sd.doomed.a1.est.1[k]=sd.doomed.a1.curr;
sd.doomed.a0.est.1[k]=sd.doomed.a0.curr;
prob.immune.est.1[k]=prob.immune.curr;
prob.treatprotected.est.1[k]=prob.treatprotected.curr;
prob.treatharmful.est.1[k]=prob.treatharmful.curr;
prob.doomed.est.1[k]=prob.doomed.curr;
xeffect.est.1[k]=xeffect.curr;
lldistvec.1[k]=lldist;

# Expected auxiliary stratification estimates
ax=aggressive.treatment*x;
auxreg1=glm(esrd~x+aggressive.treatment+ax,family=binomial);
# muhatax=predict(auxreg1,type="response");
muhatax=predict.glm(auxreg1,type="response");
a.times.muhatax=aggressive.treatment*muhatax;
auxreg2=lm(y~muhatax+aggressive.treatment+a.times.muhatax);
betahat.a[k]=coef(auxreg2)[3];
betahat.mua[k]=coef(auxreg2)[4];

# Structural nested model estimates
# Use Newton's Method
# Start at true values
psi1.curr=mean.doomed.a1-mean.doomed.a0;
psi0.curr=mean.treatprotected.a1-mean.treatprotected.a0;

# Estimate "mediating score"
tempreg=glm(esrd~x,family=binomial,subset=(aggressive.treatment==1));
mediatingscore=expit(coef(tempreg)[1]+coef(tempreg)[2]*x);

maxiters=500; # Maximum iterations in Newton-Raphson algorithm
iter=1;
dist=1;
toler=.001;
while(iter<maxiters && dist>toler){
  psiold=matrix(c(psi0.curr,psi1.curr),ncol=1);
  y0psi=y-psi0.curr*aggressive.treatment*(1-esrd)-psi1.curr*aggressive.treatment*esrd;
  tempreg=lm(y0psi~x);
  tempbetahat0=coef(tempreg)[1];
  tempbetahat1=coef(tempreg)[2];
  g1=sum((aggressive.treatment-.5)*(y0psi-tempbetahat0-tempbetahat1*x)*(1-mediatingscore));
  g2=sum((aggressive.treatment-.5)*(y0psi-tempbetahat0-tempbetahat1*x)*mediatingscore);
  Jmat=matrix(rep(0,4),ncol=2);
  Jmat[1,1]=sum(-(aggressive.treatment-.5)*aggressive.treatment*(1-esrd)*(1-mediatingscore));
  Jmat[1,2]=sum(-(aggressive.treatment-.5)*aggressive.treatment*esrd*(1-mediatingscore));
  Jmat[2,1]=sum(-(aggressive.treatment-.5)*aggressive.treatment*(1-esrd)*mediatingscore);
  Jmat[2,2]=sum(-(aggressive.treatment-.5)*aggressive.treatment*esrd*mediatingscore);
  psinew=psiold-solve(Jmat)%*%matrix(c(g1,g2),ncol=1);
  dist=(sum((psinew-psiold)^2))^.5;
  psi0.curr=psinew[1];
  psi1.curr=psinew[2];
  iter=iter+1;
}
psi0.est[k]=psi0.curr;
psi1.est[k]=psi1.curr;


# Estimates that assume ITT(immune)=ITT(treatment protected)
#                  ITT(treatment harmful)=ITT(doomed)
# EM algorithm

# Select starting values
# For now, start at truth
prob.immune.curr=prob.immune;
prob.treatprotected.curr=prob.treatprotected;
prob.treatharmful.curr=prob.treatharmful;
prob.doomed.curr=prob.doomed;
prob.immune.x1.curr=(prob.immune*probx1.immune)/totalprobx;
prob.treatprotected.x1.curr=(prob.treatprotected*probx1.treatprotected)/totalprobx;
prob.treatharmful.x1.curr=(prob.treatharmful*probx1.treatharmful)/totalprobx;
prob.doomed.x1.curr=(prob.doomed*probx1.doomed)/totalprobx;
prob.immune.x0.curr=(prob.immune*(1-probx1.immune))/(1-totalprobx);
prob.treatprotected.x0.curr=(prob.treatprotected*(1-probx1.treatprotected))/(1-totalprobx);
prob.treatharmful.x0.curr=(prob.treatharmful*(1-probx1.treatharmful))/(1-totalprobx);
prob.doomed.x0.curr=(prob.doomed*(1-probx1.doomed))/(1-totalprobx);
itt.immune.treatprotected.curr=mean.immune.a1-mean.immune.a0;
sd.immune.a1.curr=sd.immune.a1; 
mean.immune.a0.curr=mean.immune.a0; 
sd.immune.a0.curr=sd.immune.a0;
sd.treatprotected.a1.curr=sd.treatprotected.a1;
mean.treatprotected.a0.curr=mean.treatprotected.a0;
sd.treatprotected.a0.curr=sd.treatprotected.a0;
itt.treatharmful.doomed.curr=mean.treatharmful.a1-mean.treatharmful.a0;
sd.treatharmful.a1.curr=sd.treatharmful.a1;
mean.treatharmful.a0.curr=mean.treatharmful.a0;
sd.treatharmful.a0.curr=sd.treatharmful.a0;
sd.doomed.a1.curr=sd.doomed.a1;
mean.doomed.a0.curr=mean.doomed.a0;
sd.doomed.a0.curr=sd.doomed.a0;
xeffect.curr=xeffect;

maxiters=500; # Maximum iterations in EM algorithm
iter=1;
lldist=1; # Will measure difference in log likelihood between iterations
toler=.001; # Tolerance for difference in log likelihood between iterations
# for stopping EM algorithm
conv=1;

while(lldist>toler && iter<maxiters){
  # Compute log likelihood with current parameters
ll.old.p1a=sum((1-aggressive.treatment)*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*esrd*x)*log(prob.doomed.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*esrd*x)*log(prob.doomed.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x1.curr)/(prob.immune.x1.curr+prob.treatharmful.x1.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+xeffect.curr,sd=sd.treatharmful.a0.curr))));
  ll.old.p1b=sum((1-aggressive.treatment)*esrd*x*log(((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+xeffect.curr,sd=sd.treatprotected.a0.curr))));
ll.old.p2=sum(aggressive.treatment*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.immune.a0.curr+itt.immune.treatprotected.curr+xeffect.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.immune.x1.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+itt.immune.treatprotected.curr+xeffect.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*x*((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.doomed.a0.curr+itt.treatharmful.doomed.curr+xeffect.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x1.curr)/(prob.treatharmful.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+itt.treatharmful.doomed.curr+xeffect.curr,sd=sd.treatharmful.a1.curr)));
ll.old.p3a=sum((1-aggressive.treatment)*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x0.curr)/(prob.immune.x0.curr+prob.treatharmful.x0.curr)*dnorm(y,mean=mean.treatharmful.a0.curr,sd=sd.treatharmful.a0.curr))))
ll.old.p3b=sum((1-aggressive.treatment)*esrd*(1-x)*log(((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatprotected.a0.curr,sd=sd.treatprotected.a0.curr))));
ll.old.p4=sum(aggressive.treatment*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.immune.a0.curr+itt.immune.treatprotected.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.immune.x0.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+itt.immune.treatprotected.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*(1-x)*((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.doomed.a0.curr+itt.treatharmful.doomed.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x0.curr)/(prob.treatharmful.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+itt.treatharmful.doomed.curr,sd=sd.treatharmful.a1.curr)));
ll.old=ll.old.p1a+ll.old.p1b+ll.old.p2+ll.old.p3a+ll.old.p3b+ll.old.p4;

# Expectation step
# Compute probabilities for each subject's principal strata
prob.immune.subject=rep(0,n);
prob.treatprotected.subject=rep(0,n);
prob.treatharmful.subject=rep(0,n);
prob.doomed.subject=rep(0,n);
for(j in 1:n){
  if(aggressive.treatment[j]==0 && esrd[j]==0 && x[j]==1){
    prob.immune.subject[j]=prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)/(prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)+prob.treatharmful.x1.curr*dnorm(y[j],mean=mean.treatharmful.a0.curr+xeffect.curr,sd=sd.treatharmful.a0.curr));
    prob.treatharmful.subject[j]=1-prob.immune.subject[j];
  }
  if(aggressive.treatment[j]==0 && esrd[j]==0 && x[j]==0){
  prob.immune.subject[j]=prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)/(prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)+prob.treatharmful.x0.curr*dnorm(y[j],mean=mean.treatharmful.a0.curr,sd=sd.treatharmful.a0.curr));
    prob.treatharmful.subject[j]=1-prob.immune.subject[j];
  }  
  if(aggressive.treatment[j]==0 && esrd[j]==1 && x[j]==1){
    prob.doomed.subject[j]=prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)/(prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)+prob.treatprotected.x1.curr*dnorm(y[j],mean=mean.treatprotected.a0.curr+xeffect.curr,sd=sd.treatprotected.a0.curr));
    prob.treatprotected.subject[j]=1-prob.doomed.subject[j];
  }
  if(aggressive.treatment[j]==0 && esrd[j]==1 && x[j]==0){
    prob.doomed.subject[j]=prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)/(prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)+prob.treatprotected.x0.curr*dnorm(y[j],mean=mean.treatprotected.a0.curr,sd=sd.treatprotected.a0.curr));
    prob.treatprotected.subject[j]=1-prob.doomed.subject[j];
  }
  if(aggressive.treatment[j]==1 && esrd[j]==0 && x[j]==1){
    prob.immune.subject[j]=prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a0.curr+itt.immune.treatprotected.curr+xeffect.curr,sd=sd.immune.a1.curr)/(prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a0.curr+itt.immune.treatprotected.curr+xeffect.curr,sd=sd.immune.a1.curr)+prob.treatprotected.x1.curr*dnorm(y[j],mean=mean.treatprotected.a0.curr+itt.immune.treatprotected.curr+xeffect.curr,sd=sd.treatprotected.a1.curr));
    prob.treatprotected.subject[j]=1-prob.immune.subject[j];
  }
  if(aggressive.treatment[j]==1 && esrd[j]==0 && x[j]==0){
    prob.immune.subject[j]=prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a0.curr+itt.immune.treatprotected.curr,sd=sd.immune.a1.curr)/(prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a0.curr+itt.immune.treatprotected.curr,sd=sd.immune.a1.curr)+prob.treatprotected.x0.curr*dnorm(y[j],mean=mean.treatprotected.a0.curr+itt.immune.treatprotected.curr,sd=sd.treatprotected.a1.curr));
    prob.treatprotected.subject[j]=1-prob.immune.subject[j];
  }  
  if(aggressive.treatment[j]==1 && esrd[j]==1 && x[j]==1){
    prob.doomed.subject[j]=prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a0.curr+itt.treatharmful.doomed.curr+xeffect.curr,sd=sd.doomed.a1.curr)/(prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a0.curr+itt.treatharmful.doomed.curr+xeffect.curr,sd=sd.doomed.a1.curr)+prob.treatharmful.x1.curr*dnorm(y[j],mean=mean.treatharmful.a0.curr+itt.treatharmful.doomed.curr+xeffect.curr,sd=sd.treatharmful.a1.curr));
    prob.treatharmful.subject[j]=1-prob.doomed.subject[j];
  }
  if(aggressive.treatment[j]==1 && esrd[j]==1 && x[j]==0){
    prob.doomed.subject[j]=prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a0.curr+itt.treatharmful.doomed.curr,sd=sd.doomed.a1.curr)/(prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a0.curr+itt.treatharmful.doomed.curr,sd=sd.doomed.a1.curr)+prob.treatharmful.x0.curr*dnorm(y[j],mean=mean.treatharmful.a0.curr+itt.treatharmful.doomed.curr,sd=sd.treatharmful.a1.curr));
    prob.treatharmful.subject[j]=1-prob.doomed.subject[j];
  }  
}
 
# Maximization Step (Use Conditional Maximization, i.e., ECM algorithm)
# First maximize expected likelihood over xeffect, principal strata means
# and then maximize likelihood over standard deviations in each principal
# stratra
# Set up weighted regression
weighty=rep(0,2*n);
weightx=rep(0,2*n);
weight.immune=rep(0,2*n);
weight.itt.immune.treatprotected=rep(0,2*n);
weight.treatprotected=rep(0,2*n);
weight.treatharmful=rep(0,2*n);
weight.itt.treatharmful.doomed.curr=rep(0,2*n);
weight.doomed=rep(0,2*n);

for(j in 1:n){
  if(aggressive.treatment[j]==0 && esrd[j]==0){
  tempweight1=(prob.immune.subject[j]/(prob.immune.subject[j]+prob.treatharmful.subject[j]))^.5/sd.immune.a0.curr;
  weighty[((j-1)*2+1)]=y[j]*tempweight1;
  weightx[((j-1)*2+1)]=x[j]*tempweight1;
  weight.immune[((j-1)*2+1)]=tempweight1;
  tempweight2=(prob.treatharmful.subject[j]/(prob.immune.subject[j]+prob.treatharmful.subject[j]))^.5/sd.treatharmful.a0.curr;
  weighty[(j*2)]=y[j]*tempweight2;
  weightx[(j*2)]=x[j]*tempweight2;
  weight.treatharmful[(j*2)]=tempweight2;
}
  if(aggressive.treatment[j]==0 && esrd[j]==1){
  tempweight1=(prob.doomed.subject[j]/(prob.doomed.subject[j]+prob.treatprotected.subject[j]))^.5/sd.doomed.a0.curr;
  weighty[((j-1)*2+1)]=y[j]*tempweight1;
  weightx[((j-1)*2+1)]=x[j]*tempweight1;
  weight.doomed[((j-1)*2+1)]=tempweight1;
  tempweight2=(prob.treatprotected.subject[j]/(prob.doomed.subject[j]+prob.treatprotected.subject[j]))^.5/sd.treatprotected.a0.curr;
  weighty[(j*2)]=y[j]*tempweight2;
  weightx[(j*2)]=x[j]*tempweight2;
  weight.treatprotected[(j*2)]=tempweight2;
}
  if(aggressive.treatment[j]==1 && esrd[j]==0){
    tempweight1=(prob.immune.subject[j]/(prob.immune.subject[j]+prob.treatprotected.subject[j]))^.5/sd.immune.a1.curr;
    weighty[((j-1)*2+1)]=y[j]*tempweight1;
    weightx[((j-1)*2+1)]=x[j]*tempweight1;
    weight.immune[((j-1)*2+1)]=tempweight1;
    weight.itt.immune.treatprotected[((j-1)*2+1)]=tempweight1;
    tempweight2=(prob.treatprotected.subject[j]/(prob.immune.subject[j]+prob.treatprotected.subject[j]))^.5/sd.treatprotected.a1.curr;
    weighty[(j*2)]=y[j]*tempweight2;
    weightx[(j*2)]=x[j]*tempweight2;
    weight.treatprotected[(j*2)]=tempweight2;
    weight.itt.immune.treatprotected[j*2]=tempweight2;
  }
  if(aggressive.treatment[j]==1 && esrd[j]==1){
    tempweight1=(prob.doomed.subject[j]/(prob.doomed.subject[j]+prob.treatharmful.subject[j]))^.5/sd.doomed.a1.curr;
    weighty[((j-1)*2+1)]=y[j]*tempweight1;
    weightx[((j-1)*2+1)]=x[j]*tempweight1;
    weight.doomed[((j-1)*2+1)]=tempweight1;
    weight.itt.treatharmful.doomed.curr[((j-1)*2+1)]=tempweight1;
    tempweight2=(prob.treatharmful.subject[j]/(prob.treatharmful.subject[j]+prob.doomed.subject[j]))^.5/sd.treatharmful.a1.curr;
    weighty[(j*2)]=y[j]*tempweight2;
    weightx[(j*2)]=x[j]*tempweight2;
    weight.treatharmful[(j*2)]=tempweight2;
    weight.itt.treatharmful.doomed.curr[((j-1)*2+1)]=tempweight1;
  }
}
  
weightreg=lm(weighty~weightx+weight.immune+weight.itt.immune.treatprotected+weight.treatharmful+weight.itt.treatharmful.doomed.curr+weight.doomed+weight.treatprotected-1);
coef.weightreg=coef(weightreg);
  xeffect.curr=coef.weightreg[1];
  mean.immune.a0.curr=coef.weightreg[2];
  itt.immune.treatprotected.curr=coef.weightreg[3];
  mean.treatharmful.a0.curr=coef.weightreg[4];
  itt.treatharmful.doomed.curr=coef.weightreg[5];
  mean.doomed.a0.curr=coef.weightreg[6];
  mean.treatprotected.a0.curr=coef.weightreg[7];

sd.immune.a1.curr=sqrt(sum(prob.immune.subject*aggressive.treatment*(y-mean.immune.a0.curr-itt.immune.treatprotected.curr+x*xeffect.curr)^2)/sum(prob.immune.subject*aggressive.treatment))
sd.immune.a0.curr=sqrt(sum(prob.immune.subject*(1-aggressive.treatment)*(y-mean.immune.a0.curr-x*xeffect.curr)^2)/sum(prob.immune.subject*(1-aggressive.treatment)));
sd.treatprotected.a1.curr=sqrt(sum(prob.treatprotected.subject*aggressive.treatment*(y-mean.treatprotected.a0.curr-itt.immune.treatprotected.curr-x*xeffect.curr)^2)/sum(prob.treatprotected.subject*aggressive.treatment));
sd.treatprotected.a0.curr=sqrt(sum(prob.treatprotected.subject*(1-aggressive.treatment)*(y-mean.treatprotected.a0.curr-x*xeffect.curr)^2)/sum(prob.treatprotected.subject*(1-aggressive.treatment)));
sd.doomed.a1.curr=sqrt(sum(prob.doomed.subject*aggressive.treatment*(y-mean.doomed.a0.curr-itt.treatharmful.doomed.curr-x*xeffect.curr)^2)/sum(prob.doomed.subject*aggressive.treatment));
sd.doomed.a0.curr=sqrt(sum(prob.doomed.subject*(1-aggressive.treatment)*(y-mean.doomed.a0.curr-x*xeffect.curr)^2)/sum(prob.doomed.subject*(1-aggressive.treatment)));
sd.treatharmful.a1.curr=sqrt(sum(prob.treatharmful.subject*aggressive.treatment*(y-mean.treatharmful.a0.curr-itt.treatharmful.doomed.curr-x*xeffect.curr)^2)/sum(prob.treatharmful.subject*aggressive.treatment));
sd.treatharmful.a0.curr=sqrt(sum(prob.treatharmful.subject*(1-aggressive.treatment)*(y-mean.treatharmful.a0.curr-x*xeffect.curr)^2)/sum(prob.treatharmful.subject*(1-aggressive.treatment)));
  
prob.immune.x1.curr=sum(prob.immune.subject*x)/sum(x);
prob.treatprotected.x1.curr=sum(prob.treatprotected.subject*x)/sum(x);
prob.treatharmful.x1.curr=sum(prob.treatharmful.subject*x)/sum(x);
prob.doomed.x1.curr=sum(prob.doomed.subject*x)/sum(x);
prob.immune.x0.curr=sum(prob.immune.subject*(1-x))/(n-sum(x));
prob.treatprotected.x0.curr=sum(prob.treatprotected.subject*(1-x))/(n-sum(x));
prob.treatharmful.x0.curr=sum(prob.treatharmful.subject*(1-x))/(n-sum(x));
prob.doomed.x0.curr=sum(prob.doomed.subject*(1-x))/(n-sum(x));

# Compute log likelihood with updated parameters
  ll.new.p1a=sum((1-aggressive.treatment)*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*esrd*x)*log(prob.doomed.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*esrd*x)*log(prob.doomed.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x1.curr)/(prob.immune.x1.curr+prob.treatharmful.x1.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+xeffect.curr,sd=sd.treatharmful.a0.curr))));
  ll.new.p1b=sum((1-aggressive.treatment)*esrd*x*log(((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+xeffect.curr,sd=sd.treatprotected.a0.curr))));
ll.new.p2=sum(aggressive.treatment*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.immune.a0.curr+itt.immune.treatprotected.curr+xeffect.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.immune.x1.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+itt.immune.treatprotected.curr+xeffect.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*x*((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.doomed.a0.curr+itt.treatharmful.doomed.curr+xeffect.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x1.curr)/(prob.treatharmful.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+itt.treatharmful.doomed.curr+xeffect.curr,sd=sd.treatharmful.a1.curr)));
ll.new.p3a=sum((1-aggressive.treatment)*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x0.curr)/(prob.immune.x0.curr+prob.treatharmful.x0.curr)*dnorm(y,mean=mean.treatharmful.a0.curr,sd=sd.treatharmful.a0.curr))))
ll.new.p3b=sum((1-aggressive.treatment)*esrd*(1-x)*log(((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatprotected.a0.curr,sd=sd.treatprotected.a0.curr))));
ll.new.p4=sum(aggressive.treatment*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.immune.a0.curr+itt.immune.treatprotected.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.immune.x0.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+itt.immune.treatprotected.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*(1-x)*((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.doomed.a0.curr+itt.treatharmful.doomed.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x0.curr)/(prob.treatharmful.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+itt.treatharmful.doomed.curr,sd=sd.treatharmful.a1.curr)));
ll.new=ll.new.p1a+ll.new.p1b+ll.new.p2+ll.new.p3a+ll.new.p3b+ll.new.p4;

# If log likelihood has become infinite, mark that estimate has not converged
# and restart simulation
if(is.na(ll.new)==FALSE){
# If any of the standard deviations has become zero, mark that estimate
# has not converged and restart simulation
if(sd.immune.a1.curr<.00001 | sd.immune.a0.curr<.00001 | sd.treatprotected.a1.curr<.00001 | sd.treatprotected.a0.curr<.00001 | sd.doomed.a1.curr<.00001 | sd.doomed.a0.curr<.00001 | sd.treatharmful.a1.curr<.00001 | sd.treatharmful.a0.curr<.00001){
  no.notconv=no.notconv+1
  conv=0;
  
}
if(sd.immune.a1.curr>=.00001 & sd.immune.a0.curr>=.00001 & sd.treatprotected.a1.curr>=.00001 & sd.treatprotected.a0.curr>=.00001 & sd.doomed.a1.curr>=.00001 & sd.doomed.a0.curr>=.00001 & sd.treatharmful.a1.curr>=.00001 & sd.treatharmful.a0.curr>=.00001){
lldist=ll.new-ll.old;
iter=iter+1;
}
}

if(is.na(ll.new)==TRUE){
  no.notconv=no.notconv+1;
  conv=0;
  lldist=0;
}
}

if(conv==1){
itt.immune.treatprotected.est.2[k]=itt.immune.treatprotected.curr;
mean.immune.a0.est.2[k]=mean.immune.a0.curr;
mean.treatprotected.a0.est.2[k]=mean.treatprotected.a0.curr;
itt.treatharmful.doomed.est.2[k]=itt.treatharmful.doomed.curr;
mean.treatharmful.a0.est.2[k]=mean.treatharmful.a0.curr;
mean.doomed.a0.est.2[k]=mean.doomed.a0.curr;
sd.immune.a1.est.2[k]=sd.immune.a1.curr;
sd.immune.a0.est.2[k]=sd.immune.a0.curr;
sd.treatprotected.a1.est.2[k]=sd.treatprotected.a1.curr;
sd.treatprotected.a0.est.2[k]=sd.treatprotected.a0.curr;
sd.treatharmful.a1.est.2[k]=sd.treatharmful.a1.curr;
sd.treatharmful.a0.est.2[k]=sd.treatharmful.a0.curr;
sd.doomed.a1.est.2[k]=sd.doomed.a1.curr;
sd.doomed.a0.est.2[k]=sd.doomed.a0.curr;
xeffect.est.2[k]=xeffect.curr;
lldistvec.2[k]=lldist;
k=k+1;
}
}

conv=as.logical((1-is.na(mean.immune.a1.est.1-mean.immune.a0.est.1))*(1-is.na(mean.treatprotected.a1.est.1-mean.treatprotected.a0.est.1))*(1-is.na(mean.treatharmful.a1.est.1-mean.treatharmful.a0.est.1))*(1-is.na(mean.doomed.a1.est.1-mean.doomed.a0.est.1))*(1-is.na(psi0.est))*(1-is.na(psi1.est))*(1-is.na(itt.immune.treatprotected.est.2))*(1-is.na(itt.treatharmful.doomed.est.2)));
sum(conv)
mean(mean.immune.a1.est.1[conv]-mean.immune.a0.est.1[conv])
sd(mean.immune.a1.est.1[conv]-mean.immune.a0.est.1[conv])
mean(mean.treatprotected.a1.est.1[conv]-mean.treatprotected.a0.est.1[conv])
sd(mean.treatprotected.a1.est.1[conv]-mean.treatprotected.a0.est.1[conv])
mean(mean.treatharmful.a1.est.1[conv]-mean.treatharmful.a0.est.1[conv])
sd(mean.treatharmful.a1.est.1[conv]-mean.treatharmful.a0.est.1[conv])
mean(mean.doomed.a1.est.1[conv]-mean.doomed.a0.est.1[conv])
sd(mean.doomed.a1.est.1[conv]-mean.doomed.a0.est.1[conv])
mean(psi0.est[conv])
sd(psi0.est[conv])
mean(psi1.est[conv])
sd(psi1.est[conv])
mean(betahat.a[conv]+betahat.mua[conv])
# sd(betahat.a+betahat.mua)
mean(itt.immune.treatprotected.est.2[conv])
sd(itt.immune.treatprotected.est.2[conv])
mean(itt.treatharmful.doomed.est.2[conv])
sd(itt.treatharmful.doomed.est.2[conv])




# Setting I(B)

# Simulation study
sims=1; # Number of simulations
# Vectors to store estimates from simulation study
mean.immune.a1.est.1=rep(0,sims);
mean.immune.a0.est.1=rep(0,sims);
mean.treatprotected.a1.est.1=rep(0,sims);
mean.treatprotected.a0.est.1=rep(0,sims);
mean.treatharmful.a1.est.1=rep(0,sims);
mean.treatharmful.a0.est.1=rep(0,sims);
mean.doomed.a1.est.1=rep(0,sims);
mean.doomed.a0.est.1=rep(0,sims);
sd.immune.a1.est.1=rep(0,sims);
sd.immune.a0.est.1=rep(0,sims);
sd.treatprotected.a1.est.1=rep(0,sims);
sd.treatprotected.a0.est.1=rep(0,sims);
sd.treatharmful.a1.est.1=rep(0,sims);
sd.treatharmful.a0.est.1=rep(0,sims);
sd.doomed.a1.est.1=rep(0,sims);
sd.doomed.a0.est.1=rep(0,sims);
prob.immune.est.1=rep(0,sims);
prob.treatprotected.est.1=rep(0,sims);
prob.treatharmful.est.1=rep(0,sims);
prob.doomed.est.1=rep(0,sims);
xeffect.est.1=rep(0,sims);
psi0.est=rep(0,sims); # Structural nested model coefficients
psi1.est=rep(0,sims); 
betahat.a=rep(0,sims);
betahat.mua=rep(0,sims);
lldistvec.1=rep(0,sims);
no.notconv=0; # Number of times EM algorithm does not converge
itt.immune.treatprotected.est.2=rep(0,sims);
mean.immune.a0.est.2=rep(0,sims);
mean.treatprotected.a0.est.2=rep(0,sims);
itt.treatharmful.doomed.est.2=rep(0,sims);
mean.treatharmful.a0.est.2=rep(0,sims);
mean.doomed.a0.est.2=rep(0,sims);
sd.immune.a1.est.2=rep(0,sims);
sd.immune.a0.est.2=rep(0,sims);
sd.treatprotected.a1.est.2=rep(0,sims);
sd.treatprotected.a0.est.2=rep(0,sims);
sd.treatharmful.a1.est.2=rep(0,sims);
sd.treatharmful.a0.est.2=rep(0,sims);
sd.doomed.a1.est.2=rep(0,sims);
sd.doomed.a0.est.2=rep(0,sims);
xeffect.est.2=rep(0,sims);
lldistvec.2=rep(0,sims);

k=1;
while(k<=sims){
# Simulation from model with a binary covariate
# Probabilities of principal strata
prob.immune=.25;
prob.treatprotected=.4;
prob.treatharmful=.05;
prob.doomed=.3;

# Means and standard deviations of each principal strata under aggressive treatment
# for hypertension and under not aggressive treatment for hypertension
mean.immune.a1=2; # Mean of immune subjects when aggressively treated for hypertension
sd.immune.a1=1; # Variance of immune subjects when aggressively treated
mean.immune.a0=1; # Mean of immune subjects when not agressively treated
sd.immune.a0=1; # Variance of immune subjects when not aggressively treated
mean.treatprotected.a1=2.5;
sd.treatprotected.a1=1;
mean.treatprotected.a0=1.5;
sd.treatprotected.a0=1;
mean.treatharmful.a1=1.25;
sd.treatharmful.a1=1;
mean.treatharmful.a0=.75;
sd.treatharmful.a0=1;
mean.doomed.a1=1.75;
sd.doomed.a1=1;
mean.doomed.a0=1.25;
sd.doomed.a0=.25;
xeffect=.5; # Addition to mean for binary effect

# Probabilities of x for different principal strata
probx1.immune=.5;
probx1.treatprotected=.75;
probx1.treatharmful=.25;
probx1.doomed=.5;
totalprobx=probx1.immune*prob.immune+probx1.treatprotected*prob.treatprotected+probx1.treatharmful*prob.treatharmful+probx1.doomed*prob.doomed;

n=500; # total sample size
# Assume that half of subjects are randomized into aggressive treatment
# of hypertension and half into nonaggressive treatment
aggressive.treatment=c(rep(1,n/2),rep(0,n/2));
principalstrata=rmultinom(n,1,c(prob.immune,prob.treatprotected,prob.treatharmful,prob.doomed));
# Dummy variables for principal strata
principalstrata.immune=principalstrata[1,];
principalstrata.treatprotected=principalstrata[2,];
principalstrata.treatharmful=principalstrata[3,];
principalstrata.doomed=principalstrata[4,];
# Covariate X
x=principalstrata.immune*rbinom(n,1,probx1.immune)+principalstrata.treatprotected*rbinom(n,1,probx1.treatprotected)+principalstrata.treatharmful*rbinom(n,1,probx1.treatharmful)+principalstrata.doomed*rbinom(n,1,probx1.doomed);

# ESRD dummy
esrd=principalstrata.doomed+principalstrata.treatprotected*(1-aggressive.treatment)+principalstrata.treatharmful*aggressive.treatment;
# Mean and standard deviation of outcome
mean.outcome=mean.immune.a1*principalstrata.immune*aggressive.treatment+mean.immune.a0*principalstrata.immune*(1-aggressive.treatment)+mean.treatprotected.a1*principalstrata.treatprotected*aggressive.treatment+mean.treatprotected.a0*principalstrata.treatprotected*(1-aggressive.treatment)+mean.treatharmful.a1*principalstrata.treatharmful*aggressive.treatment+mean.treatharmful.a0*principalstrata.treatharmful*(1-aggressive.treatment)+mean.doomed.a1*principalstrata.doomed*aggressive.treatment+mean.doomed.a0*principalstrata.doomed*(1-aggressive.treatment)+xeffect*x;

sd.outcome=sd.immune.a1*principalstrata.immune*aggressive.treatment+sd.immune.a0*principalstrata.immune*(1-aggressive.treatment)+sd.treatprotected.a1*principalstrata.treatprotected*aggressive.treatment+sd.treatprotected.a0*principalstrata.treatprotected*(1-aggressive.treatment)+sd.treatharmful.a1*principalstrata.treatharmful*aggressive.treatment+sd.treatharmful.a0*principalstrata.treatharmful*(1-aggressive.treatment)+sd.doomed.a1*principalstrata.doomed*aggressive.treatment+sd.doomed.a0*principalstrata.doomed*(1-aggressive.treatment);

# Generate outcomes from normal distribution
# y=rnorm(n,mean=mean.outcome,sd=sd.outcome);
# Generate outcomes from gamma distribution
scalegamma=sd.outcome^2/mean.outcome;
shapegamma=mean.outcome/scalegamma;
y=rgamma(n,shape=shapegamma,scale=scalegamma);

# EM algorithm

# Select starting values
# For now, start at truth
prob.immune.curr=prob.immune;
prob.treatprotected.curr=prob.treatprotected;
prob.treatharmful.curr=prob.treatharmful;
prob.doomed.curr=prob.doomed;
prob.immune.x1.curr=(prob.immune*probx1.immune)/totalprobx;
prob.treatprotected.x1.curr=(prob.treatprotected*probx1.treatprotected)/totalprobx;
prob.treatharmful.x1.curr=(prob.treatharmful*probx1.treatharmful)/totalprobx;
prob.doomed.x1.curr=(prob.doomed*probx1.doomed)/totalprobx;
prob.immune.x0.curr=(prob.immune*(1-probx1.immune))/(1-totalprobx);
prob.treatprotected.x0.curr=(prob.treatprotected*(1-probx1.treatprotected))/(1-totalprobx);
prob.treatharmful.x0.curr=(prob.treatharmful*(1-probx1.treatharmful))/(1-totalprobx);
prob.doomed.x0.curr=(prob.doomed*(1-probx1.doomed))/(1-totalprobx);
mean.immune.a1.curr=mean.immune.a1;
sd.immune.a1.curr=sd.immune.a1; 
mean.immune.a0.curr=mean.immune.a0; 
sd.immune.a0.curr=sd.immune.a0;
mean.treatprotected.a1.curr=mean.treatprotected.a1;
sd.treatprotected.a1.curr=sd.treatprotected.a1;
mean.treatprotected.a0.curr=mean.treatprotected.a0;
sd.treatprotected.a0.curr=sd.treatprotected.a0;
mean.treatharmful.a1.curr=mean.treatharmful.a1;
sd.treatharmful.a1.curr=sd.treatharmful.a1;
mean.treatharmful.a0.curr=mean.treatharmful.a0;
sd.treatharmful.a0.curr=sd.treatharmful.a0;
mean.doomed.a1.curr=mean.doomed.a1;
sd.doomed.a1.curr=sd.doomed.a1;
mean.doomed.a0.curr=mean.doomed.a0;
sd.doomed.a0.curr=sd.doomed.a0;
xeffect.curr=xeffect;

maxiters=500; # Maximum iterations in EM algorithm
iter=1;
lldist=1; # Will measure difference in log likelihood between iterations
toler=.001; # Tolerance for difference in log likelihood between iterations
# for stopping EM algorithm
conv=1;

while(lldist>toler && iter<maxiters){
  # Compute log likelihood with current parameters
ll.old.p1a=sum((1-aggressive.treatment)*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*esrd*x)*log(prob.doomed.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*esrd*x)*log(prob.doomed.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x1.curr)/(prob.immune.x1.curr+prob.treatharmful.x1.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+xeffect.curr,sd=sd.treatharmful.a0.curr))));
  ll.old.p1b=sum((1-aggressive.treatment)*esrd*x*log(((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+xeffect.curr,sd=sd.treatprotected.a0.curr))));
ll.old.p2=sum(aggressive.treatment*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.immune.a1.curr+xeffect.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.immune.x1.curr)*dnorm(y,mean=mean.treatprotected.a1.curr+xeffect.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*x*((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.doomed.a1.curr+xeffect.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x1.curr)/(prob.treatharmful.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatharmful.a1.curr+xeffect.curr,sd=sd.treatharmful.a1.curr)));
ll.old.p3a=sum((1-aggressive.treatment)*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x0.curr)/(prob.immune.x0.curr+prob.treatharmful.x0.curr)*dnorm(y,mean=mean.treatharmful.a0.curr,sd=sd.treatharmful.a0.curr))))
ll.old.p3b=sum((1-aggressive.treatment)*esrd*(1-x)*log(((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatprotected.a0.curr,sd=sd.treatprotected.a0.curr))));
ll.old.p4=sum(aggressive.treatment*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.immune.a1.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.immune.x0.curr)*dnorm(y,mean=mean.treatprotected.a1.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*(1-x)*((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.doomed.a1.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x0.curr)/(prob.treatharmful.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatharmful.a1.curr,sd=sd.treatharmful.a1.curr)));
ll.old=ll.old.p1a+ll.old.p1b+ll.old.p2+ll.old.p3a+ll.old.p3b+ll.old.p4;

# Expectation step
# Compute probabilities for each subject's principal strata
prob.immune.subject=rep(0,n);
prob.treatprotected.subject=rep(0,n);
prob.treatharmful.subject=rep(0,n);
prob.doomed.subject=rep(0,n);
for(j in 1:n){
  if(aggressive.treatment[j]==0 && esrd[j]==0 && x[j]==1){
    prob.immune.subject[j]=prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)/(prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)+prob.treatharmful.x1.curr*dnorm(y[j],mean=mean.treatharmful.a0.curr+xeffect.curr,sd=sd.treatharmful.a0.curr));
    prob.treatharmful.subject[j]=1-prob.immune.subject[j];
  }
  if(aggressive.treatment[j]==0 && esrd[j]==0 && x[j]==0){
  prob.immune.subject[j]=prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)/(prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)+prob.treatharmful.x0.curr*dnorm(y[j],mean=mean.treatharmful.a0.curr,sd=sd.treatharmful.a0.curr));
    prob.treatharmful.subject[j]=1-prob.immune.subject[j];
  }  
  if(aggressive.treatment[j]==0 && esrd[j]==1 && x[j]==1){
    prob.doomed.subject[j]=prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)/(prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)+prob.treatprotected.x1.curr*dnorm(y[j],mean=mean.treatprotected.a0.curr+xeffect.curr,sd=sd.treatprotected.a0.curr));
    prob.treatprotected.subject[j]=1-prob.doomed.subject[j];
  }
  if(aggressive.treatment[j]==0 && esrd[j]==1 && x[j]==0){
    prob.doomed.subject[j]=prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)/(prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)+prob.treatprotected.x0.curr*dnorm(y[j],mean=mean.treatprotected.a0.curr,sd=sd.treatprotected.a0.curr));
    prob.treatprotected.subject[j]=1-prob.doomed.subject[j];
  }
  if(aggressive.treatment[j]==1 && esrd[j]==0 && x[j]==1){
    prob.immune.subject[j]=prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a1.curr+xeffect.curr,sd=sd.immune.a1.curr)/(prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a1.curr+xeffect.curr,sd=sd.immune.a1.curr)+prob.treatprotected.x1.curr*dnorm(y[j],mean=mean.treatprotected.a1.curr+xeffect.curr,sd=sd.treatprotected.a1.curr));
    prob.treatprotected.subject[j]=1-prob.immune.subject[j];
  }
  if(aggressive.treatment[j]==1 && esrd[j]==0 && x[j]==0){
    prob.immune.subject[j]=prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a1.curr,sd=sd.immune.a1.curr)/(prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a1.curr,sd=sd.immune.a1.curr)+prob.treatprotected.x0.curr*dnorm(y[j],mean=mean.treatprotected.a1.curr,sd=sd.treatprotected.a1.curr));
    prob.treatprotected.subject[j]=1-prob.immune.subject[j];
  }  
  if(aggressive.treatment[j]==1 && esrd[j]==1 && x[j]==1){
    prob.doomed.subject[j]=prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a1.curr+xeffect.curr,sd=sd.doomed.a1.curr)/(prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a1.curr+xeffect.curr,sd=sd.doomed.a1.curr)+prob.treatharmful.x1.curr*dnorm(y[j],mean=mean.treatharmful.a1.curr+xeffect.curr,sd=sd.treatharmful.a1.curr));
    prob.treatharmful.subject[j]=1-prob.doomed.subject[j];
  }
  if(aggressive.treatment[j]==1 && esrd[j]==1 && x[j]==0){
    prob.doomed.subject[j]=prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a1.curr,sd=sd.doomed.a1.curr)/(prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a1.curr,sd=sd.doomed.a1.curr)+prob.treatharmful.x0.curr*dnorm(y[j],mean=mean.treatharmful.a1.curr,sd=sd.treatharmful.a1.curr));
    prob.treatharmful.subject[j]=1-prob.doomed.subject[j];
  }  
}
 
# Maximization Step (Use Conditional Maximization, i.e., ECM algorithm)
# First maximize expected likelihood over xeffect, principal strata means
# and then maximize likelihood over standard deviations in each principal
# stratra
# Set up weighted regression
weighty=rep(0,2*n);
weightx=rep(0,2*n);
weight.immune.a1=rep(0,2*n);
weight.immune.a0=rep(0,2*n);
weight.treatprotected.a1=rep(0,2*n);
weight.treatprotected.a0=rep(0,2*n);
weight.treatharmful.a1=rep(0,2*n);
weight.treatharmful.a0=rep(0,2*n);
weight.doomed.a1=rep(0,2*n);
weight.doomed.a0=rep(0,2*n);

for(j in 1:n){
  if(aggressive.treatment[j]==0 && esrd[j]==0){
  tempweight1=(prob.immune.subject[j]/(prob.immune.subject[j]+prob.treatharmful.subject[j]))^.5/sd.immune.a0.curr;
  weighty[((j-1)*2+1)]=y[j]*tempweight1;
  weightx[((j-1)*2+1)]=x[j]*tempweight1;
  weight.immune.a0[((j-1)*2+1)]=tempweight1;
  tempweight2=(prob.treatharmful.subject[j]/(prob.immune.subject[j]+prob.treatharmful.subject[j]))^.5/sd.treatharmful.a0.curr;
  weighty[(j*2)]=y[j]*tempweight2;
  weightx[(j*2)]=x[j]*tempweight2;
  weight.treatharmful.a0[(j*2)]=tempweight2;
}
  if(aggressive.treatment[j]==0 && esrd[j]==1){
  tempweight1=(prob.doomed.subject[j]/(prob.doomed.subject[j]+prob.treatprotected.subject[j]))^.5/sd.doomed.a0.curr;
  weighty[((j-1)*2+1)]=y[j]*tempweight1;
  weightx[((j-1)*2+1)]=x[j]*tempweight1;
  weight.doomed.a0[((j-1)*2+1)]=tempweight1;
  tempweight2=(prob.treatprotected.subject[j]/(prob.doomed.subject[j]+prob.treatprotected.subject[j]))^.5/sd.treatprotected.a0.curr;
  weighty[(j*2)]=y[j]*tempweight2;
  weightx[(j*2)]=x[j]*tempweight2;
  weight.treatprotected.a0[(j*2)]=tempweight2;
}
  if(aggressive.treatment[j]==1 && esrd[j]==0){
    tempweight1=(prob.immune.subject[j]/(prob.immune.subject[j]+prob.treatprotected.subject[j]))^.5/sd.immune.a1.curr;
    weighty[((j-1)*2+1)]=y[j]*tempweight1;
    weightx[((j-1)*2+1)]=x[j]*tempweight1;
    weight.immune.a1[((j-1)*2+1)]=tempweight1;
    tempweight2=(prob.treatprotected.subject[j]/(prob.immune.subject[j]+prob.treatprotected.subject[j]))^.5/sd.treatprotected.a1.curr;
    weighty[(j*2)]=y[j]*tempweight2;
    weightx[(j*2)]=x[j]*tempweight2;
    weight.treatprotected.a1[(j*2)]=tempweight2;
  }
  if(aggressive.treatment[j]==1 && esrd[j]==1){
    tempweight1=(prob.doomed.subject[j]/(prob.doomed.subject[j]+prob.treatharmful.subject[j]))^.5/sd.doomed.a1.curr;
    weighty[((j-1)*2+1)]=y[j]*tempweight1;
    weightx[((j-1)*2+1)]=x[j]*tempweight1;
    weight.doomed.a1[((j-1)*2+1)]=tempweight1;
    tempweight2=(prob.treatharmful.subject[j]/(prob.treatharmful.subject[j]+prob.doomed.subject[j]))^.5/sd.treatharmful.a1.curr;
    weighty[(j*2)]=y[j]*tempweight2;
    weightx[(j*2)]=x[j]*tempweight2;
    weight.treatharmful.a1[(j*2)]=tempweight2;
  }
}
  
weightreg=lm(weighty~weightx+weight.immune.a0+weight.immune.a1+weight.treatharmful.a0+weight.treatharmful.a1+weight.doomed.a0+weight.doomed.a1+weight.treatprotected.a0+weight.treatprotected.a1-1);
coef.weightreg=coef(weightreg);
  xeffect.curr=coef.weightreg[1];
  mean.immune.a0.curr=coef.weightreg[2];
  mean.immune.a1.curr=coef.weightreg[3];
  mean.treatharmful.a0.curr=coef.weightreg[4];
  mean.treatharmful.a1.curr=coef.weightreg[5];
  mean.doomed.a0.curr=coef.weightreg[6];
  mean.doomed.a1.curr=coef.weightreg[7];
  mean.treatprotected.a0.curr=coef.weightreg[8];
  mean.treatprotected.a1.curr=coef.weightreg[9];

sd.immune.a1.curr=sqrt(sum(prob.immune.subject*aggressive.treatment*(y-mean.immune.a1.curr-x*xeffect.curr)^2)/sum(prob.immune.subject*aggressive.treatment))
sd.immune.a0.curr=sqrt(sum(prob.immune.subject*(1-aggressive.treatment)*(y-mean.immune.a0.curr-x*xeffect.curr)^2)/sum(prob.immune.subject*(1-aggressive.treatment)));
sd.treatprotected.a1.curr=sqrt(sum(prob.treatprotected.subject*aggressive.treatment*(y-mean.treatprotected.a1.curr-x*xeffect.curr)^2)/sum(prob.treatprotected.subject*aggressive.treatment));
sd.treatprotected.a0.curr=sqrt(sum(prob.treatprotected.subject*(1-aggressive.treatment)*(y-mean.treatprotected.a0.curr-x*xeffect.curr)^2)/sum(prob.treatprotected.subject*(1-aggressive.treatment)));
sd.doomed.a1.curr=sqrt(sum(prob.doomed.subject*aggressive.treatment*(y-mean.doomed.a1.curr-x*xeffect.curr)^2)/sum(prob.doomed.subject*aggressive.treatment));
sd.doomed.a0.curr=sqrt(sum(prob.doomed.subject*(1-aggressive.treatment)*(y-mean.doomed.a0.curr-x*xeffect.curr)^2)/sum(prob.doomed.subject*(1-aggressive.treatment)));
sd.treatharmful.a1.curr=sqrt(sum(prob.treatharmful.subject*aggressive.treatment*(y-mean.treatharmful.a1.curr-x*xeffect.curr)^2)/sum(prob.treatharmful.subject*aggressive.treatment));
sd.treatharmful.a0.curr=sqrt(sum(prob.treatharmful.subject*(1-aggressive.treatment)*(y-mean.treatharmful.a0.curr-x*xeffect.curr)^2)/sum(prob.treatharmful.subject*(1-aggressive.treatment)));
  
prob.immune.x1.curr=sum(prob.immune.subject*x)/sum(x);
prob.treatprotected.x1.curr=sum(prob.treatprotected.subject*x)/sum(x);
prob.treatharmful.x1.curr=sum(prob.treatharmful.subject*x)/sum(x);
prob.doomed.x1.curr=sum(prob.doomed.subject*x)/sum(x);
prob.immune.x0.curr=sum(prob.immune.subject*(1-x))/(n-sum(x));
prob.treatprotected.x0.curr=sum(prob.treatprotected.subject*(1-x))/(n-sum(x));
prob.treatharmful.x0.curr=sum(prob.treatharmful.subject*(1-x))/(n-sum(x));
prob.doomed.x0.curr=sum(prob.doomed.subject*(1-x))/(n-sum(x));

# Compute log likelihood with updated parameters
  ll.new.p1a=sum((1-aggressive.treatment)*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*esrd*x)*log(prob.doomed.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*esrd*x)*log(prob.doomed.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x1.curr)/(prob.immune.x1.curr+prob.treatharmful.x1.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+xeffect.curr,sd=sd.treatharmful.a0.curr))));
  ll.new.p1b=sum((1-aggressive.treatment)*esrd*x*log(((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+xeffect.curr,sd=sd.treatprotected.a0.curr))));
ll.new.p2=sum(aggressive.treatment*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.immune.a1.curr+xeffect.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.immune.x1.curr)*dnorm(y,mean=mean.treatprotected.a1.curr+xeffect.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*x*((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.doomed.a1.curr+xeffect.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x1.curr)/(prob.treatharmful.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatharmful.a1.curr+xeffect.curr,sd=sd.treatharmful.a1.curr)));
ll.new.p3a=sum((1-aggressive.treatment)*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x0.curr)/(prob.immune.x0.curr+prob.treatharmful.x0.curr)*dnorm(y,mean=mean.treatharmful.a0.curr,sd=sd.treatharmful.a0.curr))))
ll.new.p3b=sum((1-aggressive.treatment)*esrd*(1-x)*log(((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatprotected.a0.curr,sd=sd.treatprotected.a0.curr))));
ll.new.p4=sum(aggressive.treatment*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.immune.a1.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.immune.x0.curr)*dnorm(y,mean=mean.treatprotected.a1.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*(1-x)*((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.doomed.a1.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x0.curr)/(prob.treatharmful.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatharmful.a1.curr,sd=sd.treatharmful.a1.curr)));
ll.new=ll.new.p1a+ll.new.p1b+ll.new.p2+ll.new.p3a+ll.new.p3b+ll.new.p4;

# If log likelihood has become infinite, mark that estimate has not converged
# and restart simulation
if(is.na(ll.new)==FALSE){
# If any of the standard deviations has become zero, mark that estimate
# has not converged and restart simulation
if(sd.immune.a1.curr<.00001 | sd.immune.a0.curr<.00001 | sd.treatprotected.a1.curr<.00001 | sd.treatprotected.a0.curr<.00001 | sd.doomed.a1.curr<.00001 | sd.doomed.a0.curr<.00001 | sd.treatharmful.a1.curr<.00001 | sd.treatharmful.a0.curr<.00001){
  no.notconv=no.notconv+1
  conv=0;
  
}
if(sd.immune.a1.curr>=.00001 & sd.immune.a0.curr>=.00001 & sd.treatprotected.a1.curr>=.00001 & sd.treatprotected.a0.curr>=.00001 & sd.doomed.a1.curr>=.00001 & sd.doomed.a0.curr>=.00001 & sd.treatharmful.a1.curr>=.00001 & sd.treatharmful.a0.curr>=.00001){
lldist=ll.new-ll.old;
iter=iter+1;
}
}

if(is.na(ll.new)==TRUE){
  no.notconv=no.notconv+1;
  conv=0;
  lldist=0;
}
}

mean.immune.a1.est.1[k]=mean.immune.a1.curr;
mean.immune.a0.est.1[k]=mean.immune.a0.curr;
mean.treatprotected.a1.est.1[k]=mean.treatprotected.a1.curr;
mean.treatprotected.a0.est.1[k]=mean.treatprotected.a0.curr;
mean.treatharmful.a1.est.1[k]=mean.treatharmful.a1.curr;
mean.treatharmful.a0.est.1[k]=mean.treatharmful.a0.curr;
mean.doomed.a1.est.1[k]=mean.doomed.a1.curr;
mean.doomed.a0.est.1[k]=mean.doomed.a0.curr;
sd.immune.a1.est.1[k]=sd.immune.a1.curr;
sd.immune.a0.est.1[k]=sd.immune.a0.curr;
sd.treatprotected.a1.est.1[k]=sd.treatprotected.a1.curr;
sd.treatprotected.a0.est.1[k]=sd.treatprotected.a0.curr;
sd.treatharmful.a1.est.1[k]=sd.treatharmful.a1.curr;
sd.treatharmful.a0.est.1[k]=sd.treatharmful.a0.curr;
sd.doomed.a1.est.1[k]=sd.doomed.a1.curr;
sd.doomed.a0.est.1[k]=sd.doomed.a0.curr;
prob.immune.est.1[k]=prob.immune.curr;
prob.treatprotected.est.1[k]=prob.treatprotected.curr;
prob.treatharmful.est.1[k]=prob.treatharmful.curr;
prob.doomed.est.1[k]=prob.doomed.curr;
xeffect.est.1[k]=xeffect.curr;
lldistvec.1[k]=lldist;

# Expected auxiliary stratification estimates
ax=aggressive.treatment*x;
auxreg1=glm(esrd~x+aggressive.treatment+ax,family=binomial);
muhatax=predict.glm(auxreg1,type="response");
a.times.muhatax=aggressive.treatment*muhatax;
auxreg2=lm(y~muhatax+aggressive.treatment+a.times.muhatax);
betahat.a[k]=coef(auxreg2)[3];
betahat.mua[k]=coef(auxreg2)[4];

# Structural nested model estimates
# Use Newton's Method
# Start at true values
psi1.curr=mean.doomed.a1-mean.doomed.a0;
psi0.curr=mean.treatprotected.a1-mean.treatprotected.a0;

# Estimate "mediating score"
tempreg=glm(esrd~x,family=binomial,subset=(aggressive.treatment==1));
mediatingscore=expit(coef(tempreg)[1]+coef(tempreg)[2]*x);

maxiters=500; # Maximum iterations in Newton-Raphson algorithm
iter=1;
dist=1;
toler=.001;
while(iter<maxiters && dist>toler){
  psiold=matrix(c(psi0.curr,psi1.curr),ncol=1);
  y0psi=y-psi0.curr*aggressive.treatment*(1-esrd)-psi1.curr*aggressive.treatment*esrd;
  tempreg=lm(y0psi~x);
  tempbetahat0=coef(tempreg)[1];
  tempbetahat1=coef(tempreg)[2];
  g1=sum((aggressive.treatment-.5)*(y0psi-tempbetahat0-tempbetahat1*x)*(1-mediatingscore));
  g2=sum((aggressive.treatment-.5)*(y0psi-tempbetahat0-tempbetahat1*x)*mediatingscore);
  Jmat=matrix(rep(0,4),ncol=2);
  Jmat[1,1]=sum(-(aggressive.treatment-.5)*aggressive.treatment*(1-esrd)*(1-mediatingscore));
  Jmat[1,2]=sum(-(aggressive.treatment-.5)*aggressive.treatment*esrd*(1-mediatingscore));
  Jmat[2,1]=sum(-(aggressive.treatment-.5)*aggressive.treatment*(1-esrd)*mediatingscore);
  Jmat[2,2]=sum(-(aggressive.treatment-.5)*aggressive.treatment*esrd*mediatingscore);
  psinew=psiold-solve(Jmat)%*%matrix(c(g1,g2),ncol=1);
  dist=(sum((psinew-psiold)^2))^.5;
  psi0.curr=psinew[1];
  psi1.curr=psinew[2];
  iter=iter+1;
}
psi0.est[k]=psi0.curr;
psi1.est[k]=psi1.curr;


# Estimates that assume ITT(immune)=ITT(treatment protected)
#                  ITT(treatment harmful)=ITT(doomed)
# EM algorithm

# Select starting values
# For now, start at truth
prob.immune.curr=prob.immune;
prob.treatprotected.curr=prob.treatprotected;
prob.treatharmful.curr=prob.treatharmful;
prob.doomed.curr=prob.doomed;
prob.immune.x1.curr=(prob.immune*probx1.immune)/totalprobx;
prob.treatprotected.x1.curr=(prob.treatprotected*probx1.treatprotected)/totalprobx;
prob.treatharmful.x1.curr=(prob.treatharmful*probx1.treatharmful)/totalprobx;
prob.doomed.x1.curr=(prob.doomed*probx1.doomed)/totalprobx;
prob.immune.x0.curr=(prob.immune*(1-probx1.immune))/(1-totalprobx);
prob.treatprotected.x0.curr=(prob.treatprotected*(1-probx1.treatprotected))/(1-totalprobx);
prob.treatharmful.x0.curr=(prob.treatharmful*(1-probx1.treatharmful))/(1-totalprobx);
prob.doomed.x0.curr=(prob.doomed*(1-probx1.doomed))/(1-totalprobx);
itt.immune.treatprotected.curr=mean.immune.a1-mean.immune.a0;
sd.immune.a1.curr=sd.immune.a1; 
mean.immune.a0.curr=mean.immune.a0; 
sd.immune.a0.curr=sd.immune.a0;
sd.treatprotected.a1.curr=sd.treatprotected.a1;
mean.treatprotected.a0.curr=mean.treatprotected.a0;
sd.treatprotected.a0.curr=sd.treatprotected.a0;
itt.treatharmful.doomed.curr=mean.treatharmful.a1-mean.treatharmful.a0;
sd.treatharmful.a1.curr=sd.treatharmful.a1;
mean.treatharmful.a0.curr=mean.treatharmful.a0;
sd.treatharmful.a0.curr=sd.treatharmful.a0;
sd.doomed.a1.curr=sd.doomed.a1;
mean.doomed.a0.curr=mean.doomed.a0;
sd.doomed.a0.curr=sd.doomed.a0;
xeffect.curr=xeffect;

maxiters=500; # Maximum iterations in EM algorithm
iter=1;
lldist=1; # Will measure difference in log likelihood between iterations
toler=.001; # Tolerance for difference in log likelihood between iterations
# for stopping EM algorithm
conv=1;

while(lldist>toler && iter<maxiters){
  # Compute log likelihood with current parameters
ll.old.p1a=sum((1-aggressive.treatment)*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*esrd*x)*log(prob.doomed.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*esrd*x)*log(prob.doomed.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x1.curr)/(prob.immune.x1.curr+prob.treatharmful.x1.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+xeffect.curr,sd=sd.treatharmful.a0.curr))));
  ll.old.p1b=sum((1-aggressive.treatment)*esrd*x*log(((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+xeffect.curr,sd=sd.treatprotected.a0.curr))));
ll.old.p2=sum(aggressive.treatment*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.immune.a0.curr+itt.immune.treatprotected.curr+xeffect.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.immune.x1.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+itt.immune.treatprotected.curr+xeffect.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*x*((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.doomed.a0.curr+itt.treatharmful.doomed.curr+xeffect.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x1.curr)/(prob.treatharmful.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+itt.treatharmful.doomed.curr+xeffect.curr,sd=sd.treatharmful.a1.curr)));
ll.old.p3a=sum((1-aggressive.treatment)*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x0.curr)/(prob.immune.x0.curr+prob.treatharmful.x0.curr)*dnorm(y,mean=mean.treatharmful.a0.curr,sd=sd.treatharmful.a0.curr))))
ll.old.p3b=sum((1-aggressive.treatment)*esrd*(1-x)*log(((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatprotected.a0.curr,sd=sd.treatprotected.a0.curr))));
ll.old.p4=sum(aggressive.treatment*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.immune.a0.curr+itt.immune.treatprotected.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.immune.x0.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+itt.immune.treatprotected.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*(1-x)*((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.doomed.a0.curr+itt.treatharmful.doomed.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x0.curr)/(prob.treatharmful.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+itt.treatharmful.doomed.curr,sd=sd.treatharmful.a1.curr)));
ll.old=ll.old.p1a+ll.old.p1b+ll.old.p2+ll.old.p3a+ll.old.p3b+ll.old.p4;

# Expectation step
# Compute probabilities for each subject's principal strata
prob.immune.subject=rep(0,n);
prob.treatprotected.subject=rep(0,n);
prob.treatharmful.subject=rep(0,n);
prob.doomed.subject=rep(0,n);
for(j in 1:n){
  if(aggressive.treatment[j]==0 && esrd[j]==0 && x[j]==1){
    prob.immune.subject[j]=prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)/(prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)+prob.treatharmful.x1.curr*dnorm(y[j],mean=mean.treatharmful.a0.curr+xeffect.curr,sd=sd.treatharmful.a0.curr));
    prob.treatharmful.subject[j]=1-prob.immune.subject[j];
  }
  if(aggressive.treatment[j]==0 && esrd[j]==0 && x[j]==0){
  prob.immune.subject[j]=prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)/(prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)+prob.treatharmful.x0.curr*dnorm(y[j],mean=mean.treatharmful.a0.curr,sd=sd.treatharmful.a0.curr));
    prob.treatharmful.subject[j]=1-prob.immune.subject[j];
  }  
  if(aggressive.treatment[j]==0 && esrd[j]==1 && x[j]==1){
    prob.doomed.subject[j]=prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)/(prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)+prob.treatprotected.x1.curr*dnorm(y[j],mean=mean.treatprotected.a0.curr+xeffect.curr,sd=sd.treatprotected.a0.curr));
    prob.treatprotected.subject[j]=1-prob.doomed.subject[j];
  }
  if(aggressive.treatment[j]==0 && esrd[j]==1 && x[j]==0){
    prob.doomed.subject[j]=prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)/(prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)+prob.treatprotected.x0.curr*dnorm(y[j],mean=mean.treatprotected.a0.curr,sd=sd.treatprotected.a0.curr));
    prob.treatprotected.subject[j]=1-prob.doomed.subject[j];
  }
  if(aggressive.treatment[j]==1 && esrd[j]==0 && x[j]==1){
    prob.immune.subject[j]=prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a0.curr+itt.immune.treatprotected.curr+xeffect.curr,sd=sd.immune.a1.curr)/(prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a0.curr+itt.immune.treatprotected.curr+xeffect.curr,sd=sd.immune.a1.curr)+prob.treatprotected.x1.curr*dnorm(y[j],mean=mean.treatprotected.a0.curr+itt.immune.treatprotected.curr+xeffect.curr,sd=sd.treatprotected.a1.curr));
    prob.treatprotected.subject[j]=1-prob.immune.subject[j];
  }
  if(aggressive.treatment[j]==1 && esrd[j]==0 && x[j]==0){
    prob.immune.subject[j]=prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a0.curr+itt.immune.treatprotected.curr,sd=sd.immune.a1.curr)/(prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a0.curr+itt.immune.treatprotected.curr,sd=sd.immune.a1.curr)+prob.treatprotected.x0.curr*dnorm(y[j],mean=mean.treatprotected.a0.curr+itt.immune.treatprotected.curr,sd=sd.treatprotected.a1.curr));
    prob.treatprotected.subject[j]=1-prob.immune.subject[j];
  }  
  if(aggressive.treatment[j]==1 && esrd[j]==1 && x[j]==1){
    prob.doomed.subject[j]=prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a0.curr+itt.treatharmful.doomed.curr+xeffect.curr,sd=sd.doomed.a1.curr)/(prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a0.curr+itt.treatharmful.doomed.curr+xeffect.curr,sd=sd.doomed.a1.curr)+prob.treatharmful.x1.curr*dnorm(y[j],mean=mean.treatharmful.a0.curr+itt.treatharmful.doomed.curr+xeffect.curr,sd=sd.treatharmful.a1.curr));
    prob.treatharmful.subject[j]=1-prob.doomed.subject[j];
  }
  if(aggressive.treatment[j]==1 && esrd[j]==1 && x[j]==0){
    prob.doomed.subject[j]=prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a0.curr+itt.treatharmful.doomed.curr,sd=sd.doomed.a1.curr)/(prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a0.curr+itt.treatharmful.doomed.curr,sd=sd.doomed.a1.curr)+prob.treatharmful.x0.curr*dnorm(y[j],mean=mean.treatharmful.a0.curr+itt.treatharmful.doomed.curr,sd=sd.treatharmful.a1.curr));
    prob.treatharmful.subject[j]=1-prob.doomed.subject[j];
  }  
}
 
# Maximization Step (Use Conditional Maximization, i.e., ECM algorithm)
# First maximize expected likelihood over xeffect, principal strata means
# and then maximize likelihood over standard deviations in each principal
# stratra
# Set up weighted regression
weighty=rep(0,2*n);
weightx=rep(0,2*n);
weight.immune=rep(0,2*n);
weight.itt.immune.treatprotected=rep(0,2*n);
weight.treatprotected=rep(0,2*n);
weight.treatharmful=rep(0,2*n);
weight.itt.treatharmful.doomed=rep(0,2*n);
weight.doomed=rep(0,2*n);

for(j in 1:n){
  if(aggressive.treatment[j]==0 && esrd[j]==0){
  tempweight1=(prob.immune.subject[j]/(prob.immune.subject[j]+prob.treatharmful.subject[j]))^.5/sd.immune.a0.curr;
  weighty[((j-1)*2+1)]=y[j]*tempweight1;
  weightx[((j-1)*2+1)]=x[j]*tempweight1;
  weight.immune[((j-1)*2+1)]=tempweight1;
  tempweight2=(prob.treatharmful.subject[j]/(prob.immune.subject[j]+prob.treatharmful.subject[j]))^.5/sd.treatharmful.a0.curr;
  weighty[(j*2)]=y[j]*tempweight2;
  weightx[(j*2)]=x[j]*tempweight2;
  weight.treatharmful[(j*2)]=tempweight2;
}
  if(aggressive.treatment[j]==0 && esrd[j]==1){
  tempweight1=(prob.doomed.subject[j]/(prob.doomed.subject[j]+prob.treatprotected.subject[j]))^.5/sd.doomed.a0.curr;
  weighty[((j-1)*2+1)]=y[j]*tempweight1;
  weightx[((j-1)*2+1)]=x[j]*tempweight1;
  weight.doomed[((j-1)*2+1)]=tempweight1;
  tempweight2=(prob.treatprotected.subject[j]/(prob.doomed.subject[j]+prob.treatprotected.subject[j]))^.5/sd.treatprotected.a0.curr;
  weighty[(j*2)]=y[j]*tempweight2;
  weightx[(j*2)]=x[j]*tempweight2;
  weight.treatprotected[(j*2)]=tempweight2;
}
  if(aggressive.treatment[j]==1 && esrd[j]==0){
    tempweight1=(prob.immune.subject[j]/(prob.immune.subject[j]+prob.treatprotected.subject[j]))^.5/sd.immune.a1.curr;
    weighty[((j-1)*2+1)]=y[j]*tempweight1;
    weightx[((j-1)*2+1)]=x[j]*tempweight1;
    weight.immune[((j-1)*2+1)]=tempweight1;
    weight.itt.immune.treatprotected[((j-1)*2+1)]=tempweight1;
    tempweight2=(prob.treatprotected.subject[j]/(prob.immune.subject[j]+prob.treatprotected.subject[j]))^.5/sd.treatprotected.a1.curr;
    weighty[(j*2)]=y[j]*tempweight2;
    weightx[(j*2)]=x[j]*tempweight2;
    weight.treatprotected[(j*2)]=tempweight2;
    weight.itt.immune.treatprotected[j*2]=tempweight2;
  }
  if(aggressive.treatment[j]==1 && esrd[j]==1){
    tempweight1=(prob.doomed.subject[j]/(prob.doomed.subject[j]+prob.treatharmful.subject[j]))^.5/sd.doomed.a1.curr;
    weighty[((j-1)*2+1)]=y[j]*tempweight1;
    weightx[((j-1)*2+1)]=x[j]*tempweight1;
    weight.doomed[((j-1)*2+1)]=tempweight1;
    weight.itt.treatharmful.doomed[((j-1)*2+1)]=tempweight1;
    tempweight2=(prob.treatharmful.subject[j]/(prob.treatharmful.subject[j]+prob.doomed.subject[j]))^.5/sd.treatharmful.a1.curr;
    weighty[(j*2)]=y[j]*tempweight2;
    weightx[(j*2)]=x[j]*tempweight2;
    weight.treatharmful[(j*2)]=tempweight2;
    weight.itt.treatharmful.doomed[((j-1)*2+1)]=tempweight1;
  }
}
  
weightreg=lm(weighty~weightx+weight.immune+weight.itt.immune.treatprotected+weight.treatharmful+weight.itt.treatharmful.doomed+weight.doomed+weight.treatprotected-1);
coef.weightreg=coef(weightreg);
  xeffect.curr=coef.weightreg[1];
  mean.immune.a0.curr=coef.weightreg[2];
  itt.immune.treatprotected.curr=coef.weightreg[3];
  mean.treatharmful.a0.curr=coef.weightreg[4];
  itt.treatharmful.doomed.curr=coef.weightreg[5];
  mean.doomed.a0.curr=coef.weightreg[6];
  mean.treatprotected.a0.curr=coef.weightreg[7];

sd.immune.a1.curr=sqrt(sum(prob.immune.subject*aggressive.treatment*(y-mean.immune.a0.curr-itt.immune.treatprotected.curr+x*xeffect.curr)^2)/sum(prob.immune.subject*aggressive.treatment))
sd.immune.a0.curr=sqrt(sum(prob.immune.subject*(1-aggressive.treatment)*(y-mean.immune.a0.curr-x*xeffect.curr)^2)/sum(prob.immune.subject*(1-aggressive.treatment)));
sd.treatprotected.a1.curr=sqrt(sum(prob.treatprotected.subject*aggressive.treatment*(y-mean.treatprotected.a0.curr-itt.immune.treatprotected.curr-x*xeffect.curr)^2)/sum(prob.treatprotected.subject*aggressive.treatment));
sd.treatprotected.a0.curr=sqrt(sum(prob.treatprotected.subject*(1-aggressive.treatment)*(y-mean.treatprotected.a0.curr-x*xeffect.curr)^2)/sum(prob.treatprotected.subject*(1-aggressive.treatment)));
sd.doomed.a1.curr=sqrt(sum(prob.doomed.subject*aggressive.treatment*(y-mean.doomed.a0.curr-itt.treatharmful.doomed.curr-x*xeffect.curr)^2)/sum(prob.doomed.subject*aggressive.treatment));
sd.doomed.a0.curr=sqrt(sum(prob.doomed.subject*(1-aggressive.treatment)*(y-mean.doomed.a0.curr-x*xeffect.curr)^2)/sum(prob.doomed.subject*(1-aggressive.treatment)));
sd.treatharmful.a1.curr=sqrt(sum(prob.treatharmful.subject*aggressive.treatment*(y-mean.treatharmful.a0.curr-itt.treatharmful.doomed.curr-x*xeffect.curr)^2)/sum(prob.treatharmful.subject*aggressive.treatment));
sd.treatharmful.a0.curr=sqrt(sum(prob.treatharmful.subject*(1-aggressive.treatment)*(y-mean.treatharmful.a0.curr-x*xeffect.curr)^2)/sum(prob.treatharmful.subject*(1-aggressive.treatment)));
  
prob.immune.x1.curr=sum(prob.immune.subject*x)/sum(x);
prob.treatprotected.x1.curr=sum(prob.treatprotected.subject*x)/sum(x);
prob.treatharmful.x1.curr=sum(prob.treatharmful.subject*x)/sum(x);
prob.doomed.x1.curr=sum(prob.doomed.subject*x)/sum(x);
prob.immune.x0.curr=sum(prob.immune.subject*(1-x))/(n-sum(x));
prob.treatprotected.x0.curr=sum(prob.treatprotected.subject*(1-x))/(n-sum(x));
prob.treatharmful.x0.curr=sum(prob.treatharmful.subject*(1-x))/(n-sum(x));
prob.doomed.x0.curr=sum(prob.doomed.subject*(1-x))/(n-sum(x));

# Compute log likelihood with updated parameters
  ll.new.p1a=sum((1-aggressive.treatment)*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*esrd*x)*log(prob.doomed.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*esrd*x)*log(prob.doomed.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x1.curr)/(prob.immune.x1.curr+prob.treatharmful.x1.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+xeffect.curr,sd=sd.treatharmful.a0.curr))));
  ll.new.p1b=sum((1-aggressive.treatment)*esrd*x*log(((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+xeffect.curr,sd=sd.treatprotected.a0.curr))));
ll.new.p2=sum(aggressive.treatment*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.immune.a0.curr+itt.immune.treatprotected.curr+xeffect.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.immune.x1.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+itt.immune.treatprotected.curr+xeffect.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*x*((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.doomed.a0.curr+itt.treatharmful.doomed.curr+xeffect.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x1.curr)/(prob.treatharmful.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+itt.treatharmful.doomed.curr+xeffect.curr,sd=sd.treatharmful.a1.curr)));
ll.new.p3a=sum((1-aggressive.treatment)*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x0.curr)/(prob.immune.x0.curr+prob.treatharmful.x0.curr)*dnorm(y,mean=mean.treatharmful.a0.curr,sd=sd.treatharmful.a0.curr))))
ll.new.p3b=sum((1-aggressive.treatment)*esrd*(1-x)*log(((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatprotected.a0.curr,sd=sd.treatprotected.a0.curr))));
ll.new.p4=sum(aggressive.treatment*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.immune.a0.curr+itt.immune.treatprotected.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.immune.x0.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+itt.immune.treatprotected.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*(1-x)*((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.doomed.a0.curr+itt.treatharmful.doomed.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x0.curr)/(prob.treatharmful.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+itt.treatharmful.doomed.curr,sd=sd.treatharmful.a1.curr)));
ll.new=ll.new.p1a+ll.new.p1b+ll.new.p2+ll.new.p3a+ll.new.p3b+ll.new.p4;

# If log likelihood has become infinite, mark that estimate has not converged
# and restart simulation

# If log likelihood has become infinite, mark that estimate has not converged
# and restart simulation
if(is.na(ll.new)==FALSE){
# If any of the standard deviations has become zero, mark that estimate
# has not converged and restart simulation
if(sd.immune.a1.curr<.00001 | sd.immune.a0.curr<.00001 | sd.treatprotected.a1.curr<.00001 | sd.treatprotected.a0.curr<.00001 | sd.doomed.a1.curr<.00001 | sd.doomed.a0.curr<.00001 | sd.treatharmful.a1.curr<.00001 | sd.treatharmful.a0.curr<.00001){
  no.notconv=no.notconv+1
  conv=0;
  
}
if(sd.immune.a1.curr>=.00001 & sd.immune.a0.curr>=.00001 & sd.treatprotected.a1.curr>=.00001 & sd.treatprotected.a0.curr>=.00001 & sd.doomed.a1.curr>=.00001 & sd.doomed.a0.curr>=.00001 & sd.treatharmful.a1.curr>=.00001 & sd.treatharmful.a0.curr>=.00001){
lldist=ll.new-ll.old;
iter=iter+1;
}
}


if(is.na(ll.new)==TRUE){
  no.notconv=no.notconv+1;
  conv=0;
  lldist=0;
}
}

if(conv==1){
itt.immune.treatprotected.est.2[k]=itt.immune.treatprotected.curr;
mean.immune.a0.est.2[k]=mean.immune.a0.curr;
mean.treatprotected.a0.est.2[k]=mean.treatprotected.a0.curr;
itt.treatharmful.doomed.est.2[k]=itt.treatharmful.doomed.curr;
mean.treatharmful.a0.est.2[k]=mean.treatharmful.a0.curr;
mean.doomed.a0.est.2[k]=mean.doomed.a0.curr;
sd.immune.a1.est.2[k]=sd.immune.a1.curr;
sd.immune.a0.est.2[k]=sd.immune.a0.curr;
sd.treatprotected.a1.est.2[k]=sd.treatprotected.a1.curr;
sd.treatprotected.a0.est.2[k]=sd.treatprotected.a0.curr;
sd.treatharmful.a1.est.2[k]=sd.treatharmful.a1.curr;
sd.treatharmful.a0.est.2[k]=sd.treatharmful.a0.curr;
sd.doomed.a1.est.2[k]=sd.doomed.a1.curr;
sd.doomed.a0.est.2[k]=sd.doomed.a0.curr;
xeffect.est.2[k]=xeffect.curr;
lldistvec.2[k]=lldist;
k=k+1;
}
}

conv=as.logical((1-is.na(mean.immune.a1.est.1-mean.immune.a0.est.1))*(1-is.na(mean.treatprotected.a1.est.1-mean.treatprotected.a0.est.1))*(1-is.na(mean.treatharmful.a1.est.1-mean.treatharmful.a0.est.1))*(1-is.na(mean.doomed.a1.est.1-mean.treatharmful.a0.est.1))*(1-is.na(psi0.est))*(1-is.na(psi1.est))*(1-is.na(itt.immune.treatprotected.est.2))*(1-is.na(itt.treatharmful.doomed.est.2)));
sum(conv)
mean(mean.immune.a1.est.1[conv]-mean.immune.a0.est.1[conv])
sd(mean.immune.a1.est.1[conv]-mean.immune.a0.est.1[conv])
mean(mean.treatprotected.a1.est.1[conv]-mean.treatprotected.a0.est.1[conv])
sd(mean.treatprotected.a1.est.1[conv]-mean.treatprotected.a0.est.1[conv])
mean(mean.treatharmful.a1.est.1[conv]-mean.treatharmful.a0.est.1[conv])
sd(mean.treatharmful.a1.est.1[conv]-mean.treatharmful.a0.est.1[conv])
mean(mean.doomed.a1.est.1[conv]-mean.doomed.a0.est.1[conv])
sd(mean.doomed.a1.est.1[conv]-mean.doomed.a0.est.1[conv])
mean(psi0.est[conv])
sd(psi0.est[conv])
mean(psi1.est[conv])
sd(psi1.est[conv])
mean(betahat.a[conv]+betahat.mua[conv])
# sd(betahat.a+betahat.mua)
mean(itt.immune.treatprotected.est.2[conv])
sd(itt.immune.treatprotected.est.2[conv])
mean(itt.treatharmful.doomed.est.2[conv])
sd(itt.treatharmful.doomed.est.2[conv])


# Setting II(A)

# Simulation study
sims=1; # Number of simulations
# Vectors to store estimates from simulation study
mean.immune.a1.est.1=rep(0,sims);
mean.immune.a0.est.1=rep(0,sims);
mean.treatprotected.a1.est.1=rep(0,sims);
mean.treatprotected.a0.est.1=rep(0,sims);
mean.treatharmful.a1.est.1=rep(0,sims);
mean.treatharmful.a0.est.1=rep(0,sims);
mean.doomed.a1.est.1=rep(0,sims);
mean.doomed.a0.est.1=rep(0,sims);
sd.immune.a1.est.1=rep(0,sims);
sd.immune.a0.est.1=rep(0,sims);
sd.treatprotected.a1.est.1=rep(0,sims);
sd.treatprotected.a0.est.1=rep(0,sims);
sd.treatharmful.a1.est.1=rep(0,sims);
sd.treatharmful.a0.est.1=rep(0,sims);
sd.doomed.a1.est.1=rep(0,sims);
sd.doomed.a0.est.1=rep(0,sims);
prob.immune.est.1=rep(0,sims);
prob.treatprotected.est.1=rep(0,sims);
prob.treatharmful.est.1=rep(0,sims);
prob.doomed.est.1=rep(0,sims);
xeffect.est.1=rep(0,sims);
psi0.est=rep(0,sims); # Structural nested model coefficients
psi1.est=rep(0,sims); 
betahat.a=rep(0,sims);
betahat.mua=rep(0,sims);
lldistvec.1=rep(0,sims);
no.notconv=0; # Number of times EM algorithm does not converge
itt.immune.treatprotected.est.2=rep(0,sims);
mean.immune.a0.est.2=rep(0,sims);
mean.treatprotected.a0.est.2=rep(0,sims);
itt.treatharmful.doomed.est.2=rep(0,sims);
mean.treatharmful.a0.est.2=rep(0,sims);
mean.doomed.a0.est.2=rep(0,sims);
sd.immune.a1.est.2=rep(0,sims);
sd.immune.a0.est.2=rep(0,sims);
sd.treatprotected.a1.est.2=rep(0,sims);
sd.treatprotected.a0.est.2=rep(0,sims);
sd.treatharmful.a1.est.2=rep(0,sims);
sd.treatharmful.a0.est.2=rep(0,sims);
sd.doomed.a1.est.2=rep(0,sims);
sd.doomed.a0.est.2=rep(0,sims);
xeffect.est.2=rep(0,sims);
lldistvec.2=rep(0,sims);

k=1;
while(k<=sims){
# Simulation from model with a binary covariate
# Probabilities of principal strata
prob.immune=.25;
prob.treatprotected=.4;
prob.treatharmful=.05;
prob.doomed=.3;

# Means and standard deviations of each principal strata under aggressive treatment
# for hypertension and under not aggressive treatment for hypertension
mean.immune.a1=2; # Mean of immune subjects when aggressively treated for hypertension
sd.immune.a1=1; # Variance of immune subjects when aggressively treated
mean.immune.a0=1; # Mean of immune subjects when not agressively treated
sd.immune.a0=1; # Variance of immune subjects when not aggressively treated
mean.treatprotected.a1=2.5;
sd.treatprotected.a1=1;
mean.treatprotected.a0=1;
sd.treatprotected.a0=1;
mean.treatharmful.a1=1.25;
sd.treatharmful.a1=1;
mean.treatharmful.a0=.75;
sd.treatharmful.a0=1;
mean.doomed.a1=2.25;
sd.doomed.a1=1;
mean.doomed.a0=1.25;
sd.doomed.a0=.25;
xeffect=.5; # Addition to mean for binary effect

# Probabilities of x for different principal strata
probx1.immune=.5;
probx1.treatprotected=.75;
probx1.treatharmful=.25;
probx1.doomed=.5;
totalprobx=probx1.immune*prob.immune+probx1.treatprotected*prob.treatprotected+probx1.treatharmful*prob.treatharmful+probx1.doomed*prob.doomed;

n=500; # total sample size
# Assume that half of subjects are randomized into aggressive treatment
# of hypertension and half into nonaggressive treatment
aggressive.treatment=c(rep(1,n/2),rep(0,n/2));
principalstrata=rmultinom(n,1,c(prob.immune,prob.treatprotected,prob.treatharmful,prob.doomed));
# Dummy variables for principal strata
principalstrata.immune=principalstrata[1,];
principalstrata.treatprotected=principalstrata[2,];
principalstrata.treatharmful=principalstrata[3,];
principalstrata.doomed=principalstrata[4,];
# Covariate X
x=principalstrata.immune*rbinom(n,1,probx1.immune)+principalstrata.treatprotected*rbinom(n,1,probx1.treatprotected)+principalstrata.treatharmful*rbinom(n,1,probx1.treatharmful)+principalstrata.doomed*rbinom(n,1,probx1.doomed);

# ESRD dummy
esrd=principalstrata.doomed+principalstrata.treatprotected*(1-aggressive.treatment)+principalstrata.treatharmful*aggressive.treatment;
# Mean and standard deviation of outcome
mean.outcome=mean.immune.a1*principalstrata.immune*aggressive.treatment+mean.immune.a0*principalstrata.immune*(1-aggressive.treatment)+mean.treatprotected.a1*principalstrata.treatprotected*aggressive.treatment+mean.treatprotected.a0*principalstrata.treatprotected*(1-aggressive.treatment)+mean.treatharmful.a1*principalstrata.treatharmful*aggressive.treatment+mean.treatharmful.a0*principalstrata.treatharmful*(1-aggressive.treatment)+mean.doomed.a1*principalstrata.doomed*aggressive.treatment+mean.doomed.a0*principalstrata.doomed*(1-aggressive.treatment)+xeffect*x;

sd.outcome=sd.immune.a1*principalstrata.immune*aggressive.treatment+sd.immune.a0*principalstrata.immune*(1-aggressive.treatment)+sd.treatprotected.a1*principalstrata.treatprotected*aggressive.treatment+sd.treatprotected.a0*principalstrata.treatprotected*(1-aggressive.treatment)+sd.treatharmful.a1*principalstrata.treatharmful*aggressive.treatment+sd.treatharmful.a0*principalstrata.treatharmful*(1-aggressive.treatment)+sd.doomed.a1*principalstrata.doomed*aggressive.treatment+sd.doomed.a0*principalstrata.doomed*(1-aggressive.treatment);

# Generate outcomes from normal distribution
y=rnorm(n,mean=mean.outcome,sd=sd.outcome);
# Generate outcomes from gamma distribution
# scalegamma=sd.outcome^2/mean.outcome;
# shapegamma=mean.outcome/scalegamma;
# y=rgamma(n,shape=shapegamma,scale=scalegamma);

# EM algorithm

# Select starting values
# For now, start at truth
prob.immune.curr=prob.immune;
prob.treatprotected.curr=prob.treatprotected;
prob.treatharmful.curr=prob.treatharmful;
prob.doomed.curr=prob.doomed;
prob.immune.x1.curr=(prob.immune*probx1.immune)/totalprobx;
prob.treatprotected.x1.curr=(prob.treatprotected*probx1.treatprotected)/totalprobx;
prob.treatharmful.x1.curr=(prob.treatharmful*probx1.treatharmful)/totalprobx;
prob.doomed.x1.curr=(prob.doomed*probx1.doomed)/totalprobx;
prob.immune.x0.curr=(prob.immune*(1-probx1.immune))/(1-totalprobx);
prob.treatprotected.x0.curr=(prob.treatprotected*(1-probx1.treatprotected))/(1-totalprobx);
prob.treatharmful.x0.curr=(prob.treatharmful*(1-probx1.treatharmful))/(1-totalprobx);
prob.doomed.x0.curr=(prob.doomed*(1-probx1.doomed))/(1-totalprobx);
mean.immune.a1.curr=mean.immune.a1;
sd.immune.a1.curr=sd.immune.a1; 
mean.immune.a0.curr=mean.immune.a0; 
sd.immune.a0.curr=sd.immune.a0;
mean.treatprotected.a1.curr=mean.treatprotected.a1;
sd.treatprotected.a1.curr=sd.treatprotected.a1;
mean.treatprotected.a0.curr=mean.treatprotected.a0;
sd.treatprotected.a0.curr=sd.treatprotected.a0;
mean.treatharmful.a1.curr=mean.treatharmful.a1;
sd.treatharmful.a1.curr=sd.treatharmful.a1;
mean.treatharmful.a0.curr=mean.treatharmful.a0;
sd.treatharmful.a0.curr=sd.treatharmful.a0;
mean.doomed.a1.curr=mean.doomed.a1;
sd.doomed.a1.curr=sd.doomed.a1;
mean.doomed.a0.curr=mean.doomed.a0;
sd.doomed.a0.curr=sd.doomed.a0;
xeffect.curr=xeffect;

maxiters=500; # Maximum iterations in EM algorithm
iter=1;
lldist=1; # Will measure difference in log likelihood between iterations
toler=.001; # Tolerance for difference in log likelihood between iterations
# for stopping EM algorithm
conv=1;

while(lldist>toler && iter<maxiters){
  # Compute log likelihood with current parameters
ll.old.p1a=sum((1-aggressive.treatment)*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*esrd*x)*log(prob.doomed.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*esrd*x)*log(prob.doomed.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x1.curr)/(prob.immune.x1.curr+prob.treatharmful.x1.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+xeffect.curr,sd=sd.treatharmful.a0.curr))));
  ll.old.p1b=sum((1-aggressive.treatment)*esrd*x*log(((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+xeffect.curr,sd=sd.treatprotected.a0.curr))));
ll.old.p2=sum(aggressive.treatment*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.immune.a1.curr+xeffect.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.immune.x1.curr)*dnorm(y,mean=mean.treatprotected.a1.curr+xeffect.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*x*((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.doomed.a1.curr+xeffect.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x1.curr)/(prob.treatharmful.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatharmful.a1.curr+xeffect.curr,sd=sd.treatharmful.a1.curr)));
ll.old.p3a=sum((1-aggressive.treatment)*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x0.curr)/(prob.immune.x0.curr+prob.treatharmful.x0.curr)*dnorm(y,mean=mean.treatharmful.a0.curr,sd=sd.treatharmful.a0.curr))))
ll.old.p3b=sum((1-aggressive.treatment)*esrd*(1-x)*log(((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatprotected.a0.curr,sd=sd.treatprotected.a0.curr))));
ll.old.p4=sum(aggressive.treatment*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.immune.a1.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.immune.x0.curr)*dnorm(y,mean=mean.treatprotected.a1.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*(1-x)*((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.doomed.a1.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x0.curr)/(prob.treatharmful.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatharmful.a1.curr,sd=sd.treatharmful.a1.curr)));
ll.old=ll.old.p1a+ll.old.p1b+ll.old.p2+ll.old.p3a+ll.old.p3b+ll.old.p4;

# Expectation step
# Compute probabilities for each subject's principal strata
prob.immune.subject=rep(0,n);
prob.treatprotected.subject=rep(0,n);
prob.treatharmful.subject=rep(0,n);
prob.doomed.subject=rep(0,n);
for(j in 1:n){
  if(aggressive.treatment[j]==0 && esrd[j]==0 && x[j]==1){
    prob.immune.subject[j]=prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)/(prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)+prob.treatharmful.x1.curr*dnorm(y[j],mean=mean.treatharmful.a0.curr+xeffect.curr,sd=sd.treatharmful.a0.curr));
    prob.treatharmful.subject[j]=1-prob.immune.subject[j];
  }
  if(aggressive.treatment[j]==0 && esrd[j]==0 && x[j]==0){
  prob.immune.subject[j]=prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)/(prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)+prob.treatharmful.x0.curr*dnorm(y[j],mean=mean.treatharmful.a0.curr,sd=sd.treatharmful.a0.curr));
    prob.treatharmful.subject[j]=1-prob.immune.subject[j];
  }  
  if(aggressive.treatment[j]==0 && esrd[j]==1 && x[j]==1){
    prob.doomed.subject[j]=prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)/(prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)+prob.treatprotected.x1.curr*dnorm(y[j],mean=mean.treatprotected.a0.curr+xeffect.curr,sd=sd.treatprotected.a0.curr));
    prob.treatprotected.subject[j]=1-prob.doomed.subject[j];
  }
  if(aggressive.treatment[j]==0 && esrd[j]==1 && x[j]==0){
    prob.doomed.subject[j]=prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)/(prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)+prob.treatprotected.x0.curr*dnorm(y[j],mean=mean.treatprotected.a0.curr,sd=sd.treatprotected.a0.curr));
    prob.treatprotected.subject[j]=1-prob.doomed.subject[j];
  }
  if(aggressive.treatment[j]==1 && esrd[j]==0 && x[j]==1){
    prob.immune.subject[j]=prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a1.curr+xeffect.curr,sd=sd.immune.a1.curr)/(prob.immune.x1.curr*dnorm(y[j],mean=mean.immune.a1.curr+xeffect.curr,sd=sd.immune.a1.curr)+prob.treatprotected.x1.curr*dnorm(y[j],mean=mean.treatprotected.a1.curr+xeffect.curr,sd=sd.treatprotected.a1.curr));
    prob.treatprotected.subject[j]=1-prob.immune.subject[j];
  }
  if(aggressive.treatment[j]==1 && esrd[j]==0 && x[j]==0){
    prob.immune.subject[j]=prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a1.curr,sd=sd.immune.a1.curr)/(prob.immune.x0.curr*dnorm(y[j],mean=mean.immune.a1.curr,sd=sd.immune.a1.curr)+prob.treatprotected.x0.curr*dnorm(y[j],mean=mean.treatprotected.a1.curr,sd=sd.treatprotected.a1.curr));
    prob.treatprotected.subject[j]=1-prob.immune.subject[j];
  }  
  if(aggressive.treatment[j]==1 && esrd[j]==1 && x[j]==1){
    prob.doomed.subject[j]=prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a1.curr+xeffect.curr,sd=sd.doomed.a1.curr)/(prob.doomed.x1.curr*dnorm(y[j],mean=mean.doomed.a1.curr+xeffect.curr,sd=sd.doomed.a1.curr)+prob.treatharmful.x1.curr*dnorm(y[j],mean=mean.treatharmful.a1.curr+xeffect.curr,sd=sd.treatharmful.a1.curr));
    prob.treatharmful.subject[j]=1-prob.doomed.subject[j];
  }
  if(aggressive.treatment[j]==1 && esrd[j]==1 && x[j]==0){
    prob.doomed.subject[j]=prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a1.curr,sd=sd.doomed.a1.curr)/(prob.doomed.x0.curr*dnorm(y[j],mean=mean.doomed.a1.curr,sd=sd.doomed.a1.curr)+prob.treatharmful.x0.curr*dnorm(y[j],mean=mean.treatharmful.a1.curr,sd=sd.treatharmful.a1.curr));
    prob.treatharmful.subject[j]=1-prob.doomed.subject[j];
  }  
}
 
# Maximization Step (Use Conditional Maximization, i.e., ECM algorithm)
# First maximize expected likelihood over xeffect, principal strata means
# and then maximize likelihood over standard deviations in each principal
# stratra
# Set up weighted regression
weighty=rep(0,2*n);
weightx=rep(0,2*n);
weight.immune.a1=rep(0,2*n);
weight.immune.a0=rep(0,2*n);
weight.treatprotected.a1=rep(0,2*n);
weight.treatprotected.a0=rep(0,2*n);
weight.treatharmful.a1=rep(0,2*n);
weight.treatharmful.a0=rep(0,2*n);
weight.doomed.a1=rep(0,2*n);
weight.doomed.a0=rep(0,2*n);

for(j in 1:n){
  if(aggressive.treatment[j]==0 && esrd[j]==0){
  tempweight1=(prob.immune.subject[j]/(prob.immune.subject[j]+prob.treatharmful.subject[j]))^.5/sd.immune.a0.curr;
  weighty[((j-1)*2+1)]=y[j]*tempweight1;
  weightx[((j-1)*2+1)]=x[j]*tempweight1;
  weight.immune.a0[((j-1)*2+1)]=tempweight1;
  tempweight2=(prob.treatharmful.subject[j]/(prob.immune.subject[j]+prob.treatharmful.subject[j]))^.5/sd.treatharmful.a0.curr;
  weighty[(j*2)]=y[j]*tempweight2;
  weightx[(j*2)]=x[j]*tempweight2;
  weight.treatharmful.a0[(j*2)]=tempweight2;
}
  if(aggressive.treatment[j]==0 && esrd[j]==1){
  tempweight1=(prob.doomed.subject[j]/(prob.doomed.subject[j]+prob.treatprotected.subject[j]))^.5/sd.doomed.a0.curr;
  weighty[((j-1)*2+1)]=y[j]*tempweight1;
  weightx[((j-1)*2+1)]=x[j]*tempweight1;
  weight.doomed.a0[((j-1)*2+1)]=tempweight1;
  tempweight2=(prob.treatprotected.subject[j]/(prob.doomed.subject[j]+prob.treatprotected.subject[j]))^.5/sd.treatprotected.a0.curr;
  weighty[(j*2)]=y[j]*tempweight2;
  weightx[(j*2)]=x[j]*tempweight2;
  weight.treatprotected.a0[(j*2)]=tempweight2;
}
  if(aggressive.treatment[j]==1 && esrd[j]==0){
    tempweight1=(prob.immune.subject[j]/(prob.immune.subject[j]+prob.treatprotected.subject[j]))^.5/sd.immune.a1.curr;
    weighty[((j-1)*2+1)]=y[j]*tempweight1;
    weightx[((j-1)*2+1)]=x[j]*tempweight1;
    weight.immune.a1[((j-1)*2+1)]=tempweight1;
    tempweight2=(prob.treatprotected.subject[j]/(prob.immune.subject[j]+prob.treatprotected.subject[j]))^.5/sd.treatprotected.a1.curr;
    weighty[(j*2)]=y[j]*tempweight2;
    weightx[(j*2)]=x[j]*tempweight2;
    weight.treatprotected.a1[(j*2)]=tempweight2;
  }
  if(aggressive.treatment[j]==1 && esrd[j]==1){
    tempweight1=(prob.doomed.subject[j]/(prob.doomed.subject[j]+prob.treatharmful.subject[j]))^.5/sd.doomed.a1.curr;
    weighty[((j-1)*2+1)]=y[j]*tempweight1;
    weightx[((j-1)*2+1)]=x[j]*tempweight1;
    weight.doomed.a1[((j-1)*2+1)]=tempweight1;
    tempweight2=(prob.treatharmful.subject[j]/(prob.treatharmful.subject[j]+prob.doomed.subject[j]))^.5/sd.treatharmful.a1.curr;
    weighty[(j*2)]=y[j]*tempweight2;
    weightx[(j*2)]=x[j]*tempweight2;
    weight.treatharmful.a1[(j*2)]=tempweight2;
  }
}
  
weightreg=lm(weighty~weightx+weight.immune.a0+weight.immune.a1+weight.treatharmful.a0+weight.treatharmful.a1+weight.doomed.a0+weight.doomed.a1+weight.treatprotected.a0+weight.treatprotected.a1-1);
coef.weightreg=coef(weightreg);
  xeffect.curr=coef.weightreg[1];
  mean.immune.a0.curr=coef.weightreg[2];
  mean.immune.a1.curr=coef.weightreg[3];
  mean.treatharmful.a0.curr=coef.weightreg[4];
  mean.treatharmful.a1.curr=coef.weightreg[5];
  mean.doomed.a0.curr=coef.weightreg[6];
  mean.doomed.a1.curr=coef.weightreg[7];
  mean.treatprotected.a0.curr=coef.weightreg[8];
  mean.treatprotected.a1.curr=coef.weightreg[9];

sd.immune.a1.curr=sqrt(sum(prob.immune.subject*aggressive.treatment*(y-mean.immune.a1.curr-x*xeffect.curr)^2)/sum(prob.immune.subject*aggressive.treatment))
sd.immune.a0.curr=sqrt(sum(prob.immune.subject*(1-aggressive.treatment)*(y-mean.immune.a0.curr-x*xeffect.curr)^2)/sum(prob.immune.subject*(1-aggressive.treatment)));
sd.treatprotected.a1.curr=sqrt(sum(prob.treatprotected.subject*aggressive.treatment*(y-mean.treatprotected.a1.curr-x*xeffect.curr)^2)/sum(prob.treatprotected.subject*aggressive.treatment));
sd.treatprotected.a0.curr=sqrt(sum(prob.treatprotected.subject*(1-aggressive.treatment)*(y-mean.treatprotected.a0.curr-x*xeffect.curr)^2)/sum(prob.treatprotected.subject*(1-aggressive.treatment)));
sd.doomed.a1.curr=sqrt(sum(prob.doomed.subject*aggressive.treatment*(y-mean.doomed.a1.curr-x*xeffect.curr)^2)/sum(prob.doomed.subject*aggressive.treatment));
sd.doomed.a0.curr=sqrt(sum(prob.doomed.subject*(1-aggressive.treatment)*(y-mean.doomed.a0.curr-x*xeffect.curr)^2)/sum(prob.doomed.subject*(1-aggressive.treatment)));
sd.treatharmful.a1.curr=sqrt(sum(prob.treatharmful.subject*aggressive.treatment*(y-mean.treatharmful.a1.curr-x*xeffect.curr)^2)/sum(prob.treatharmful.subject*aggressive.treatment));
sd.treatharmful.a0.curr=sqrt(sum(prob.treatharmful.subject*(1-aggressive.treatment)*(y-mean.treatharmful.a0.curr-x*xeffect.curr)^2)/sum(prob.treatharmful.subject*(1-aggressive.treatment)));
  
prob.immune.x1.curr=sum(prob.immune.subject*x)/sum(x);
prob.treatprotected.x1.curr=sum(prob.treatprotected.subject*x)/sum(x);
prob.treatharmful.x1.curr=sum(prob.treatharmful.subject*x)/sum(x);
prob.doomed.x1.curr=sum(prob.doomed.subject*x)/sum(x);
prob.immune.x0.curr=sum(prob.immune.subject*(1-x))/(n-sum(x));
prob.treatprotected.x0.curr=sum(prob.treatprotected.subject*(1-x))/(n-sum(x));
prob.treatharmful.x0.curr=sum(prob.treatharmful.subject*(1-x))/(n-sum(x));
prob.doomed.x0.curr=sum(prob.doomed.subject*(1-x))/(n-sum(x));

# Compute log likelihood with updated parameters
  ll.new.p1a=sum((1-aggressive.treatment)*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*esrd*x)*log(prob.doomed.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*esrd*x)*log(prob.doomed.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x1.curr)/(prob.immune.x1.curr+prob.treatharmful.x1.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+xeffect.curr,sd=sd.treatharmful.a0.curr))));
  ll.new.p1b=sum((1-aggressive.treatment)*esrd*x*log(((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+xeffect.curr,sd=sd.treatprotected.a0.curr))));
ll.new.p2=sum(aggressive.treatment*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.immune.a1.curr+xeffect.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.immune.x1.curr)*dnorm(y,mean=mean.treatprotected.a1.curr+xeffect.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*x*((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.doomed.a1.curr+xeffect.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x1.curr)/(prob.treatharmful.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatharmful.a1.curr+xeffect.curr,sd=sd.treatharmful.a1.curr)));
ll.new.p3a=sum((1-aggressive.treatment)*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*(1-esrd)*(1-x))*log(prob.immune.x0.curr+prob.treatprotected.x0.curr)+sum(aggressive.treatment*esrd*(1-x))*log(prob.doomed.x0.curr+prob.treatharmful.x0.curr)+sum((1-aggressive.treatment)*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.immune.a0.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x0.curr)/(prob.immune.x0.curr+prob.treatharmful.x0.curr)*dnorm(y,mean=mean.treatharmful.a0.curr,sd=sd.treatharmful.a0.curr))))
ll.new.p3b=sum((1-aggressive.treatment)*esrd*(1-x)*log(((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.doomed.a0.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatprotected.a0.curr,sd=sd.treatprotected.a0.curr))));
ll.new.p4=sum(aggressive.treatment*(1-esrd)*(1-x)*log(((prob.immune.x0.curr/(prob.immune.x0.curr+prob.treatprotected.x0.curr))*dnorm(y,mean=mean.immune.a1.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x0.curr)/(prob.treatprotected.x0.curr+prob.immune.x0.curr)*dnorm(y,mean=mean.treatprotected.a1.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*(1-x)*((prob.doomed.x0.curr/(prob.doomed.x0.curr+prob.treatharmful.x0.curr))*dnorm(y,mean=mean.doomed.a1.curr,sd=sd.doomed.a1.curr)+(prob.treatharmful.x0.curr)/(prob.treatharmful.x0.curr+prob.doomed.x0.curr)*dnorm(y,mean=mean.treatharmful.a1.curr,sd=sd.treatharmful.a1.curr)));
ll.new=ll.new.p1a+ll.new.p1b+ll.new.p2+ll.new.p3a+ll.new.p3b+ll.new.p4;


# If log likelihood has become infinite, mark that estimate has not converged
# and restart simulation
if(is.na(ll.new)==FALSE){
# If any of the standard deviations has become zero, mark that estimate
# has not converged and restart simulation
if(sd.immune.a1.curr<.00001 | sd.immune.a0.curr<.00001 | sd.treatprotected.a1.curr<.00001 | sd.treatprotected.a0.curr<.00001 | sd.doomed.a1.curr<.00001 | sd.doomed.a0.curr<.00001 | sd.treatharmful.a1.curr<.00001 | sd.treatharmful.a0.curr<.00001){
  no.notconv=no.notconv+1
  conv=0;
  
}
if(sd.immune.a1.curr>=.00001 & sd.immune.a0.curr>=.00001 & sd.treatprotected.a1.curr>=.00001 & sd.treatprotected.a0.curr>=.00001 & sd.doomed.a1.curr>=.00001 & sd.doomed.a0.curr>=.00001 & sd.treatharmful.a1.curr>=.00001 & sd.treatharmful.a0.curr>=.00001){
lldist=ll.new-ll.old;
iter=iter+1;
}
}

if(is.na(ll.new)==TRUE){
  no.notconv=no.notconv+1;
  conv=0;
  lldist=0;
}
}

mean.immune.a1.est.1[k]=mean.immune.a1.curr;
mean.immune.a0.est.1[k]=mean.immune.a0.curr;
mean.treatprotected.a1.est.1[k]=mean.treatprotected.a1.curr;
mean.treatprotected.a0.est.1[k]=mean.treatprotected.a0.curr;
mean.treatharmful.a1.est.1[k]=mean.treatharmful.a1.curr;
mean.treatharmful.a0.est.1[k]=mean.treatharmful.a0.curr;
mean.doomed.a1.est.1[k]=mean.doomed.a1.curr;
mean.doomed.a0.est.1[k]=mean.doomed.a0.curr;
sd.immune.a1.est.1[k]=sd.immune.a1.curr;
sd.immune.a0.est.1[k]=sd.immune.a0.curr;
sd.treatprotected.a1.est.1[k]=sd.treatprotected.a1.curr;
sd.treatprotected.a0.est.1[k]=sd.treatprotected.a0.curr;
sd.treatharmful.a1.est.1[k]=sd.treatharmful.a1.curr;
sd.treatharmful.a0.est.1[k]=sd.treatharmful.a0.curr;
sd.doomed.a1.est.1[k]=sd.doomed.a1.curr;
sd.doomed.a0.est.1[k]=sd.doomed.a0.curr;
prob.immune.est.1[k]=prob.immune.curr;
prob.treatprotected.est.1[k]=prob.treatprotected.curr;
prob.treatharmful.est.1[k]=prob.treatharmful.curr;
prob.doomed.est.1[k]=prob.doomed.curr;
xeffect.est.1[k]=xeffect.curr;
lldistvec.1[k]=lldist;

# Expected auxiliary stratification estimates
ax=aggressive.treatment*x;
auxreg1=glm(esrd~x+aggressive.treatment+ax,family=binomial);
muhatax=predict.glm(auxreg1,type="response");
a.times.muhatax=aggressive.treatment*muhatax;
auxreg2=lm(y~muhatax+aggressive.treatment+a.times.muhatax);
betahat.a[k]=coef(auxreg2)[3];
betahat.mua[k]=coef(auxreg2)[4];

# Structural nested model estimates
# Use Newton's Method
# Start at true values
psi1.curr=mean.doomed.a1-mean.doomed.a0;
psi0.curr=mean.treatprotected.a1-mean.treatprotected.a0;

# Estimate "mediating score"
tempreg=glm(esrd~x,family=binomial,subset=(aggressive.treatment==1));
mediatingscore=expit(coef(tempreg)[1]+coef(tempreg)[2]*x);

maxiters=500; # Maximum iterations in Newton-Raphson algorithm
iter=1;
dist=1;
toler=.001;
while(iter<maxiters && dist>toler){
  psiold=matrix(c(psi0.curr,psi1.curr),ncol=1);
  y0psi=y-psi0.curr*aggressive.treatment*(1-esrd)-psi1.curr*aggressive.treatment*esrd;
  tempreg=lm(y0psi~x);
  tempbetahat0=coef(tempreg)[1];
  tempbetahat1=coef(tempreg)[2];
  g1=sum((aggressive.treatment-.5)*(y0psi-tempbetahat0-tempbetahat1*x)*(1-mediatingscore));
  g2=sum((aggressive.treatment-.5)*(y0psi-tempbetahat0-tempbetahat1*x)*mediatingscore);
  Jmat=matrix(rep(0,4),ncol=2);
  Jmat[1,1]=sum(-(aggressive.treatment-.5)*aggressive.treatment*(1-esrd)*(1-mediatingscore));
  Jmat[1,2]=sum(-(aggressive.treatment-.5)*aggressive.treatment*esrd*(1-mediatingscore));
  Jmat[2,1]=sum(-(aggressive.treatment-.5)*aggressive.treatment*(1-esrd)*mediatingscore);
  Jmat[2,2]=sum(-(aggressive.treatment-.5)*aggressive.treatment*esrd*mediatingscore);
  psinew=psiold-solve(Jmat)%*%matrix(c(g1,g2),ncol=1);
  dist=(sum((psinew-psiold)^2))^.5;
  psi0.curr=psinew[1];
  psi1.curr=psinew[2];
  iter=iter+1;
}
psi0.est[k]=psi0.curr;
psi1.est[k]=psi1.curr;


# Estimates that assume ITT(immune)=ITT(treatment protected)
#                  ITT(treatment harmful)=ITT(doomed)
# EM algorithm

# Select starting values
# For now, start at truth
prob.immune.curr=prob.immune;
prob.treatprotected.curr=prob.treatprotected;
prob.treatharmful.curr=prob.treatharmful;
prob.doomed.curr=prob.doomed;
prob.immune.x1.curr=(prob.immune*probx1.immune)/totalprobx;
prob.treatprotected.x1.curr=(prob.treatprotected*probx1.treatprotected)/totalprobx;
prob.treatharmful.x1.curr=(prob.treatharmful*probx1.treatharmful)/totalprobx;
prob.doomed.x1.curr=(prob.doomed*probx1.doomed)/totalprobx;
prob.immune.x0.curr=(prob.immune*(1-probx1.immune))/(1-totalprobx);
prob.treatprotected.x0.curr=(prob.treatprotected*(1-probx1.treatprotected))/(1-totalprobx);
prob.treatharmful.x0.curr=(prob.treatharmful*(1-probx1.treatharmful))/(1-totalprobx);
prob.doomed.x0.curr=(prob.doomed*(1-probx1.doomed))/(1-totalprobx);
itt.immune.treatprotected.curr=mean.immune.a1-mean.immune.a0;
sd.immune.a1.curr=sd.immune.a1; 
mean.immune.a0.curr=mean.immune.a0; 
sd.immune.a0.curr=sd.immune.a0;
sd.treatprotected.a1.curr=sd.treatprotected.a1;
mean.treatprotected.a0.curr=mean.treatprotected.a0;
sd.treatprotected.a0.curr=sd.treatprotected.a0;
itt.treatharmful.doomed.curr=mean.treatharmful.a1-mean.treatharmful.a0;
sd.treatharmful.a1.curr=sd.treatharmful.a1;
mean.treatharmful.a0.curr=mean.treatharmful.a0;
sd.treatharmful.a0.curr=sd.treatharmful.a0;
sd.doomed.a1.curr=sd.doomed.a1;
mean.doomed.a0.curr=mean.doomed.a0;
sd.doomed.a0.curr=sd.doomed.a0;
xeffect.curr=xeffect;

maxiters=500; # Maximum iterations in EM algorithm
iter=1;
lldist=1; # Will measure difference in log likelihood between iterations
toler=.001; # Tolerance for difference in log likelihood between iterations
# for stopping EM algorithm
conv=1;

while(lldist>toler && iter<maxiters){
  # Compute log likelihood with current parameters
ll.old.p1a=sum((1-aggressive.treatment)*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*esrd*x)*log(prob.doomed.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*(1-esrd)*x)*log(prob.immune.x1.curr+prob.treatprotected.x1.curr)+sum(aggressive.treatment*esrd*x)*log(prob.doomed.x1.curr+prob.treatharmful.x1.curr)+sum((1-aggressive.treatment)*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatharmful.x1.curr))*dnorm(y,mean=mean.immune.a0.curr+xeffect.curr,sd=sd.immune.a0.curr)+(prob.treatharmful.x1.curr)/(prob.immune.x1.curr+prob.treatharmful.x1.curr)*dnorm(y,mean=mean.treatharmful.a0.curr+xeffect.curr,sd=sd.treatharmful.a0.curr))));
  ll.old.p1b=sum((1-aggressive.treatment)*esrd*x*log(((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.doomed.a0.curr+xeffect.curr,sd=sd.doomed.a0.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.doomed.x1.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+xeffect.curr,sd=sd.treatprotected.a0.curr))));
ll.old.p2=sum(aggressive.treatment*(1-esrd)*x*log(((prob.immune.x1.curr/(prob.immune.x1.curr+prob.treatprotected.x1.curr))*dnorm(y,mean=mean.immune.a0.curr+itt.immune.treatprotected.curr+xeffect.curr,sd=sd.immune.a1.curr)+(prob.treatprotected.x1.curr)/(prob.treatprotected.x1.curr+prob.immune.x1.curr)*dnorm(y,mean=mean.treatprotected.a0.curr+itt.immune.treatprotected.curr+xeffect.curr,sd=sd.treatprotected.a1.curr))))+sum(aggressive.treatment*esrd*x*((prob.doomed.x1.curr/(prob.doomed.x1.curr+prob.treatharmful.x1.curr))*dno