############################################################################################
#
# Example: Gibbs Sampler, Poisson-Exponential Model, Hierarchical Bayes
# 
# Professor M. Zhu
# for teaching illustration
#
# Current version = Oct 2022
#
############################################################################################

rm(list=ls(all=TRUE))
plot.only=FALSE

my.gibbs.PoissonExp<-function(x, a=0.001, b=0.001, t.max=20000, burn.in=15000)
{
  n=length(x)
  lambda<-matrix(0,n,t.max)
  theta<-numeric(t.max)
  lambda[,1]<-x                     # initialize lambda with MLE

  # Gibbs iteration here
  for (t in 2:t.max){
    theta[t]<-rgamma(1, shape=n+a, rate=sum(lambda[,(t-1)])+b) 
    lambda[,t]<-rgamma(n, shape=x+1, rate=theta[t]+1)
  }

  return(list(x=x, 
              lambda=lambda[,(burn.in+1):t.max],
              theta=theta[(burn.in+1):t.max],
              theta.empBayes=1/mean(x)))
}

my.plot <- function(obj, which.plot=NULL)
{
  # unpack 'obj'
  x=obj$x
  lambda=obj$lambda
  theta=obj$theta
  theta.empBayes = obj$theta.empBayes
  
  # prepare to plot
  if (is.null(which.plot)) {        # if not specified, simply plot all
    which.plot=1:length(x)
  }
  plotside=ceiling(sqrt(length(which.plot))) 
  par(mfrow=c(plotside,plotside),mar=c(2,2,2,2)+0.1)

  # plot posterior of lambda_i 
  for (i in 1:(length(which.plot))){

   # grab posterior sample from Gibbs
   smpl=lambda[which.plot[i],]

   # compute posterior density based on empirical Bayes
   xx=seq(min(smpl),max(smpl),len=100)
   yy=dgamma(xx,shape=x[which.plot[i]]+1,rate=theta.empBayes+1)

   # histogram of posterior sample, while allowing some extra space on the top
   hist(smpl,prob=TRUE,ylim=c(0,max(yy)),
     main=substitute(paste(lambda[ii], '; ', X[ii], '=', val),
                     list(ii=which.plot[i], val=x[which.plot[i]])),
     xlab='',ylab='',border=FALSE,col='grey')->obj

   # add posterior density based on empirical Bayes to compare
   lines(xx, yy, lwd=2, col='blue')
  }

  # plot posterior of theta
  hist(theta,prob=TRUE,main=expression(theta), 
    xlab='',ylab='',border=FALSE,col='grey') -> obj
  # add empirical Bayes estimate
  abline(v=theta.empBayes,col='blue',lwd=2)	 
} 

if (!plot.only){
mydata=c(rep(3,2),
         rep(2,8),
         rep(1,20),
         rep(0,70))		 
my.gibbs.PoissonExp(mydata, a=0.001, b=0.001)->junk
}

# total of 100 dimensions, so randomly choose 2 from each 'category' to plot
my.plot(junk, which.plot=sort(      
              c(1:2,
                sample(3:10,size=2),
                sample(11:30,size=2),
                sample(31:100,size=2))))