#  Stage 1 GLS-PL pooled estimation
#  Stage 2 Estimation of population parameters using linear
#  mixed model software

#  This program outputs a data set that may then be input to
#  linear mixed model software like SAS proc mixed or Splus lme()
#  to carry out stage 2

#  function to print out matrices pretty

write.matrix <- function(x,file="",sep=" "){
  x <- as.matrix(x)
  p <- ncol(x)
  cat(dimnames(x)[[2]],format(t(x)),file=file,sep=c(rep(sep,p-1),
    "\n"),append=T)
}
	
#  put the mean function you want here 

meanfunc <- function(x,b1,b2,dose){

	tinf <- 240
	cl <- exp(b1)
	v <- exp(b2)
        t1 <- x<=tinf
        t2 <- tinf*(1-t1)+t1*x
        f1 <- (dose/cl)*(1-exp(-cl*t2/v))*exp(-cl*(1-t1)*(x-tinf)/v)

#  compute analytical dervivatives -- create the gradient matrix X

        t3 <- (1-t1)*(x-tinf)
	temp1 <- (dose/cl)*exp(-cl*t3/v)
        temp2 <- (dose/(v^2))*exp(-cl*t3/v)
        
   meangrad <- array(0,c(length(x),2),list(NULL,c("b1","b2")))
   meangrad[,"b1"] <- temp1*(-(1-exp(-cl*t2/v))*(1/cl + t3/v)+(t2/v)*exp(-cl*t2/v))*cl 
   meangrad[,"b2"] <- temp2*((1-exp(-cl*t2/v))*t3-exp(-cl*t2/v)*t2)*v
 
   attr(f1,"gradient") <- meangrad
   f1

}

meanfunc2 <- function(x,b1,b2,dose){

	tinf <- 240
	cl <- exp(b1)
	v <- exp(b2)
        t1 <- x<=tinf
        t2 <- tinf*(1-t1)+t1*x
        f1 <- (dose/cl)*(1-exp(-cl*t2/v))*exp(-cl*(1-t1)*(x-tinf)/v)
        f1
}

weightfunc <- function(x,b1,b2,dose,mut){
	
   f1 <- meanfunc2(x,b1,b2,dose)
   weightf <- f1/mut
   tinf <- 240
   cl <- exp(b1)
   v <- exp(b2)
   t1 <- x<=tinf
   t2 <- tinf*(1-t1)+t1*x
   t3 <- (1-t1)*(x-tinf)
   temp1 <- (dose/cl)*exp(-cl*t3/v)
   temp2 <- (dose/(v^2))*exp(-cl*t3/v)

#  compute analytical dervivatives -- create the gradient matrix X
        
   weightgrad <- array(0,c(length(x),2),list(NULL,c("b1","b2")))
   weightgrad[,"b1"] <- temp1*(-(1-exp(-cl*t2/v))*(1/cl + t3/v)+(t2/v)*exp(-cl*t2/v))*cl/mut
   weightgrad[,"b2"] <- temp2*((1-exp(-cl*t2/v))*t3-exp(-cl*t2/v)*t2)*v/mut
   attr(weightf,"gradient") <- weightgrad
   weightf
}
	
#  set start values, etc

# max number of iterations for fitting algorithm

cmax <- 20

#  start values

bstart <- list(b1=-6.0,b2=-2.0)

tstart <- list(theta=0.5)
p <- 2

#  name of output file

outfile <- "argstage1.Rout"
cat("ARGATROBAN DATA -- STAGE 1",file=outfile,"\n","\n","\n",append=F)

#  read in data

thedata <- scan("argconc.dat")
thedata <- matrix(thedata,ncol=5,byrow=T);

indiv <- thedata[,2]   # individual indicator
ds <- thedata[,3]      # dose
xs <- thedata[,4]      # time
ys <- thedata[,5]      # concentration

n <- length(xs)
uindiv <- unique(indiv)
m <- length(uindiv)
		
# form individual 2nd stage covariate design matrices Ai

# these are all identity matrices in this case

aimat <- NULL
r <- p

a1 <- diag(p)
	
for(i in 1:m){

      aimat <- rbind(aimat,a1)
	
}


#######################################################################

#  "mean" function for the "trick"

pltrkfunc <- function(resid,mudot,mu,theta){
   trk <- resid*((mudot/mu)**theta)
   trkgrad <- array(0,c(length(mu),1),list(NULL,c("theta")))
   trkgrad[,"theta"] <- trk*log(mudot/mu)
   attr(trk,"gradient") <- trkgrad   #  analytic derivative
   trk
}


#  Stage 1 -- Pooled GLS-PL estimation

#  Step 1 -- initial fit by OLS for each indiv

cat("INITIAL OLS ESTIMATION",file=outfile,"\n","\n",append=T)

bolsmat <- NULL     #  matrix to contain OLS estimates
muvec <- NULL       #  vector to contain all predicted values
residvec <- NULL    #  vector to contain all residuals
		
for (i in 1:m){

id <- uindiv[i]
y <- ys[indiv==id]
x <- xs[indiv==id]
d <- ds[indiv==id]	

olsdat <- data.frame(x,y,d)

ols.fit <- nls(y ~ meanfunc(x,b1,b2,d),olsdat,bstart)
bols <- coef(ols.fit)
	
mu <- meanfunc2(x,bols[1],bols[2],d)
resid <- y-mu

bols <- matrix(bols,ncol=p,byrow=T)	
bolsmat <- rbind(bolsmat,bols)

muvec <- c(muvec,mu)
residvec <- c(residvec,resid)	

}

bols <- cbind(uindiv,bolsmat)

muvec <- matrix(muvec,ncol=1)
residvec <- matrix(residvec,ncol=1)
		
cat("OLS estimates",file=outfile,"\n","\n",append=T)

write.matrix(round(bols,6),file=outfile)
	
#  pooled OLS estimate of sigma

sigma <- sqrt(sum(residvec*residvec)/(n-m*p))

cat(file=outfile,"\n","\n",append=T)	
cat("OLS pooled estimate of sigma ",sigma,"\n","\n",file=outfile,append=T)
			
# begin iteration loop for GLS-PL
				
cat("GLS-PL POOLED ESTIMATION",file=outfile,"\n","\n",append=T)
	
for (k in 1:cmax){

#  compute the geometric mean and predicted values
	
#mudot <- prod(muvec)**(1/n)
mudot <- exp((1/n)*sum(log(muvec)))
mudot <- rep(mudot,length(muvec))
dummy <- rep(0,length(muvec))
pldat <- data.frame(residvec,muvec,mudot)

#  Step 2 -- estimate theta using the PL trick
	
pl.fit <- nls(dummy ~ pltrkfunc(residvec,mudot,muvec,theta),pldat,
	tstart,control=list(maxiter=500))
theta <- coef(pl.fit)
	
cat("Iteration ",k,"\n",file=outfile,append=T)
cat("Estimate of theta ",theta,"\n",file=outfile,append=T)	

#  Step 3 -- update estimates of beta for each individual

bglsmat <- NULL
new.muvec <- NULL
new.residvec <- NULL

for (i in 1:m){

id <- uindiv[i]
y <- ys[indiv==id]
x <- xs[indiv==id]
d <- ds[indiv==id]	

mu <- muvec[indiv==id]
	
mut <- mu^theta
ymut <- y/mut
	
glsdat <- data.frame(x,ymut,d)

glspl.fit <- nls(ymut ~ weightfunc(x,b1,b2,d,mut),glsdat,
	bstart,control=list(maxiter=100))

bgls <- coef(glspl.fit)

mu <- meanfunc2(x,bgls[1],bgls[2],d)
resid <- y-mu

bgls <- matrix(bgls,ncol=p,byrow=T)	
bglsmat <- rbind(bglsmat,bgls)

new.muvec <- c(new.muvec,mu)
new.residvec <- c(new.residvec,resid)	

}

bgls <- cbind(uindiv,bglsmat)
muvec <- matrix(new.muvec,ncol=1)
residvec <- matrix(new.residvec,ncol=1)
		
}

#  Compute final estimate of sigma^2

g <- muvec**theta

sigma2 <- sum((residvec/g)**2)/(n-m*p)
sigma <- sqrt(sigma2)

cat(file=outfile,"\n","\n","\n",append=T)
cat("Final GLS estimates after ",cmax," iterations",
	file=outfile,"\n","\n",append=T)

write.matrix(round(bgls,6),file=outfile)

cat(file=outfile,"\n","\n",append=T)

cat("Final pooled estimate of sigma ",sigma,"\n","\n",file=outfile,append=T)
	
#  now transform problem to allow use of mix model software to fit 2nd stage

mixdat <- NULL
	
for (i in 1:m){

id <- uindiv[i]
y <- ys[indiv==id]
x <- xs[indiv==id]
d <- ds[indiv==id]	

thisid <- rep(i,p)

#  get covariance matrix for bglsi

bglsi <- bgls[i,2:(p+1)]
mui <- meanfunc2(x,bglsi[1],bglsi[2],d)
muit <- mui^theta

xmati <- attr(weightfunc(x,bglsi[1],bglsi[2],d,muit),"gradient")

cmati <- round(sigma^2*solve(t(xmati) %*% xmati),6)	

#  cholesky decomp of inverse 

cinvhalf <- chol(round(solve(cmati),6))

#  Create the "data" for stage 2 -- response, "Xi" and "Zi"

ai <- aimat[((i-1)*p+1):(i*p),]

respi <- cinvhalf%*%bglsi

thisxi <- cinvhalf%*%ai
	
thiszi <- cinvhalf
	
thisdat <- cbind(thisid,respi,as.matrix(thisxi),as.matrix(thiszi))

mixdat <- rbind(mixdat,thisdat)

}

write.matrix(mixdat,file="argmixsig.dat")