########## R script: glmmPQLadditve ##########

# For illustrating how glmmPQL() can be used
# to fit and additive model 

# Set the type of model to be fit.

do.additive.model <- T

# Specify high level settings

family.val <- "poisson"

basis <- "truncated lines"

# Load the MASS library to access glmmPQL()
# load the nlme library to access groupedData()

library(MASS)
library(nlme) 

set.seed(14536)

n <- 500

x1 <- runif(n)
x2 <- runif(n)

f1 <- function(x)
{
  temp <- sin(pi*x)
  temp <- temp-mean(temp)
  return(temp)
}
  
f2 <- function(x)
  {
    temp <- sin(2*pi*x)
    temp <- temp-mean(temp)
    return(temp)
  }
mu <- exp(f1(x1)+f2(x2))

if (family.val=="poisson")
   y <- rpois(n,mu)

num.add.comps <- 2

num.knots <- rep(15,num.add.comps)

knots <- list()
knots[[1]] <- quantile(unique(x1),
                       seq(0,1,
                           length=(num.knots[1]+2))[-c(1,(num.knots[1]+2))])
knots[[2]] <- quantile(unique(x2),
                       seq(0,1,
                           length=(num.knots[2]+2))[-c(1,(num.knots[2]+2))])

# Construct X and Z matrix

Z.inds <- list()
Z.inds[[1]] <-
  (num.add.comps+2):(num.knots[1]+(num.add.comps+1))
Z.inds[[2]] <-
  (Z.inds[[1]][num.knots[1]]+1):(sum(num.knots)+(num.add.comps+1))

X <- cbind(rep(1,n),x1,x2)

Z.1 <- outer(x1,knots[[1]],"-")
Z.1 <- Z.1*(Z.1>0)
Z.1 <- cbind(Z.1)

Z.2 <- outer(x2,knots[[2]],"-")
Z.2 <- Z.2*(Z.2>0)
Z.2 <- cbind(Z.2)

C.mat <- cbind(X,Z.1,Z.2)
cols <- dim(C.mat)[2]
rescale.mean <- apply(C.mat[,2:cols],2,mean)

C.mat[,2:cols] <- scale(C.mat[,2:cols],scale=F)
X <- cbind(rep(1,length=n),C.mat[,2:(num.add.comps+1)])

Z.1 <- C.mat[,Z.inds[[1]]]
Z.2 <- C.mat[,Z.inds[[2]]]

# Obtain the fit using glmmPQL()

print("Note: currently just for Poisson")

re.block.val <- list(1:num.knots[1],1:num.knots[2])
Z.block  <-  list()
Z.block[[1]] <-
  as.formula(paste("~Z.1[,c(",
                   paste(re.block.val[[1]],collapse=","),")]-1"))
Z.block[[2]] <-
  as.formula(paste("~Z.2[,c(",
                   paste(re.block.val[[1]],collapse=","),")]-1"))


groups.1 <- rep(1,length=length(y))
data.fr <- groupedData( y ~ X[,-1] | groups.1,
     	           data = data.frame( y,X,Z.1,Z.2))
fit <- glmmPQL(y~-1+X,
               data=data.fr,
               random=
               list(groups.1 = pdBlocked(Z.block,
                      pdClass=c("pdIdent","pdIdent"))),family=poisson)


u1.hat <- as.vector(unlist(fit$coef$ran))[Z.inds[[1]]-(num.add.comps+1)]
u2.hat <- as.vector(unlist(fit$coef$ran))[Z.inds[[2]]-(num.add.comps+1)]
beta.hat <- as.vector(unlist(fit$coef$fix))

# plotting

grid.size <- 101
x1.grid <- seq(0,1,length=grid.size)
x2.grid <- seq(0,1,length=grid.size)
ones.grid <- rep(1,grid.size)

X.grid <- cbind(ones.grid,x1.grid,x2.grid)
Z1.grid <- outer(x1.grid,knots[[1]],"-")
Z1.grid <- Z1.grid*(Z1.grid>0)
Z2.grid <- outer(x2.grid,knots[[2]],"-")
Z2.grid <- Z2.grid*(Z2.grid>0)

C.grid <- cbind(X.grid,Z1.grid,Z2.grid)
C.grid[,2:cols] <- scale(C.grid[,2:cols],scale=F,center=rescale.mean)

X.grid <- cbind(rep(1,length=grid.size),
           C.grid[,2:(num.add.comps+1)])
Z1.grid <- C.grid[,Z.inds[[1]]]
Z2.grid <- C.grid[,Z.inds[[2]]]


eta1.hat.grid <- as.vector(X.grid[,2]*beta.hat[2]+Z1.grid%*%u1.hat)
eta1.grid <- f1(x1.grid)

eta2.hat.grid <- as.vector(X.grid[,3]*beta.hat[3]+Z2.grid%*%u2.hat)
eta2.grid <- f2(x2.grid)

if (family.val=="poisson")
  {
    mu1.hat.grid <-  exp(eta1.hat.grid)
    mu2.hat.grid <-  exp(eta2.hat.grid)
    mu1.grid <-  exp(eta1.grid)
    mu2.grid <-  exp(eta2.grid)
  }

x11()
par(bty="l",mfrow=c(2,1))

plot(c(x1.grid),c(mu1.hat.grid),bty="l",xlab="x1",
     ylab="f1(x)",type="n",ylim=range(c(mu1.grid,mu1.hat.grid)),
     main="Additive Model Component One and Fit")
lines(x1.grid,mu1.grid) 
lines(x1.grid,mu1.hat.grid,lwd=2,col="red")

plot(c(x2.grid),c(mu2.hat.grid),bty="l",xlab="x2",
     ylab="mu2(x)",type="n",ylim=range(c(mu2.grid,mu2.hat.grid)),
     main="Additive Model Component Two and Fit")
lines(x2.grid,mu2.grid) 
lines(x2.grid,mu2.hat.grid,lwd=2,col="red")


########## End of glmmPQL-tst ##########