############################################################
# implementing metropolis-hastings (univariate setting)
#   1) define log_target(x)
#   2) define rprop(xn)
#   3) define log_r(y, xn)
#   4) run mh_1d(...)
############################################################

# requires mcmcse
install.packages('mcmcse')

# ---------- skeleton (used everywhere) ----------
mh_1d <- function(rprop, log_r, N = 20000, x0 = 0, burn = 0, seed = 1) {
  set.seed(seed)
  stopifnot(burn >= 0, burn < N)
  
  x <- numeric(N + 1); x[1] <- x0
  acc <- logical(N)
  
  for (n in 1:N) {
    xn <- x[n]
    y  <- rprop(xn)
    
    if (log(runif(1)) < min(0, log_r(y, xn))) {
      x[n + 1] <- y; acc[n] <- TRUE
    } else {
      x[n + 1] <- xn; acc[n] <- FALSE
    }
  }
  
  keep_idx <- (burn + 1):(N + 1)
  list(
    x = x[keep_idx],
    acc_rate = mean(acc),
    burn = burn,
    N = N
  )
}

# ---------- example: N(mu,1) with RW proposals ----------
mu <- 2.5

# define log_target(x) up to a constant
log_target <- function(x) -0.5 * (x - mu)^2

# define proposal y ~ q(.|xn)
s <- 1.0
rprop <- function(xn) rnorm(1, mean = xn, sd = s)

# define log acceptance ratio log_r(y, xn)
log_r <- function(y, xn) log_target(y) - log_target(xn)

# ---------- run MH ----------
N <- 50000
nburn <- 10000
x0 <- 0
out <- mh_1d(rprop, log_r, N = N, x0 = x0, burn = nburn, seed = 1)

# ---------- estimated mean ----------
x_mcmc <- out$x
mean(x_mcmc)

# ---------- diagnostics ----------
plot(x_mcmc, type = "l", xlab = "n", ylab = "xn", main = "trace plot")
abline(h = mu, lty = 2, col = "red")

acf(x_mcmc, main = "autocorrelation")

hist(x_mcmc, breaks = 40, main = "histogram", xlab = "x")
abline(v = mu, lty = 2, col = "red")

rm <- cumsum(x_mcmc) / seq_along(x_mcmc)
plot(rm, type = "l", main = "running mean", xlab = "iteration", ylab = "mean")
abline(h = mu, lty = 2, col = "red")

# ---------- compare with iid sample ----------
set.seed(1)
x_iid <- rnorm(N - burn, mean = mu, sd = 1)

# similar estimates
mean(x_iid)
mean(x_mcmc)

# but larger standard error for mcmc
mcmcse::mcse(x_iid)
mcmcse::mcse(x_mcmc)

# because autocorrelation -> lower effective sample size
mcmcse::ess(x_iid)
mcmcse::ess(x_mcmc)

# ---------- tuning the step ----------
# find the value that optimizes effective sample size
rprop <- function(xn) rnorm(1, mean = xn, sd = 0.5)
out <- mh_1d(rprop, log_r, N = 20000, x0 = 0, burn = nburn, seed = 1)
out$acc_rate
x_mcmc <- out$x
mcmcse::ess(x_mcmc)


############################################################
# MH template for notecard problems (fill in the blanks)
############################################################

# 1) Define the target (log, unnormalized)
log_target <- function(x) {
  ...  # notecard gives f(x) or log f(x) up to a constant
}

# 2) Define the proposal sampler Y ~ q(. | x)
rprop <- function(x) {
  ...  # e.g., rnorm(1, mean = x, sd = s) for RW proposal
}

# 3) Define the log acceptance ratio log r = log[ pi(y) q(x|y) / (pi(x) q(y|x)) ]
# IMPORTANT: arguments are (y, x) = (proposed, current)
log_r <- function(y, x) {
  ...  # usually: log_target(y) - log_target(x) + log_q(x|y) - log_q(y|x)
}

# 4) Choose starting value + tuning parameter(s)
x0   <- ...  # reasonable start (notecard may suggest one)
s    <- ...  # step size if using RW; otherwise set proposal parameters
burn <- ...  # e.g., 5000
N    <- ...  # e.g., 20000

# 5) Run MH
out <- mh_1d(
  rprop = rprop,
  log_r = log_r,
  N = N,
  x0 = x0,
  burn = burn,
  seed = 1
)

# 6) Diagnostics
out$acc_rate
plot(out$x, type="l", xlab="iteration (post-burn)", ylab="x", main="trace")
hist(out$x, breaks=40, main=NULL, xlab="x")
acf(out$x, main="ACF")

# 7) Additional notecard-specific goals, if any (usually an estimate of some kind)