#MODEL
# y[i]~N(mu[i],tau1)
# mu[1]~N(0,tau2)
# mu[t]~N(rho*mu[t-1],tau3)
# rho~Unif(min.rho,max.rho)
# tau1,tau2,tau3~Gamma(a,b)

AR1<-function(y,
              min.rho=0,max.rho=1,a=.1,b=.1,
              iters=20000,burn=1000,update=100){

   n<-length(y)

   #INITIAL VALUES:
   mu<-rep(0,n)
   tau1<-tau2<-tau3<-.1
   rho<-(max.rho+min.rho)/2
   fitted<-rep(0,n)

   #Keep track of these samples:
   dev<-rep(0,iters)
   keepers<-matrix(0,iters,4)
   colnames(keepers)<-c("rho","tau1","tau2","tau3")

   for(i in 1:iters){

    #UPDATE THE AR RANDOM EFFECTS
    if(max.rho>0){for(t in 1:n){
      VVV<-tau1
      MMM<-tau1*y[t]
      if(t==1){VVV<-VVV+tau2}
      if(t>1){VVV<-VVV+tau3}
      if(t<n){VVV<-VVV+tau3*rho^2}
      if(t>1){MMM<-MMM+tau3*rho*mu[t-1]}
      if(t<n){MMM<-MMM+tau3*rho*mu[t+1]}
      mu[t]<-rnorm(1,MMM/VVV,1/sqrt(VVV))      
    }}
    if(max.rho==0){mu<-rep(0,n)}

    #UPDATE THE HYPERPARAMETERS:
    tau1<-rgamma(1,n/2+a,sum((y-mu)^2)/2+b)
    tau2<-rgamma(1,1/2+a,sum(mu[1]^2)/2+b)
    tau3<-rgamma(1,(n-1)/2+a,sum((mu[-1]-rho*mu[-n])^2)/2+b)

   if(max.rho>min.rho){
     VVV<-tau3*sum(mu[-n]^2)
     MMM<-tau3*sum(mu[-1]*mu[-n])
     rho<-rtnorm(1,MMM/VVV,1/sqrt(VVV),min.rho,max.rho)
    }


    keepers[i,]<-c(rho,tau1,tau2,tau3)
    dev[i]<--2*sum(dnorm(y,mu,1/sqrt(tau1),log=T))

    if(i>burn){
      fitted<-fitted+mu/(iters-burn)
    }
    if(i%%update==0){
      plot(y,main=i)
      lines(mu)   
    }

   }

   #COMPUTE DIC:
   dbar<-mean(dev[burn:iters])
   shat<-median(1/sqrt(keepers[burn:iters,2]))
   dhat<--2*sum(dnorm(y,fitted,shat,log=T))
   pD<-dbar-dhat
   DIC<-dbar+pD

list(DIC=DIC,pD=pD,dbar=dbar,dev=dev,
     keepers=keepers,fitted=fitted)}

#GENERATE A TRUNCATED NORMAL:
rtnorm<-function(n,mu,sigma,L,U){
   l.prob<-pnorm(L,mu,sigma)
   u.prob<-pnorm(U,mu,sigma)
qnorm(runif(n,l.prob,u.prob),mu,sigma)}



#####################################################
####    FIT THE MODEL TO THE MOTORCYCLE DATA     ####:
#####################################################

library(MASS)
y<-mcycle$accel

fit1<-AR1(y,min.rho=0,max.rho=0)
fit2<-AR1(y,min.rho=1,max.rho=1)
fit3<-AR1(y,min.rho=0,max.rho=1)

print(c(fit1$dbar,fit1$pD,fit1$DIC))
print(c(fit2$dbar,fit2$pD,fit2$DIC))
print(c(fit3$dbar,fit3$pD,fit3$DIC))

plot(y)
lines(fit2$fitted,lty=1)
lines(fit3$fitted,lty=2)
legend("topleft",c("Random Walk","AR(1)"),
   lty=c(1,2),cex=0.8,inset=0.05)