################################################################
#
#  Model:  x1,...,xn iid t(mu,sig,nu)
#  mu: location in R
#  sig: scale in R+
#  nu: degrees of freedom (known)
#
#  Prior: p(mu,sig)=p(mu)p(sig)
#         mu   ~ N(0,4)
#         sig2 ~ IG(1,1)
#
################################################################
loglike = function(mu,sig){
  sum(dt((x-mu)/sig,df=nu,log=TRUE))-n*log(sig)
}

logpost = function(mu,sig){
  loglike(mu,sig) + dnorm(mu,0,2,log=TRUE) +
  dgamma(1/sig,1,1,log=TRUE)-2*log(sig)
}

##########################
# Simulating some data
##########################
set.seed(123456)
n    = 100
nu   = 3
mu0  = 0
sig0 = 1
x    = mu0+sig0*rt(n,df=nu)

pdf(file="Student-t-model-SIR-RWM-dataaugmentation.pdf",width=12,height=8)

par(mfrow=c(2,2))
plot(x,xlab="Observations",main="")
title(paste("x1,...,xn iid t(",mu0,",",sig0^2,",",nu,")",sep=""))
hist(x,seq(min(x),max(x),length=15),prob=TRUE,main="")

mus = seq(-10,10,length=1000)
sigs = seq(0.1,20,length=1000)

plot(mus,dnorm(mus,0,2),col=6,lwd=2,type="l",
     xlab=expression(mu),ylab="Prior density")
abline(v=mu0,lwd=2)
plot(sigs,dgamma(1/sigs^2,1,1)/sigs^2,col=6,lwd=2,type="l",
     xlab=expression(sigma),ylab="Prior density")
abline(v=sig0,lwd=2)


############################################
# Sampling importance resampling (SIR
# Here, the proposal is the prior
############################################
M    = 1000000
N    = 10000
mus  = rnorm(M,0,2)
sigs = 1/rgamma(M,1,1)
w    = rep(0,M)
for (i in 1:M)
  w[i] = loglike(mus[i],sigs[i])
w = exp(w-max(w))
ind = sample(1:M,size=N,prob=w,replace=TRUE)

draws.sir = cbind(mus[ind],sigs[ind])

mus  = seq(-10,10,length=1000)
sigs = seq(0.1,20,length=1000)

# Range such that all graphs are comparable
mur  = c(-0.6,0.6)
sigr = c(0.5,1.5)

par(mfrow=c(1,2))
hist(draws.sir[,1],xlab=expression(mu),
     main="SIR\nProposal=prior",prob=TRUE,xlim=mur)
abline(v=mu0,col=2,lwd=2)
lines(mus,dnorm(mus,0,2),col=6,lwd=2)
legend("topleft",legend=c(paste("M=",M,sep=""),
       paste("N=",N,sep="")),bty="n")
hist(draws.sir[,2],xlab=expression(sigma),main="",prob=TRUE,xlim=sigr)
abline(v=sig0,col=2,lwd=2)
lines(sigs,dgamma(1/sigs^2,1,1)/sigs^2,col=6,lwd=2)



############################################
# Random-walk Metropolis
############################################
logpost = function(mu,sig){
  sum(dt((x-mu)/sig,df=nu,log=TRUE))-n*log(sig)+
  dnorm(mu,0,2,log=TRUE)+
  dgamma(1/sig,1,1,log=TRUE)-2*log(sig)
}

M0 = 1000
M  = 10000
L  = 50
niter = M0+M*L
draws = matrix(0,niter,2)

mu  = 0
sig = 2

for (i in 1:niter){
  mu.d  = rnorm(1,mu,0.1)
  alpha = min(0,logpost(mu.d,sig)-logpost(mu,sig))
  if (log(runif(1))<alpha){
    mu = mu.d
  }

  sig.d  = rnorm(1,sig,0.1)
  alpha = min(0,logpost(mu,sig.d)-logpost(mu,sig))
  if (log(runif(1))<alpha){
    sig = sig.d
  }
  draws[i,] = c(mu,sig)
}
ind = seq(M0+1,niter,by=L)

draws.rwm = draws[ind,]

mus  = seq(-10,10,length=1000)
sigs = seq(0.1,20,length=1000)

par(mfrow=c(2,3))
ts.plot(draws.rwm[,1],xlab="Iterations",ylab="",main=expression(mu))
acf(draws.rwm[,1],main="Random-walk Metropolis")
hist(draws.rwm[,1],prob=TRUE,xlab=expression(mu),main="",xlim=mur)
abline(v=mu0,col=2,lwd=2)
lines(mus,dnorm(mus,0,2),col=6,lwd=2)
ts.plot(draws.rwm[,2],xlab="Iterations",ylab="",main=expression(sigma))
acf(draws.rwm[,2],main="")
hist(draws.rwm[,2],prob=TRUE,xlab=expression(mu),main="",xlim=sigr)
abline(v=sig0,col=2,lwd=2)
lines(sigs,dgamma(1/sigs^2,1,1)/sigs^2,col=6,lwd=2)



############################################
# data augmentation
# x|lambda,mu,sig ~ N(mu,lambda*sig2)
# lambda sim IG(nu/2,nu/2)
############################################
M0     = 1000
M      = 10000
L      = 5
niter  = M0+L*M
draws  = matrix(0,niter,2)

lambda = rep(1,n)
sig    = 1

for (i in 1:niter){
  x1 = x/lambda
  ilam = sum(1/lambda)
  mu = rnorm(1,sum(x1)/ilam,sig/sqrt(ilam))
  sig = sqrt(1/rgamma(1,n/2,sum((x-mu)^2/lambda)/2))
  lambda = 1/rgamma(n,(nu+1)/2,(nu+(x-mu)^2/sig^2)/2)
  draws[i,] = c(mu,sig)
}

ind = seq(M0+1,niter,by=L)

draws.gibbs = draws[ind,]

mus  = seq(-10,10,length=1000)
sigs = seq(0.1,20,length=1000)

par(mfrow=c(2,3))
ts.plot(draws.gibbs[,1],xlab="Iterations",ylab="",main=expression(mu))
acf(draws.gibbs[,1],main="Gibbs samplers")
hist(draws.gibbs[,1],prob=TRUE,xlab=expression(mu),main="",xlim=mur)
abline(v=mu0,col=2,lwd=2)
lines(mus,dnorm(mus,0,2),col=6,lwd=2)
ts.plot(draws.gibbs[,2],xlab="Iterations",ylab="",main=expression(sigma))
acf(draws.gibbs[,2],main="")
hist(draws.gibbs[,2],prob=TRUE,xlab=expression(mu),main="",xlim=sigr)
abline(v=sig0,col=2,lwd=2)
lines(sigs,dgamma(1/sigs^2,1,1)/sigs^2,col=6,lwd=2)



##################################################################
# Comparing SIR, random-walk Metropolis and Gibbs sampler
##################################################################
par(mfrow=c(1,2))
plot(density(draws.sir[,1]),xlab=expression(mu),main="",lwd=2)
lines(density(draws.rwm[,1]),col=2,lwd=2)
lines(density(draws.gibbs[,1]),col=3,lwd=2)
lines(mus,dnorm(mus,0,2),col=6,lwd=2)
abline(v=mu0,col=4,lwd=2)
legend("topleft",legend="Prior",lwd=2,col=6,bty="n")

plot(density(draws.sir[,2]),xlab=expression(sigma),main="",lwd=2)
lines(density(draws.rwm[,2]),col=2,lwd=2)
lines(density(draws.gibbs[,2]),col=3,lwd=2)
lines(sigs,dgamma(1/sigs^2,1,1)/sigs^2,col=6,lwd=2)
abline(v=sig0,col=4,lwd=2)
legend("topright",legend=c("SIR","RWM","GIBBS"),
       col=1:3,lwd=2,lty=1,bty="n")

dev.off()