# ============================================================================
# Metropolis-Hastings: two targets to code up
# ============================================================================
#
# For each one: work out the acceptance ratio, fill in the three functions,
# run it, tune the proposal, and estimate the quantity asked for.

metropolis <- function(x0, n, log_density, propose, log_ratio) {
  draws <- matrix(NA_real_, nrow = n, ncol = length(x0))
  draws[1, ] <- x0
  accepted <- 0

  for (t in 2:n) {
    current  <- draws[t - 1, ]
    proposed <- propose(current)
    if (log(runif(1)) < log_ratio(current, proposed, log_density)) {
      draws[t, ] <- proposed
      accepted <- accepted + 1
    } else {
      draws[t, ] <- current
    }
  }

  list(draws = drop(draws), rate = accepted / (n - 1))
}

ess <- function(x) {
  rho <- acf(x, lag.max = 1000, plot = FALSE)$acf[-1]
  cut <- which(rho < 0.05)[1]
  if (is.na(cut)) cut <- length(rho)
  length(x) / (1 + 2 * sum(rho[1:cut]))
}

look <- function(ch, burn = 1000, target = NULL) {
  x <- ch$draws[-(1:burn)]
  plot(x, type = "l", xlab = "t", ylab = "x")
  acf(x, lag.max = 60, main = "")
  hist(x, breaks = 50, freq = FALSE, main = "", xlab = "x")
  if (!is.null(target)) curve(target(x), add = TRUE, lwd = 2)
  c(rate = ch$rate, ess = ess(x), mean = mean(x))
}

## ---------------------------------------------------------------------------
## EXAMPLE: WHAT THE THREE FUNCTIONS LOOK LIKE
## ---------------------------------------------------------------------------
##
## Target is standard normal, proposal is a symmetric Gaussian random walk.
## Each exercise below needs these same three pieces and nothing else. The
## names are yours to choose; they are passed to metropolis() by position.

log_target <- function(x) dnorm(x, log = TRUE)

propose <- function(x) rnorm(1, x, 2.4)

log_ratio <- function(x, y, log_density) log_density(y) - log_density(x)

chain <- metropolis(0, 50000, log_target, propose, log_ratio)

look(chain, target = dnorm)

## The proposal is symmetric, so q cancels and log_ratio is just a difference
## of log densities. Neither target below has that luxury: the first is not a
## density you can call, and the second has an asymmetric proposal.

## ---------------------------------------------------------------------------
## 1. A SURPRISING OBSERVATION
## ---------------------------------------------------------------------------
##
## Prior  theta ~ Cauchy(0, 1)
## Data   y | theta ~ N(theta, 1),  and you observe y = 6
##
## Estimate the posterior mean of theta. Before running anything, write down
## where you expect it to be. Had the prior been N(0, 1) the answer would be
## exactly 3.

log_posterior <- function(theta) {
  ...
}

propose <- function(theta) {
  ...
}

log_ratio <- function(x, y, log_density) {
  ...
}

chain <- metropolis(0, 50000, log_posterior, propose, log_ratio)

look(chain)

## ---------------------------------------------------------------------------
## 2. A RATE WITH NO CONJUGATE PRIOR
## ---------------------------------------------------------------------------
##
## Counts  y ~ Poisson(lambda), observed below
## Prior   lambda ~ LogNormal(0, 1)
##
## A Gamma prior would be conjugate; this one is not. Since lambda > 0, use a
## multiplicative proposal, lambda* = lambda * exp(sigma * Z) with Z ~ N(0, 1).
## This proposal is NOT symmetric, so work out the Hastings terms.
##
## Estimate the posterior mean of lambda and P(lambda > 4).

y <- c(3, 1, 4, 1, 5)

log_posterior <- function(lambda) {
  ...
}

propose <- function(lambda) {
  ...
}

log_ratio <- function(x, y, log_density) {
  ...
}

chain <- metropolis(2, 50000, log_posterior, propose, log_ratio)

look(chain)

## ---------------------------------------------------------------------------
## Checks
## ---------------------------------------------------------------------------
##
## Both targets are one-dimensional, so you can integrate numerically. Compare
## your estimate against quadrature, and your proposal scale against the
## acceptance rate you were aiming for.