# ZINB with random effect
# last modified on 15/12/2000; 5/11/2002
#
# Original Splus code provided by Drs. Andy Lee and Kelvin Yau
#
# Modifications for R by Dave Atkins
#
# For relevant papers, see Drs. Lee and Yau websites:
#
# http://fbstaff.cityu.edu.hk/mskyau/
# http://www.publichealth.curtin.edu.au/html/about_staffprofile.cfm?ID=482
#
## ---------------------------------------
#  MLE of scale parameter alpha of NB regression with weights w
# ---------------------------------------

agetk.ml <- function(y, mu, w) # calculate the k
{
loglik <- function(th,y,mu,w)
{
u <- exp(th)/(exp(th)+mu)
(sum(w*(log(gamma(y+exp(th))/gamma(y+1)/gamma(exp(th)))+exp(th)*log(u)+log((1-u)^y))))
}
objm <- optimize(loglik,lower =-8, upper =5, y=y,mu=mu,w=w,maximum=T)
res <- objm$maximum
1/exp(res)
}

## ---------------------------------------

#  GLMM of Poisson regression with weights mzk

# ---------------------------------------

wreml.poi <- function(y, mzk, x, z, beta1, va1, sig2, fam="Poisson", epsilon=1e-3)
{
M <- ncol(z);n <- length(y)
X <- cbind(1,x);p1 <- ncol(X)
zero1 <- matrix(0,ncol=p1,nrow=M)	
X1 <- rbind(X,zero1)
Z <- rbind(z,diag(M))
XX <- cbind(X1,Z)

itmax <- 1000;
alfa0 <- c(beta1,va1)
beta <- beta1 ; va <- va1
flag <- 0
for(iter in 1:itmax)
{
theta <- as.vector(X%*%beta+z%*%va)
lamda <- exp(theta)
w1 <- mzk*lamda
w <- c(w1,rep(1/sig2,M))
mu <- lamda

w.sq <- w^0.5
zy <- c((theta+(y-mu)/mu),rep(0,M))*w.sq
zx <- XX*w.sq
tfit <- lm.fit(zx,zy) # Dave: change to lm.fit
  
Alfa <- coef(tfit)
beta <- Alfa[1:p1]
va <- Alfa[(p1+1):(p1+M)]
if(max(abs(Alfa-alfa0))<epsilon) { flag <- 1;break}
alfa0 <- Alfa;
#cat(beta, iter,'\n')
}
if(flag)  result <- list(beta = beta, va = va)
else stop("error: not reach the convergence")
}

# ---------------------------------------

#  GLMM of Logistic regression

# ---------------------------------------
wreml.logit <- function(y,x,z,alfa1,yu1,sig1,famaly="logistic",epsilon=1e-3)
{
M <- ncol(z);n <- length(y)
X <- cbind(1,x)
p1 <- ncol(X)
zero1 <- matrix(0,ncol=p1,nrow=M)	
X1 <- rbind(X,zero1)
Z <- rbind(z,diag(M))
XX <- cbind(X1,Z)

itmax <- 1000
alfa0 <- c(alfa1,yu1)
alfa <- alfa1
yu <- yu1;flag <- 0
for(iter in 1:itmax)
{
theta <- as.vector(X%*%alfa+z%*%yu)
w1 <- exp(theta)/(1+exp(theta))^2
w <- c(w1,rep(1/sig1,M))
mu <- exp(theta)/(1+exp(theta))

w.sq <- w^0.5
zy <- c((theta+(y-mu)/w1),rep(0,M))*w.sq
zx <- XX*w.sq
tfit <- lm.fit(zx,zy) # Dave: change to lm.fit
  
Alfa <- coef(tfit)
alfa <- Alfa[1:p1]
yu <- Alfa[(p1+1):(p1+M)]

if(max(abs(Alfa-alfa0))<epsilon) { flag <- 1;break}
alfa0 <- Alfa
}
if(flag) reslt <- list(alfa=alfa,yu=yu)
else stop("error:not reach the convergence")
}

hznb <- function(y,X,G,Z,pai,mu,th,sig=NULL,sigu=NULL)
{

yzero <- ifelse(y>0,0,1)
ksi <- pai/(1-pai)
u <- th/(th+mu)
ep1 <- u^th
ep2 <- (ksi+ep1)^2

# information matrix

w11<--(yzero*ksi*ep1/ep2-ksi/(1+ksi)^2)
w12<--(yzero*th*ksi*ep1*(1-u)/ep2)

w22 <- (-yzero*th*ksi*(ksi+ep1+(1-u)*th*ep1/u)/ep2+th+(1-yzero)*y)*u*(1-u)

# second derivtive of alpha
#B<-log(u)+(1-u)
#B1<-(1-u)^2/th
#w23<-(1-u-(1-yzero)*y*u/th+yzero*ksi*(u/(ksi+ep1)/th+ep1*B/ep2))*(1-u)
#-------------------------------------
pa <- ncol(G)
pb <- ncol(X)
p <- ncol(Z)

ww11 <- t(matrix(rep(w11,pa),ncol=pa))
ww12 <- t(matrix(rep(w12,pa),ncol=pa))
ww22 <- t(matrix(rep(w22,pb),ncol=pb))
m11 <- (t(G)*ww11)
m12 <- (t(G)*ww12)
m22 <- (t(X)*ww22)

I11 <- m11%*%G
I12 <- m12%*%X

I22 <- m22%*%X

if(!is.null(sig))
{
z22 <- t(Z)*t(matrix(rep(w22,p),ncol=p))
I14 <- m12%*%Z
I24 <- m22%*%Z
I44 <- z22%*%Z+diag(1/sig,p)
}
if(!is.null(sigu))
{
z11 <- t(Z)*t(matrix(rep(w11,p),ncol=p))
z12 <- (matrix(rep(w12,p),ncol=p))*Z

I13 <- m11%*%Z
I23 <- t(X)%*%z12
I33 <- z11%*%Z+diag(1/sigu,p)
}

if(is.null(sigu)&is.null(sig))
{
	V1 <- cbind(I11,I12)
	V2 <- cbind(t(I12),I22)
	V <- rbind(V1,V2)
}
if((!is.null(sigu))&(!is.null(sig)))
{
	I34 <- t(Z)%*%z12

	V1 <- cbind(I11,I12,I13,I14)
	V2 <- cbind(t(I12),I22,I23,I24)
	V3 <- cbind(t(I13),t(I23),I33,I34)
	V4 <- cbind(t(I14),t(I24),t(I34),I44)
	V <- rbind(V1,V2,V3,V4)
	
	M2 <- diag(0,(pa+pb+p+p))
	M2[pa+pb+1:p,pa+pb+1:p] <- diag(1/sigu,p)
	M2[(pa+pb+p)+1:p,(pa+pb+p)+1:p] <- diag(1/sig,p)

	H <- V-M2
}

if(is.null(sigu)&(!is.null(sig)))
{	
	V1 <- cbind(I11,I12,I14)
	V2 <- cbind(t(I12),I22,I24)
	V4 <- cbind(t(I14),t(I24),I44)
	V <- rbind(V1,V2,V4)
}
if((!is.null(sigu))&(is.null(sig)))
{	
	V1 <- cbind(I11,I12,I13)
	V2 <- cbind(t(I12),I22,I23)
	V3 <- cbind(t(I13),t(I23),I33)

	V <- rbind(V1,V2,V3)
}

	IV <- solve(V)
if((!is.null(sigu))&(!is.null(sig)))
{
	df <- length(y)-sum(diag(IV%*%H))
   list(dd=diag(IV),df=df) 	
}

else dd <- diag(IV)

				
}

################### Main Estimation Function ###################################

zinbmix <- function(y, x.p=NULL, rv=NULL, random, x.l=NULL, model)
{	
itmax <- 1000
n <- length(y)
yz <- ifelse(y > 0, 0, 1)
ct0 <-  list(epsilon = 0.001, maxit = 50, trace = F)
if(!is.null(x.l))
{
x.l <- as.matrix(x.l)	
G <- cbind(1,x.l)
alfa <- coef(glm(yz ~ x.l, family = binomial(link = logit), 
                  na.action = na.omit, control = ct0))
}
else
{       		 
alfa <- coef(glm(yz ~ 1, family = binomial(link = logit), 
                  na.action = na.omit, control = ct0))
G <- as.matrix(rep(1,n))
}
if(!is.null(x.p))
{
x.p <- as.matrix(x.p)	
X <- cbind(1,x.p)
beta <- coef(glm(y ~ x.p, family = poisson(link =log), 
                  na.action = na.omit, control = ct0))
}
else
{       	 
beta <- coef(glm(y ~ 1, family = poisson(link = log), 
                  na.action = na.omit, control = ct0))
X <- as.matrix(rep(1,n))
}
	
pa <- ncol(G)
pb <- ncol(X)
m <- ncol(rv)

#initial value
ZK1 <- ifelse(y > 0, 1, 0)
th <- 1

yu <- rep(0., m)
va <- rep(0, m)
sigu <- 1.2
sig2 <- 0.1
names(sig2) <- "RandomEffect"
flag <- 0
# beginning of outer loop
for( ie in 1:itmax)
{
for (iter in 1:itmax)
{
 if(is.null(x.l))
	theta <- as.vector(exp(G*alfa))
 else	
	{
		if(model == "rnb" | model == "zinb" )
		theta <- as.vector(exp(G %*% alfa))
		else theta <- as.vector(exp(G %*% alfa+rv%*%yu))
	}	
	
if(is.null(x.p))
	mu <- as.vector(exp(X*beta))
else	
	{
		if(model == "rlg" | model == "zinb" )
		mu <- as.vector(exp(X %*% beta))
		else mu <- as.vector(exp(X%*%beta+rv%*%va))
	}	
	k <- agetk.ml(y,mu,(1-ZK1))
	th <- 1/k
	#  E-step
	ZK <- ifelse(y > 0, 0, 1/(1+1/theta*(th/(mu+th))^th))
	#  M-step
	wmm <- 1/(1+k*mu)*(1-ZK) #weight

   if(!is.null(x.l))
	{
	if((model != "rnb")&(model != "zinb"))
	{
      lgt <- wreml.logit(ZK,x.l,rv,alfa,yu,sigu)
		alfa <- lgt$alfa
		yu <- lgt$yu
	}	
	else
			alfa <- coef(glm(ZK ~ x.l, family = binomial(link = logit), 
                  na.action = na.omit, control = ct0))
	}
	else alfa <- coef(glm(ZK ~ 1, family = binomial(link = logit), 
                  na.action = na.omit, control = ct0))

   if(!is.null(x.p))