## beta distribution Be(a,b)
##
## density: p(x) = x^{a+1}(1-x)^{b+1} / B(a,b)
##   where B(a,b) is the beta function
##
## mean: a/(a+b)
## mode: (a-1)/(a+b-2)  -- for a>1,b>1
## variance: ab/((a+b)^2(a+b+1))


## rejection sampling

## we can sample uniform [0,1] where the envelope constant is given by
## the mode

beta.propto <- function(x,a,b) {
  if (0 <= x && x <= 1) return (x^(a-1)*(1-x)^(b-1))
  return (0)
}
beta.mode   <- function(a,b)   (a-1)/(a+b-2)
beta.mean   <- function(a,b)   a/(a+b)
beta.var    <- function(a,b)   a*b/((a+b)^2*(a+b+1))
beta.sd     <- function(a,b)   sqrt(beta.var(a,b))

sample.beta <- function(no.samples, a,b) {
  sample.one <- function(a,b) {
    envelope <- beta.propto(beta.mode(a,b), a, b)
    while (TRUE) {
      x <- runif(1)
      u <- runif(1)
      if (u*envelope < beta.propto(x, a, b)) return (x)
    }
  }
  replicate(no.samples, sample.one(a,b))  
}

comparison.plot <- function(a, b, no.samples) {
  grid.points <- seq(0,1,by=0.01)
  beta.dist <- dbeta(grid.points, a, b)
  
  opar <- par(mfrow=c(2,1))
  plot(grid.points, beta.dist, type='l')
  hist(sample.beta(no.samples,a,b), xlim=c(0,1), probability=TRUE)
  par(opar)
}

rejection.fraction <- function(a,b, n=100) {
  sample.fraction <- function(a,b) {
    envelope <- beta.propto(beta.mode(a,b), a, b)
    reject.count <- 0 ; trials <- 1
    while (TRUE) {
      x <- runif(1)
      u <- runif(1)
      if (u*envelope < beta.propto(x, a, b)) return (reject.count / trials)
      reject.count <- reject.count + 1
      trials <- trials + 1
    }
  }
  mean(replicate(n, sample.fraction(a,b)))
}

rejection.plot <- function(as, bs) {
  fracs <- outer(as,bs,function (as,bs) mapply(rejection.fraction,as,bs))
  filled.contour(as, bs, fracs)
}

estimate.plot <- function(p, a, b, no.estimates, no.samples) {
  estimates <- replicate(no.estimates, mean(p(no.samples, a, b)))
  hist(estimates, xlim=c(0,1))
  abline(v=beta.mean(a,b), col='red')
}

estimate.comparison.plot <- function(a,b, no.estimates, no.samples) {
  opar <- par(mfrow=c(2,1))
  estimate.plot(rbeta, a, b, no.estimates, no.samples)
  estimate.plot(sample.beta, a, b, no.estimates, no.samples)
  par(opar)
}



## importance sampling
comparison.plot <- function(a,b) {
  grid.points <- seq(0,1,by=0.01)
  beta.dist <- dbeta(grid.points, a, b)
  norm.dist <- dnorm(grid.points, mean=beta.mode(a,b), sd=beta.sd(a,b))

  plot(grid.points, beta.dist, type='l', col='blue',
       ylim=range(beta.dist,norm.dist), xlab='x', ylab='Density',
       main='Comparison of Beta and Normal densities')
  lines(grid.points, norm.dist, col='red')
  legend('topleft',
         legend=c(expression(Be(alpha,beta)), expression(N(mu,sigma))),
         fill=c('blue','red'))           
}

importance.sample <- function(no.samples, a,b) {
  bmean <- beta.mean(a,b)
  std.dev <- beta.sd(a,b)
  xs <- rnorm(no.samples, mean=bmean, sd=std.dev)
  betas <- sapply(xs, function(x) beta.propto(x, a, b))
  norms <- dnorm(xs, mean=bmean, sd=std.dev)
  ratios <-  betas / norms
  sum(xs * ratios) / sum(ratios)
}

estimate.plot <- function(p, a, b, no.estimates, no.samples) {
  estimates <- replicate(no.estimates, mean(p(no.samples, a, b)))
  hist(estimates, xlim=c(0,1))
  abline(v=beta.mean(a,b), col='red')
}

estimate.comparison.plot <- function(a,b, no.estimates, no.samples) {
  opar <- par(mfrow=c(2,1))

  estimates <- replicate(no.estimates, mean(rbeta(no.samples, a, b)))
  hist(estimates, xlim=c(0,1), main='True samples')
  abline(v=beta.mean(a,b), col='red')

  estimates <- replicate(no.estimates, importance.sample(no.samples, a, b))
  hist(estimates, xlim=c(0,1), main='Importance samples')
  abline(v=beta.mean(a,b), col='red')

  par(opar)
}


## MCMC
MCMC.sample <- function(no.samples, a, b) {
  propose.next <- function(x) {
    y <- rnorm(1, mean=x, sd=0.1)
    while (y < 0 || 1 < y) {
      if (y < 0) y = - y
      if (1 < y) y = 2 - y
    }
    y
  }

  acceptance.prob <- function(x,y) beta.propto(y,a,b)/beta.propto(x,a,b)

  sample.next <- function(x) {
    y <- propose.next(x)
    u <- runif(1)
    if (u < acceptance.prob(x,y)) return (y)
    else return (x)
  }

  samples <- vector('numeric', no.samples)
  samples[1] <- runif(1)
  for (i in 2:no.samples) {
    samples[i] <- sample.next(samples[i-1])
  }

  samples
}

comparison.plot <- function(a, b, no.samples) {
  grid.points <- seq(0,1,by=0.01)
  beta.dist <- dbeta(grid.points, a, b)
  
  opar <- par(mfrow=c(2,1))
  plot(grid.points, beta.dist, type='l')
  hist(MCMC.sample(no.samples,a,b), xlim=c(0,1), probability=TRUE)
  par(opar)
}