# This is Example 7.3 (pages 251-2)  of Gamerman and Lopes (2006) 
# on non-linear regression models.  The actual data was presented on the 
# temporal evolution of the dry weight of onion bulbs.

g = function(x,beta){
  beta[1]+beta[2]*beta[3]^x
}

loglike = function(beta,sigma){
  return(sum(dnorm(y,g(x,beta),sigma,log=TRUE)))
}

minusloglike = function(theta){
  beta  = c(theta[1],theta[2],1/(1+exp(-theta[3])))
  sigma = exp(theta[4])
  return(-sum(dnorm(y,g(x,beta),sigma,log=TRUE)))
}

x = c(1.0,1.5,1.5,1.5,2.5,4.0,5.0,5.0,7.0,8.0,8.5,9.0,9.5,9.5,10.0,
12.0,12.0,13.0,13.0,14.5,15.5,15.5,16.5,17.0,22.5,29.0,31.5)
y = c(1.80,1.85,1.87,1.77,2.02,2.27,2.15,2.26,2.47,2.19,2.26,2.40,2.39,2.41,
2.50,2.32,2.32,2.43,2.47,2.56,2.65,2.47,2.64,2.56,2.70,2.72,2.57)

pdf(file="nonlinear-regression-dryweightofonionbulbs.pdf")
par(mfrow=c(1,1))
plot(x,y,xlab="time",ylab="dry weight of onion bulbs")

# How I selected the initial values for the minimizer
# When x = 30 => y=2.6 ~ beta0
# When x =  0 => y=1.8 ~ beta0+beta1 => beta1 ~ -0.8
# When x = 10 => y=2.4 ~ beta0+beta1*beta2^(10) ~ 2.4 => beta2 ~ 0.87
# sigma ~ sqrt(mean((y-(2.6-0.8*0.87^x))^2))

beta0     = c(2.6,-0.8,0.87)
sigma0    = 0.1
theta0    = c(beta0[1:2],log(beta0[3]/(1-beta0[3])),log(sigma0))
estimate  = nlm(minusloglike,theta0)$estimate
beta.mle  = c(estimate[1:2],1/(1+exp(-estimate[3])))
sigma.mle = exp(estimate[4])

par(mfrow=c(1,1))
plot(x,y,xlab="time",ylab="dry weight of onion bulbs")
xx = seq(min(x),max(x),length=100)
lines(xx,g(xx,beta.mle),col=2,lwd=2)


# Non-informative and independent 
# prior distributions for (beta0,beta1,beta2,sigma)
#
# beta0 ~ N(2.6,0.25)
# beta1 ~ N(-1.0,0.25)
# beta2 ~ U(0.6,1.0)
# sigma ~ gamma(2,10)
d0    = c(2.6,-1.0)
D0    = diag(0.25,2)
L0    = 0.6
U0    = 1.0
a0    = 2
b0    = 10

# Sampling importance resampling
# ------------------------------
set.seed(54321)
M          = 100000
beta       = matrix(0,M,3)
beta[,1:2] = matrix(d0,M,2,byrow=TRUE)+matrix(rnorm(M*2),M,2)%*%chol(D0)
beta[,3] = runif(M,L0,U0)
sigma    = rgamma(M,a0,b0)

w = rep(0,M)
for (i in 1:M)
  w[i] = loglike(beta[i,],sigma[i])
w = exp(w-max(w))
ind = sample(1:M,size=M,replace=TRUE,prob=w)

par(mfrow=c(2,2))
hist(beta[ind,1],prob=TRUE,main=expression(beta[0]),xlab="")
lines(density(beta[,1]),col=2,lwd=2)
hist(beta[ind,2],prob=TRUE,main=expression(beta[1]),xlab="")
lines(density(beta[,2]),col=2,lwd=2)
hist(beta[ind,3],prob=TRUE,main=expression(beta[2]),xlab="")
abline(h=1/0.4,col=2,lwd=2)
hist(sigma[ind],prob=TRUE,main=expression(sigma),xlab="")
lines(density(sigma),col=2,lwd=2)


# discarding the prior
beta = beta[ind,]
sigma = sigma[ind]
draws.sir = cbind(beta,sigma)

# Posterior predictive
N        = 100
xx       = seq(0,50,length=N)
post.sir = matrix(0,M,N)
for (i in 1:N)
  post.sir[,i] = draws.sir[,1]+ draws.sir[,2]*draws.sir[,3]^xx[i]+rnorm(M,0, draws.sir[,4])
qs.sir = apply(post.sir,2,quantile,c(0.025,0.5,0.975))

fit.mle   = g(xx,beta.mle)

par(mfrow=c(1,1))
plot(x,y,xlab="time",ylab="dry weight of onion bulbs",
     ylim=range(fit.mle,qs),xlim=range(x,xx),pch=16)
lines(xx,fit.mle,col=2,lwd=2)
lines(xx,qs.sir[2,],col=4,lwd=2)
lines(xx,qs.sir[1,],col=4,lwd=2,lty=2)
lines(xx,qs.sir[3,],col=4,lwd=2,lty=2)
abline(v=min(x),lty=2)
abline(v=max(x),lty=2)
title("Posterior predictives")


D0    = diag(0.25,2)
d0    = c(2.6,-1.0)
iD0   = solve(D0)
iD0d0 = iD0%*%d0
L0    = 0.6
U0    = 1.0
a0    = 2
b0    = 10


# MCMC  set up
set.seed(12345)
M0    = 1000
M     = 25000
niter = M0+M
draws = matrix(0,niter,4)

# initial values
beta  = beta.mle
sigma = sigma.mle
sig2  = sigma^2

for (iter in 1:niter){
  X = cbind(1,beta[3]^x)

  # Sampling beta1,beta2 jointly (Gibbs step)
  var = solve(iD0+t(X)%*%X/sig2)
  mean = var%*%(iD0d0+t(X)%*%y/sig2)
  beta[1:2] = mean+t(chol(var))%*%rnorm(2)

  # Sampling beta3 (RW-MH step)
  beta3 = rnorm(1,beta[3],0.1)
  num   = sum(dnorm(y,g(x,c(beta[1:2],beta3)),sigma,log=TRUE))
  den   = sum(dnorm(y,g(x,beta),sigma,log=TRUE))
  logaccept = min(0,num-den)
  if (log(runif(1))<logaccept){
    beta[3] = beta3
  }

  # Sampling sigma (RW-MH step)
  sig = rnorm(1,sigma,0.005)
  num = sum(dnorm(y,g(x,beta),sig,log=TRUE))+dgamma(sig,a0,b0,log=TRUE)
  den = sum(dnorm(y,g(x,beta),sigma,log=TRUE))+dgamma(sigma,a0,b0,log=TRUE)
  logaccept = min(0,num-den)
  if (log(runif(1))<logaccept){
    sigma = sig
  }
  
  # Storing the draws
  draws[iter,] = c(beta,sigma)
}
draws.mcmc = draws[(M0+1):niter,]

post.mcmc = matrix(0,M,N)
for (i in 1:N)
  post.mcmc[,i] = draws.mcmc[,1]+draws.mcmc[,2]*draws.mcmc[,3]^xx[i]+
                  rnorm(M,0,draws.mcmc[,4])
qs.mcmc = apply(post.mcmc,2,quantile,c(0.025,0.5,0.975))


param = c("beta1","beta2","beta3","sigma")
par(mfrow=c(3,4),mai=c(0.6,0.6,0.2,0.05))
for (i in 1:4)
 ts.plot(draws.mcmc[,i],main=param[i],ylab="",xlab="Iterations")
for (i in 1:4)
 acf(draws.mcmc[,i],main="")
for (i in 1:4){
  hist(draws.sir[,i],prob=TRUE,xlab="",ylab="Posterior density",main="")
  lines(density(draws.mcmc[,i]),col=2)
  box()
}

par(mfrow=c(1,1),mai=c(1.02,0.82,0.82,0.42))
plot(x,y,xlab="time",ylab="dry weight of onion bulbs",
     ylim=range(fit.mle,qs),xlim=range(x,xx),pch=16)

lines(xx,qs.sir[2,],col=4,lwd=2)
lines(xx,qs.sir[1,],col=4,lwd=2,lty=2)
lines(xx,qs.sir[3,],col=4,lwd=2,lty=2)

lines(xx,qs.mcmc[2,],col=6,lwd=2)
lines(xx,qs.mcmc[1,],col=6,lwd=2,lty=2)
lines(xx,qs.mcmc[3,],col=6,lwd=2,lty=2)

abline(v=min(x),lty=2)
abline(v=max(x),lty=2)
title("Posterior predictives")
legend("bottomright",legend=c("SIR","MCMC"),col=c(4,6),lty=1,lwd=2)

dev.off()