#install.packages("bayesreg")
#library("bayesreg")
#library("mvtnorm")

set.seed(1233456)
x.iid    = FALSE
n        = 500
p        = 50
k        = 5
sig.true = 5
beta.true = c(rep(1,k),rep(0,p-k))

if (x.iid){
  X = rmvnorm(n)
}else{
  B = matrix(runif(p*q),p,q)
  S = diag(1,p)
  Var = B%*%t(B)+S
  X = rmvnorm(n,sigma=Var)
}
y = X%*%beta.true + rnorm(n,0,sig.true)


par(mfrow=c(1,1))
plot(lm(y~X-1)$coef,ylab="beta",xlab="index",pch=16)
points(beta.true,col=2,pch=16)


beta.mle = solve(t(X)%*%X)%*%t(X)%*%y
y.hat    = X%*%beta.mle
e.hat    = y-y.hat
sig.mle  = sqrt(mean(e.hat^2))

table = round(rbind(cbind(beta.mle,lm(y~X-1)$coef),
                    c(sig.mle,(n-p)/n*summary(lm(y~X-1))$sigma)),4)
colnames(table)      = c("Our function","lm function")
rownames(table)[1]   = "constant"
rownames(table)[p+1] = "sigma"
table


summary(lm(y~X-1))




# Prior knowledge 
b = rep(0,p)
V = diag(100,p)
c = 5/2
d = 5/2

# Gibbs sampler
sigma2 = 1 
M0     = 1000
M      = 1000
thin   = 1
niter  = M0+M*thin
betas  = matrix(0,niter,p)
sigs   = rep(0,niter)
for (iter in 1:niter){
  V1 = solve(t(X)%*%X/sigma2+solve(V))
  b1 = V1%*%(t(X)%*%y/sigma2+solve(V)%*%b)
  beta = b1 + t(chol(V1))%*%rnorm(p)
  c1 = c+n/2
  d1 = d + sum((y-X%*%beta)^2)/2
  sigma2 = 1/rgamma(1,c1,d1)
  betas[iter,] = beta
  sigs[iter] = sqrt(sigma2)
}
ind = seq(M0+1,niter,by=thin)
betas = betas[ind,]
sigs = sigs[ind]

mean.beta = apply(betas,2,mean)
table = round(cbind(beta.mle,mean.beta),4)
colnames(table) = c("ML","Bayes")
rownames(table)[1] = "Constant"
table


par(mfrow=c(3,4))
for (i in 1:11){
  ts.plot(betas[,i],ylab="",main=paste("beta(",i,")",sep=""))
abline(h=beta.true[i],col=2)
}
ts.plot(sigs,ylab="",main="sigma")

par(mfrow=c(3,4))
for (i in 1:11){
  hist(betas[,i],xlab="",prob=TRUE,main=paste("beta(",i,")",sep=""))
  abline(v=beta.true[i],col=2,lwd=2)
}
hist(sigs,xlab="",prob=TRUE,main="sigma")
abline(v=sig.true,col=2,lwd=2)


# Shrinkage/sparsity priors: ridge, Laplace & horseshoe
M0    = 1000
M     = 1000
skip  = 10
data1 = data.frame(y=y,X=X)

fit.r = bayesreg(y~X,data=data1,model="normal",prior="ridge",n.samples=M,burnin=M0,thin=skip)
fit.l = bayesreg(y~X,data=data1,model="normal",prior="lasso",n.samples=M,burnin=M0,thin=skip)
fit.h = bayesreg(y~X,data=data1,model="normal",prior="horseshoe",n.samples=M,burnin=M0,thin=skip)

quant.r = t(apply(fit.r$beta,1,quantile,c(0.05,0.5,0.95)))
quant.l = t(apply(fit.l$beta,1,quantile,c(0.05,0.5,0.95)))
quant.h = t(apply(fit.h$beta,1,quantile,c(0.05,0.5,0.95)))



range = range(quant.r,quant.l,quant.h)
par(mfrow=c(2,2))
plot(quant.r[,2],lwd=2,pch=16,ylim=range,ylab="beta")
title("Ridge regression")
for (i in 1:p)
  segments(i,quant.r[i,1],i,quant.r[i,3])
points(beta.true,col=4,pch=16)

plot(quant.l[,2],lwd=2,pch=16,ylim=range,ylab="beta")
title("Lasso regression")
for (i in 1:p)
  segments(i,quant.l[i,1],i,quant.l[i,3])
points(beta.true,col=4,pch=16)

plot(quant.r[,2],lwd=2,pch=16,ylim=range,ylab="beta")
title("Horseshoe regression")
for (i in 1:p)
  segments(i,quant.h[i,1],i,quant.h[i,3])
points(beta.true,col=4,pch=16)

range = sqrt(range(fit.r$sigma2,fit.l$sigma2,fit.h$sigma2))
plot(density(sqrt(fit.r$sigma2)),xlab="",main="Sigma")
lines(density(sqrt(fit.l$sigma2)),col=2)
lines(density(sqrt(fit.h$sigma2)),col=3)
abline(v=sig.true,col=4,lwd=3)


par(mfrow=c(3,5))
for (i in 1:15){
plot(density(fit.r$beta[i,]),main="")
lines(density(fit.l$beta[i,]),col=2)
lines(density(fit.h$beta[i,]),col=3)
abline(v=beta.true[i],col=4)
}


names = c("ridge","lasso","horseshoe")
par(mfrow=c(3,5))
for (i in 1:15){
  boxplot(fit.r$beta[i,],fit.l$beta[i,],fit.h$beta[i,],outline=FALSE,names=names,main="")
  title(paste("beta(",i,")",sep=""))
  abline(h=beta.true[i],col=4)
}














################################################
# Monte Carlo exercise
################################################
set.seed(1233456)
x.iid    = FALSE
n        = 200
p        = 50
k        = 5
sig.true = 5
beta.true = c(rep(1,k),rep(0,p-k))

nrep  = 10
error = array(0,c(3,nrep,p))
S     = diag(1,p)
for (r in 1:nrep){
  if (x.iid){
    X = rmvnorm(n,sigma=S)
  }else{
    B = matrix(runif(p*q),p,q)
    Var = B%*%t(B)+S
    X = rmvnorm(n,sigma=Var)
  }
  y = X%*%beta.true + rnorm(n,0,sig.true)
  M0    = 2000
  M     = 2000
  skip  = 1
  data1 = data.frame(y=y,X=X)
  fit.r = bayesreg(y~X,data=data1,model="normal",prior="ridge",n.samples=M,burnin=M0,thin=skip)
  fit.l = bayesreg(y~X,data=data1,model="normal",prior="lasso",n.samples=M,burnin=M0,thin=skip)
  fit.h = bayesreg(y~X,data=data1,model="normal",prior="horseshoe",n.samples=M,burnin=M0,thin=skip)
  error[1,r,] = fit.r$mu.beta-beta.true
  error[2,r,] = fit.l$mu.beta-beta.true
  error[3,r,] = fit.h$mu.beta-beta.true
}

par(mfrow=c(3,5))
for (i in 1:15){
  boxplot(error[1,,i],error[2,,i],error[3,,i],names=names,ylim=range(error))
  abline(h=0,col=2)
}

mse = matrix(0,p,3)
mad = matrix(0,p,3)
for (i in 1:3){
  mse[,i] = apply(error[i,,]^2,2,mean)
  mad[,i] = apply(abs(error[i,,]),2,mean)
}
rmse = matrix(0,p,2)
rmse[,1]= mse[,2]/mse[,1]
rmse[,2]= mse[,3]/mse[,1]
rmad = matrix(0,p,2)
rmad[,1]= mad[,2]/mad[,1]
rmad[,2]= mad[,3]/mad[,1]


par(mfrow=c(2,2))
plot(mse[,1],type="h",ylim=range(mse),ylab="MSE")
lines(1:p+0.25,mse[,2],type="h",col=2)
lines(1:p+0.5,mse[,3],type="h",col=3)
legend("topright",legend=c("ridge","lasso","horseshoe"),col=1:3,lty=1)

plot(mad[,1],type="h",ylim=range(mad),ylab="MAD")
lines(1:p+0.25,mad[,2],type="h",col=2)
lines(1:p+0.5,mad[,3],type="h",col=3)


plot(rmse[,1],ylim=range(rmse),ylab="Relative MSE",pch=16)
points(rmse[,2],col=2,pch=16)
abline(h=1,col=3)
legend("topright",legend=c("lasso/ridge","horseshoe/ridge"),col=1:2,pch=16)

plot(rmad[,1],ylim=range(rmad),ylab="Relative MAD",pch=16)
points(rmad[,2],col=2,pch=16)
abline(h=1,col=3)