# Introduction to Bayesian modeling
# Rick Condit, 19 March 2006
# NCEAS-UCSB

# Functions used

# For use with metrop1step and dbinom
binomialMetrop <- function(testtheta, S, N)
{
 if(testtheta<=0 | testtheta>=1) return(-Inf)
 return(dbinom(x=S, size=N, prob=testtheta, log=TRUE))
}


# A function to work through a Metropolis for a single-parameter binomial distribution.
showBinomialMCMC <- function(thetastart, S, N, steps=25, stepsize=0.1, loggraph=FALSE, pause=TRUE, showstep=1)
{
 theta <- seq(0, 1, by=0.01)
 postlike <- dbinom(x=S, size=N, prob=theta, log=loggraph)

 graphics.off()

 x11(height=8, width=8, xpos=900, ypos=1)
 plot(theta, postlike, type="l", main="binomial pseudo-posterior", ylab="prob")

 thetaMCMC <- numeric()
 thetaMCMC[1] <- thetastart
 points(thetaMCMC[1], postlike[theta==thetaMCMC[1]], col="blue", cex=1.7)
 points(thetaMCMC[1], postlike[theta==thetaMCMC[1]], col="blue", pch=16)
 if(pause) browser()

 for(i in 2:steps)
  {
   metropResult <- metrop1step(func=binomialMetrop, start.param=thetaMCMC[i-1], scale.param=stepsize, adjust=1.01, target=0.25, S=S, N=N)

   thetaTest <- metropResult[6]
   origLike <- dbinom(x=S, size=N, prob=thetaMCMC[i-1], log=loggraph)

   if(thetaTest>0 & thetaTest<1) newLike <- dbinom(x=S, size=N, prob=thetaTest, log=loggraph)
   else newLike <- (-Inf)

   if(loggraph) likeratio <- newLike-origLike
   else likeratio <- newLike/origLike

   points(thetaTest, newLike, col="red", cex=1.7)
   if(i%%showstep==0)
     cat("Start at ", round(thetaMCMC[i-1], 3), "(", round(origLike, 3), ")", round(thetaTest, 3), "(", round(newLike, 3), ") ; Ratio: ",
         round(likeratio, 3), "\n")
   if(pause) browser()

   thetaMCMC[i] <- metropResult[1]
   acceptLike <- dbinom(x=S, size=N, prob=thetaMCMC[i], log=loggraph)

   if(metropResult[3]==0) points(thetaTest, newLike, col="white", cex=1.7)
   else points(thetaMCMC[i], acceptLike, col="blue", cex=1.7)
   if(pause) browser()

   points(thetaMCMC[i-1], origLike, col="blue", cex=1, pch=16)
   points(thetaMCMC[i], acceptLike, col="blue", cex=1, pch=16)
   if(pause) browser()
  }

 x11(height=8, width=8, xpos=900, ypos=500)
 hist(thetaMCMC, breaks=theta)

 x11(height=8, width=8, xpos=200, ypos=1)
 plot(thetaMCMC)

 return(thetaMCMC)
}


# Takes a single metropolis step on a single parameter for any given likelihood function.
# The arguments start.param and scale.param are atomic (single values), as are adjust and target. 
# The ellipses handle all other arguments to the function. The function func must accept the test 
# parameter as the first argument, plus any additional arguments which come in the ellipses.

# Note the metropolis rule: if rejected, the old value is returned to be re-used. The return value
# includes a one if accepted, zero if rejected.

# The step size, refered to as scale.param, is adjusted following Helene's rule. 
# For every acceptance, scale.param is multiplied
# by adjust, which is a small number > 1 (1.01, 1.02, 1.1 all seem to work). For every rejection, scale.param
# is multiplied by (1/adjust) raised to a power (AdjExp) that is based on the target acceptance rate.
# When the target acceptance rate is 0.25, which is recommended for any model with > 4 parameters,
# AdjExp=3. It's easy to see how this system arrives at an equilibrium acceptance rate=target.

# The program calling metrop1step has to keep track of the scaling parameter: submitting it each time
# metrop1step is called, and saving the adjusted value for the next call. Given many parameters, a
# scale must be stored separately for every one.

# Note the return value is a vector of 6:
# 1) the new parameter value;
# 2) the new scale (step size);
# 3) a zero or a one to keep track of the acceptance rate;
# 4) the likelihood of original parameter (if rejected) or new parameter (if accepted)
# 5) the likelihood of original parameter (if accepted) or new parameter (if rejected)
# 6) the new parameter tested (whether accepted or not)

metrop1step <- function(func, start.param, scale.param, adjust, target, ...)
{
 origlike <- func(start.param, ...)
 newval <- rnorm(1, mean=start.param, sd=scale.param)
 newlike <- func(newval, ...)
 
 AdjExp <- (1-target)/target
  
 likeratio <- exp(newlike-origlike)
 if(runif(1)<likeratio) 
  {
   newscale <- scale.param*adjust^AdjExp
   return(c(newval, newscale, 1, newlike, origlike, newval))
  }
 else
  {
   newscale <- scale.param*(1/adjust)
   return(c(start.param, newscale, 0, origlike, newlike, newval))
  }

}


# Use a Gibbs sampler to alternately update the mu and SD estimates for a normal distribution, using Metropolis steps.
normalPosterior.Gibbs <- function(data, steps, showstep=10, pause=TRUE, adjuststep=FALSE)
{
 muMCMC <- sdMCMC <- stepsizeMu <- stepsizeSD <- acceptMu <- acceptSD <- llikeMu <- llikeSD <- rep(1, steps)
 muMCMC[1] <- mean(data)
 sdMCMC[1] <- sd(data)
 stepsizeMu[1] <- mean(data)/10
 stepsizeSD[1] <- sd(data)/10
 llikeMu[1] <- llikeSD[1] <- sum(dnorm(data, mean=muMCMC[1], sd=sdMCMC[1], log=TRUE))

 for(i in 2:steps)
  {
   metropMu <- metrop1step(func=dnormMu, start.param=muMCMC[i-1], x=data, sd=sdMCMC[i-1], scale.param=stepsizeMu[i-1], adjust=1.01, target=0.25)
   muMCMC[i] <- metropMu[1]
   acceptMu[i] <- metropMu[3]
   if(adjuststep) stepsizeMu[i] <- metropMu[2]
   else stepsizeMu[i] <- stepsizeMu[i-1]
   llikeMu[i] <- metropMu[4]

   if(pause) browser()

   metropSD <- metrop1step(func=dnormSD, start.param=sdMCMC[i-1], x=data, mean=muMCMC[i], scale.param=stepsizeSD[i-1], adjust=1.01, target=0.25)
   sdMCMC[i] <- metropSD[1]
   acceptSD[i] <- metropSD[3]
   if(adjuststep) stepsizeSD[i] <- metropSD[2]
   else stepsizeSD[i] <- stepsizeSD[i-1]
   llikeSD[i] <- metropSD[4]

   if(pause) browser()

   if(i%%showstep==0) cat("mu: ", round(muMCMC[i], 2), "sd: ", round(sdMCMC[i], 2), "Llike: ", round(metropSD[4], 3), "\n")
  }

 return(data.frame(muMCMC, sdMCMC, stepsizeMu, stepsizeSD, llikeMu, llikeSD, acceptMu, acceptSD))
}



# A function that returns the normal density, but accepts the mean as the first parameter. This allows metrop1step to
# create a posterior distribution for the mean.
dnormMu <- function(mean, x, sd) return(sum(dnorm(x, mean, sd, log=TRUE)))

# A function that returns the normal density, but accepts the SD as the first parameter. This allows metrop1step to
# create a posterior distribution for the mean.
dnormSD <- function(sd, x, mean)
{
 if(sd<=0) return(-Inf)
# browser()
 return(sum(dnorm(x, mean, sd, log=TRUE)))
}


# A function that calculates normal density with multiple observations and multiple alternate values of the mean
calcNormalManyMu <- function(x, mean, sd)
{
 loglike <- matrix(ncol=length(mean), nrow=length(x))
 for(i in 1:length(x)) loglike[i, ] <- dnorm(x[i], mean=mean, sd=sd, log=TRUE)
 return(colSums(loglike))
}

# A function that calculates normal density with multiple observations and multiple alternate values of the SD
calcNormalManySD <- function(x, mean, sd)
{
 loglike <- matrix(ncol=length(sd), nrow=length(x))
 for(i in 1:length(x)) loglike[i, ]=dnorm(x[i], mean=mean, sd=sd, log=TRUE)
 return(colSums(loglike))
}



# The standard Gibbs sampler for a normal distribution with unknown mean and variance.
# These need to be called alternately to estimate both parameters. A main function to carry out
# alternate calls is not included.

# Gibbs draw for the mean, given an estimated standard deviation sigma
Gibbs.normalmean <- function(y, sigma)
{
 y <- y[!is.na(y)]
 if(sd(y)==0) return(mean(y))

 samplemean <- mean(y)
 n <- length(y)

 return(rnorm(1, mean=samplemean, sd=sqrt(sigma^2/n)))
}

# Gibbs draw for the SD (http://www.biostat.jhsph.edu/~fdominic/teaching/BM/3-4.pdf)
Gibbs.normalSD <- function(y)
{
 y <- y[!is.na(y)]
 if(sd(y)==0) return(0)
 
 n <- length(y)
 vr <- var(y)

 newvar <- rinvgamma(1, shape=(n-1)/2, scale=(n-1)*vr/2)

 return(sqrt(newvar))
}