# Real data: mcycle (MASS)
# A data frame giving a series of measurements of head acceleration
# in a simulated motorcycle accident, used to test crash helmets.
# x: time in miliseconds after impact
# y: head accelaration (in g)
# Source: Silverman (1985) JRSS, 47,1-52.

rm(list=ls())
install.packages("mvtnorm")
library("MASS")
library("mvtnorm")


n = nrow(mcycle)
x = mcycle$times
y = mcycle$accel
xt = x[1:132]
yt = y[1:132]
n = length(xt)
xt = (xt-mean(xt))/sqrt(var(xt))
yt = (yt-mean(yt))/sqrt(var(yt))

plot(xt,yt,xlab="time in miliseconds after impact",ylab="head accelaration (in g)")
legend("topleft",legend=c(
"Real data: mcycle (MASS)",
"A data frame giving a series of measurements of",
"head acceleration in a simulated motorcycle accident,",
"used to test crash helmets.",
"x: time in miliseconds after impact",
"y: head accelaration (in g)",
"Source: Silverman (1985) JRSS, 47,1-52."),bty="n",cex=0.75)

p = 14
X = xt

for (i in 2:p)
X = cbind(X,xt^i)


###############
# ML estimation
###############
r2 = rep(0,p)
adj.r2 = rep(0,p)
for (i in 1:p){
  fit = lm(yt~X[,1:i])
  r2[i] = summary(fit)$r.squared
  adj.r2[i] = summary(fit)$adj.r.squared
}
ind   = 1:p
ord1  = ind[adj.r2==max(adj.r2)]  
ord3  = ind[r2==max(r2)]  
  

fit1 = lm(yt~X[,1:ord1])
fit3 = lm(yt~X[,1:ord3])

par(mfrow=c(1,2))
plot(adj.r2,xlab="Order of the polynomial regression",ylab="R2 and adjusted R2",pch=16,ylim=range(adj.r2,r2))
points(r2,pch=16,col=2)
abline(v=ord1)
abline(v=ord3,col=2)
legend("topleft",legend=c("R2","Adjusted R2"),col=2:1,pch=16,bty="n")
plot(xt,yt,xlab="time in miliseconds after impact",ylab="head accelaration (in g)")
points(xt,fit$fit,col=2,pch=16)
legend("bottomright",legend=paste("Order of the polynomial=",ord1,sep=""),bty="n")


#######################
# Bayesian estimation
#######################
Bayes.regression = function(y,X,b0,B0,nu0,s20){
  p      = ncol(X)
  nu0s20 = nu0*s20
  iB0    = solve(B0)
  iB0b0  = iB0%*%b0
  B1     = solve(iB0+t(X)%*%X)
  tcB1   = t(chol(B1))
  b1     = B1%*%(iB0b0+t(X)%*%y)
  nu1    = nu0 + (n-p)
  s21    = (nu0s20 + t(y-X%*%b1)%*%y + t(b0-b1)%*%iB0b0)/nu1
  se     = sqrt(diag(s21[1,1]*B1))
  m      = X%*%b0
  V      = s20*(diag(1,n)+X%*%B0%*%t(X))
  pred   = dmvt(y,delta=m,sigma=V,df=nu0,log = TRUE)
  return(list(b1=b1,B1=B1,nu1=nu1,s21=s21,se=se,pred=pred))
}

b0  = lm(yt~X)$coef
B0  = diag(100,p+1)
nu0 = 5
s20 = 1
XX = cbind(1,X)

pred = rep(0,p+1)
for (i in 2:(p+1))
  pred[i] = Bayes.regression(yt,XX[,1:i],b0[1:i],B0[1:i,1:i],nu0,s20)$pred
ord2 = ind[pred==max(pred[2:(p+1)])]  

par(mfrow=c(1,2))
plot(adj.r2,xlab="Order of the polynomial regression",ylab="R2 and adjusted R2",pch=16,ylim=range(adj.r2,r2))
title("Maximum likelihood estimation")
points(r2,pch=16,col=2)
abline(v=ord1,lwd=2)
abline(v=ord3,col=2,lwd=2)
legend("topleft",legend=c("R2","Adjusted R2"),col=2:1,pch=16,bty="n")
plot(2:p,pred[2:p],xlab="Order of the polynomial regression",ylab="Log prior predictive",pch=16)
abline(v=ord2,col=2,lwd=2)
title("Bayesian inference")


fit.bayes = Bayes.regression(yt,XX[,1:ord2],b0[1:ord2],B0[1:ord2,1:ord2],nu0,s20)
b1    = fit.bayes$b1
B1    = fit.bayes$B1
draws = rmvnorm(1000,mean=b1,sigma=B1)
funcs = XX[,1:ord2]%*%t(draws)
qfunc = apply(funcs,1,quantile,c(0.025,0.5,0.975))

par(mfrow=c(1,1))
plot(xt,yt,xlab="time in miliseconds after impact",ylab="head accelaration (in g)",pch=16)
lines(xt,fit1$fit,col=2,lwd=2)
lines(xt,fit3$fit,col=3,lwd=2)
lines(xt,qfunc[1,],col=4,lwd=2,lty=2)
lines(xt,qfunc[2,],col=4,lwd=2)
lines(xt,qfunc[3,],col=4,lwd=2,lty=2)
legend("bottomright",legend=c(
paste("ML(Adj.R2)     - order of the polynomial=",ord1,sep=""),
paste("ML(R2) - order of the polynomial=",ord3,sep=""),
paste("Bayes      - order of the polynomial=",ord2,sep="")),col=2:4,pch=16,bty="n")




pmp=exp(pred[2:(p+1)])/sum(exp(pred[2:(p+1)]))

tab = cbind(2:(p+1),100*round(pmp,3))
colnames(tab) = c("order","PMP")
tab