# Metropolis for normal mu and sigma:   metropolis.v2
# edps 590 BAY
# Fall 2019
# C.J.Anderson
#
# Set:
#  y = name of variable that want the mean and sigma of
#
# Input:
#  start = length 2 with the starting value of mu and sigma
#  jump.sd = value of standard deviation 
#  iterations= number of iterations to run
#
# Output:
#   An mcmc object where Var 1 is the mean and Var 2 is the variance
#
# Notes: 
#     o starting values for mu should be between min(y) and max(y)
#
# 
# Example are after code for function
#


library(coda)

## log Likelihood
lnlike <- function(parm) {
  mu = parm[1]
  std= parm[2]
  pred = mu
  onelike =  dnorm(y,mean=pred, sd=std, log=TRUE) 
  lnlikes = sum(onelike)
  return(lnlikes)
  }

## Prior
prior = function(parm) {
    mu = parm[1]
    std= parm[2]
    lo = min(y)                                      
    hi = max(y)                                      
    s.max = 10*std**2                                   
    mu.prior  = dunif(mu,min=lo,max=hi,log=TRUE)
    std.prior = dunif(std,min=0,max=s.max)
    return(mu.prior+std.prior)
    }

## Proposal 
jump.dist = function(parm,jump.sd) {
    return(rnorm(2,mean=parm,sd=c(jump.sd,jump.sd)))  
    }

## Metropolis alogrithm
metroNorm2 = function(start,jump.sd,iterations) {
    chain = array(dim= c(iterations+1,2) )
    chain[1,] = start
    for (i in 1:iterations) {
       propose = jump.dist(chain[i,],jump.sd)
       r = lnlike(propose) + prior(propose) - lnlike(chain[i,]) - prior(chain[i,])
   u <-  log( runif(1)  )
   if (r < u ) {chain[i+1,] = chain[i,] }
   if (r > u ) {chain[i+1,] = propose } 
   
   }
    return(mcmc(chain))
}

###########################################################
# Make some data
mu = 2
std= 1
N  = 20


# Values that need to be set
y <- rnorm(N,mean=mu,sd=std)  
  
jump.sd <- .2

# Needed for input
start <- c(4.00,1.0)      # make sure start min(y) < mu start < max(y)
iterations <- 5000

chain1 <- metroNorm2(start,jump.sd,iterations)

#
# How similar are sample statistics and posterior statistics?
#
mean(chain1[200:iterations,1])
mean(y)
mean(chain1[200:iterations,2])
var(y)

#
#  Random walk of first 20 iterations w/ labels
#
par(mfrow=c(1,1))
plot(chain1[1:20,1],chain1[1:20,2],type="l",
     main="Random Walk of 1st 20",
	 xlab="Mean of Y",
	 ylab="Variance of Y",
	 )
ipoint <- seq(1:20)
text(chain1[1:20,1], chain1[1:20,2], labels=ipoint,offset=10,cex=2)

#
#  Random walk of first 20 iterations
#
par(mfrow=c(2,2))
plot(chain1[1:20,1],chain1[1:20,2],type="l",
     main="Random Walk of 1st 20",
	 xlab="Mean of Y",
	 ylab="Variance of Y",
	 xlim=c(1,4),
	 ylim=c(.5,2.5)
	 )
#
#  Random walk of first 100 iterations
#
plot(chain1[1:100,1],chain1[1:100,2],type="l",
     main="Random Walk of 1st 100",
	 xlab="Mean of Y",
	 ylab="Variance of Y",
	 xlim=c(1,4),
	 ylim=c(.5,2.5)
	 )
#
#  Random walk of first 500 iterations
#
plot(chain1[1:500,1],chain1[1:500,2],type="l",
     main="Random Walk of 1st 500",
	 xlab="Mean of Y",
	 ylab="Variance of Y",
	 xlim=c(1,4),
	 ylim=c(.5,2.5)
	 )

#
#  Random walk of all iterations
#
plot(chain1[1:5000,1],chain1[1:5000,2],type="l",
     main="Random Walk of All",
	 xlab="Mean of Y",
	 ylab="Variance of Y"
	 )

#
#  Random walk of all iterations
#
par(mfrow=c(1,1))
plot(chain1[1:5000,1],chain1[1:5000,2],type="l",
     main="Random Walk Blown Up",
	 xlab="Mean of Y",
	 ylab="Variance of Y"
	 )
	 
#
#  Some diagonistics 
#

par(mfrow=c(2,2))
traceplot(chain1)


par(mfrow=c(2,1))
par(ask=TRUE)
traceplot(chain1,smooth=FALSE,type='l', xlab='Iterations',ylab="Parameter Values ")

autocorr.plot(chain1, iterations, auto.layout=TRUE)

geweke.diag(chain1,frac1=0.1,frac2=0.5)


ess<- effectiveSize(chain1)
ess

# what if we treat first 200 as burnin
chain1.b <- chain1[200:5000,]
geweke.diag(chain1.b,frac1=0.1,frac2=0.5)
ess<- effectiveSize(chain1.b)
ess

cumuplot(chain1,probs=c(.25,.50,.75),lwd=c(1,2),lty=c(2,1),col=c("blue","red"), main="25th, 50th and 75th percentiles")


######################################################
# Increase chain length to 20,000 and run 4 chains
######################################################

iterations<- 20000
start1 <- c(4.00,1.0)      # make sure start min(y) < mu start < max(y)
start2 <- c(0.5,2.0)      # make sure start min(y) < mu start < max(y)
start3 <- c(1,1.0)      # make sure start min(y) < mu start < max(y)
start4 <- c(2,4.0)      # make sure start min(y) < mu start < max(y)

chain1 <- metroNorm2(start1,jump.sd,iterations)
chain2 <- metroNorm2(start2,jump.sd,iterations)
chain3 <- metroNorm2(start3,jump.sd,iterations)
chain4 <- metroNorm2(start4,jump.sd,iterations)

x<- seq(1:20001)
plot(x,chain1[x,1],col="blue",type='l',main="Mixing of Means",ylim=c(.5,4.0),ylab="Estimated Mean",xlab="iteration")
lines(x,chain2[x,1],col="red",type='l')
lines(x,chain3[x,1],col="green",type='l')
lines(x,chain4[x,1],col="orange",type='l')

plot(x,chain1[x,2],col="blue",type='l',main="Mixing of Variances",ylim=c(.5,4.0),ylab="Estimated Mean",xlab="iteration")
lines(x,chain2[x,2],col="red",type='l')
lines(x,chain3[x,2],col="green",type='l')
lines(x,chain4[x,2],col="orange",type='l')

#
# Combine chains into one ojbect
#
all.chains <- mcmc.list(chain1,chain2,chain3,chain4)

# library("coda")
gelman.plot(all.chains)

cumuplot(all.chains, probs=c(0.025,0.5,0.975), ylab="",lty=c(2,1), lwd=c(1,2), 
   type="l", ask=dev.interactive(),auto.layout=TRUE)

autocorr.plot(all.chains, iterations, auto.layout=TRUE)

geweke.diag(all.chains,frac1=0.1,frac2=0.5)

effectiveSize(all.chains)

summary(all.chains)