#################################################################
# Uebung zu Schaetzen und Testen 1 im WiSe 15/16 
# R-Loesung zur Uebung 11, Aufgabe 22
# Autor: David Ruegamer
#################################################################

library(MASS)
library(MCMCpack)

set.seed(42)

### Simuliere Groessen

n = 100
p = 2
X = matrix(c(rep(1, n), runif(n*p, -10,10)), ncol=3)
mu0 = c(1,2,3)
Sigma0 = 1.2*diag(p+1)
a=2
b=1

beta <- mvrnorm(1, mu0, Sigma0)
sigmaSq <- rinvgamma(n, shape=a, scale=b)
eps <- rnorm(n, 0, sigmaSq)

y = as.numeric(X%*%beta + eps)

### Frequentistisches Ergebnis:

(betaML <- coef(lm(y ~ X - 1)))
(sigmaSqML <- summary(lm(y~X))$sigma)

### Erstelle FCs als Funktionen

fullCondSigmaSq <- function(a, b, y, X, beta){
    
  n = length(y)
  atilde = a + n/2
  btilde = as.numeric(b + 0.5*crossprod(y-X%*%beta))
  
  # 1 / gamma-verteilte ZV ergibt IG-verteilte ZV  
  rinvgamma(1, shape=(atilde), scale=(btilde))  
  
}

fullCondBeta <- function(sigmaSq, Sigma0, mu0, y, X){
  
  n = length(y)
  
  Sigma1 = solve( 1/sigmaSq*crossprod(X) + solve(Sigma0) )
  mu1 = Sigma1 %*% (1/sigmaSq * crossprod(X,y) + solve(Sigma0)%*%mu0)
  
  # eigtl. statt solve besser Inverse ueber QR oder Aehnliches berechnen!
  
  mvrnorm(1, mu1, Sigma1)  
  
}

### Erstelle Gibbs-Sampler

# hier mit Startwerten nur fuer beta

gs <- function(nrIter=100, startWertBeta = c(0,0,0), plot=TRUE, waitSec = 2){
  
  # nrIter = Anzahl an Iterationen
  # startWertBeta = Startwert(e) fuer beta-Koeffizienten
  # plot %in% c("beta","sigma","none") => welche MK soll geplottet werden
  # waitSec = Anzahl an Sekunden, die zwischen zwei Iterationen
  #           gewartet werden soll
  
  betaPost <- matrix(rep(NA,nrIter*(p+1)),ncol=p+1)
  sigmaSqPost <- rep(NA,nrIter)
  
  betaPost[1,] = startWertBeta
  sigmaSqPost[1] = fullCondSigmaSq(a,b,y,X,betaPost[1,])
  
  par(mfrow=c(2,2))
  
  if(plot){
    
    plot(betaPost[,1]~1, xlim=c(0,nrIter), type="l", xlab="Iteration",ylab=expression(beta[0]))
    plot(betaPost[,2]~1, xlim=c(0,nrIter), type="l", xlab="Iteration",ylab=expression(beta[1]))
    plot(betaPost[,3]~1, xlim=c(0,nrIter), type="l", xlab="Iteration",ylab=expression(beta[2]))
    plot(sigmaSqPost~1, xlim=c(0,nrIter), type="l", xlab="Iteration",ylab=expression(sigma^2))
    
  }
    
    
  Sys.sleep(waitSec)
  
  for(i in 2:nrIter){
    
    betaPost[i,] <- fullCondBeta(sigmaSqPost[i-1], Sigma0, mu0, y, X)
    sigmaSqPost[i] = fullCondSigmaSq(a,b,y,X,betaPost[i,])
    
    if(plot){
      
      plot(betaPost[,1]~1, xlim=c(0,nrIter), type="l", xlab="Iteration", ylab=expression(beta[0]))
      plot(betaPost[,2]~1, xlim=c(0,nrIter), type="l", xlab="Iteration", ylab=expression(beta[1]))
      plot(betaPost[,3]~1, xlim=c(0,nrIter), type="l", xlab="Iteration", ylab=expression(beta[2]))
      plot(sigmaSqPost~1, xlim=c(0,nrIter), type="l", xlab="Iteration", ylab=expression(sigma^2))
      
    }
    
    Sys.sleep(waitSec)
    
  }
  
  return(list(betaPost=betaPost,sigmaSqPost=sigmaSqPost))
  
}

### run

nIterShow = 50

res <- gs(nrIter=nIterShow,waitSec=1)

plot(res$betaPost[10:nIterShow,1]~c(10:nIterShow), type="l", ylab=expression(beta[0]))
plot(res$betaPost[10:nIterShow,2]~c(10:nIterShow), type="l", ylab=expression(beta[1]))
plot(res$betaPost[10:nIterShow,3]~c(10:nIterShow), type="l", ylab=expression(beta[2]))
plot(res$sigmaSqPost[10:nIterShow]~c(10:nIterShow), type="l", ylab=expression(sigma^2))

betaML
colMeans(res$betaPost[10:nIterShow,])