#---------------------------------------------
# inferring p, predicting n1  --> Metropolis
#
# GdA, May 2020
#---------------------------------------------

dbg <- !TRUE

#model parameters
n0 = 20
x0 = 10
n1 = 10

# Metropolis parameters
N = 10000
Dx = 2
Dp = 0.1

# unnormalized distribution 
uf <- function(p, x1, n0, x0, n1) {
  if(p<0 || p>1)     return(0)
  if(x1 <0 || x1>n1) return(0)
  return( p^x0*(1-p)^(n0-x0) * p^x1*(1-p)^(n1-x1) /
          (factorial(x1)*factorial(n1-x1)) ) 
}

# proposal functions
pr.p  <- function(p.o, Dp) p.o + runif(1, -Dp, +Dp)  
pr.x1 <- function(x1.o, Dx) x1.o + sample(-Dx:Dx)[1]

# inits
p <- x1 <- numeric(N)

# initial p and x1 
p[1] = rbeta(1,1,1)           # uniform
x1[1] = rbinom(1, n1, p[1])
if(dbg) cat(sprintf("1: %f, %d\n", p[1], x1[1]))

# Metropolis
for (i in 2:N) {
   p.p   <- pr.p(p[i-1], Dp)     # proposals
   x1.p  <- pr.x1(x1[i-1], Dx)
   if(dbg) cat(sprintf("Proposals: %f, %d\n", p.p, x1.p))
   A <- min(1, uf(p.p, x1.p, n0, x0, n1) /
               uf(p[i-1], x1[i-1], n0, x0, n1)  )
   if(dbg) cat(sprintf("A: %f", A))
   if ( runif(1) <= A ) {
      if(dbg) cat("  -> moves \n")     
      p[i]  <- p.p
      x1[i] <- x1.p
   } else {
      if(dbg) cat("  -> stays \n")        
      p[i]  <- p[i-1]
      x1[i] <- x1[i-1] 
   }
   if(dbg) cat(sprintf("%d: %f, %d\n\n", i, p[i], x1[i]))
   if(dbg) pausa()
}

# analysis and plots of the chain
par(mfcol=c(2,2))
hist(p, nc=50, col='cyan', xlim=c(0,1))
barplot(table(x1)/N, col='cyan', xlab='x1')    #,xlim=c(0,n1))
plot(x1, p, col='blue', cex=1)
plot(NULL,  xaxt = 'n', yaxt = 'n', bty = 'n', pch = '',
     ylab = '', xlab = '', xlim=c(0,10), ylim=c(0,10))
text(4,9, sprintf("mean(p) = %.3f\n sigma(p) =  %.3f",
                  mean(p), sd(p)), cex=1.5,  col='blue')
text(4,5, sprintf("mean(x1) = %.3f\n sigma(x1) =  %.3f",
                  mean(x1), sd(x1)), cex=1.5, col='blue')
text(4,1, sprintf("rho(p,x1) = %.3f", cor(p,x1)), cex=1.5, col='blue')
par(mfcol=c(1,1))