###########################################################################
# Homework 3
# Time Series
# Phd in Business Economics, Insper
# February 2021
# 
# http://hedibert.org/wp-content/uploads/2021/02/hw3-timeseries-2020.pdf
#
##########################################################################
library(forecast)
library(astsa)


unemp.q = read.csv("LRUN64TTUSQ156S.csv",header=TRUE)[,2]
date.q  = seq(1970,2019+3/4,by=1/4)

par(mfrow=c(1,1))
plot(date.q,unemp.q,xlab="",ylab="unemployment rate",type="h",lwd=2)
title("Quarterly U.S. civilian unemployment rate, seasonally adjusted\n  https://fred.stlouisfed.org")
lines(date.q,unemp.q,col=2,lwd=2)

par(mfrow=c(1,2))
acf(unemp.q,main="",lag.max=20)
pacf(unemp.q,main="",lag.max=20)

auto.arima(unemp.q)

# Series: unemp.q 
# ARIMA(2,0,0) with non-zero mean 
#
# Coefficients:
#          ar1      ar2    mean
#       1.6533  -0.6870  6.0415
# s.e.  0.0515   0.0519  0.5047
#
# sigma^2 estimated as 0.0618:  log likelihood=-6.14
# AIC=20.28   AICc=20.48   BIC=33.47



#####################################################################################
#  Fitting a TAR(2) model
#####################################################################################


# Setting up y and X for the AR(2) part of the model
ts = diff(unemp.q)
n  = length(ts)
y  = ts[3:n]
X  = cbind(1,ts[2:(n-1)],ts[1:(n-2)])
n  = nrow(X)
ymin = sort(y)[10]
ymax = sort(y)[(n-10)]


#  MAXIMUM LIKELIHOOD INFERENCE
###############################
loglike = function(gamma){  
  ind    = X[,2]<gamma
  y1     = y[ind]
  X1     = X[ind,]
  n1     = length(y1)
  b1     = solve(t(X1)%*%X1)%*%t(X1)%*%y1
  sig21  = sum((y1-X1%*%b1)^2)/n1
  y2     = y[ind==FALSE]
  X2     = X[ind==FALSE,]
  n2     = length(y2)
  b2     = solve(t(X2)%*%X2)%*%t(X2)%*%y2
  sig22  = sum((y2-X2%*%b2)^2)/n2
  return(sum(dnorm(y1,X1%*%b1,sqrt(sig21),log=TRUE))+sum(dnorm(y2,X2%*%b2,sqrt(sig22),log=TRUE)))          
}

gammas = sort(y[((y<ymax)&(y>ymin))])
N = length(gammas)
llike  = rep(0,N)
for (i in 1:N)
  llike[i] = loglike(gammas[i])
  
par(mfrow=c(1,1))
plot(gammas,llike,type="h")


gamma.mle = gammas[llike==max(llike)]
ind    = X[,2]<gamma.mle
y1     = y[ind]
X1     = X[ind,]
n1     = length(y1)
b1     = solve(t(X1)%*%X1)%*%t(X1)%*%y1
sig21  = sum((y1-X1%*%b1)^2)/(n1-3)
y2     = y[ind==FALSE]
X2     = X[ind==FALSE,]
n2     = length(y2)
b2     = solve(t(X2)%*%X2)%*%t(X2)%*%y2
sig22  = sum((y2-X2%*%b2)^2)/(n2-3)

mle = c(b1,sig21,b2,sig22,gamma.mle)

#####################################################################################
#
#  BAYESIAN INFERENCE
#
#####################################################################################
# Prior hyperparameters
b0    = rep(0,3)
B0    = diag(3)
iB0   = solve(B0)
iB0b0 = iB0%*%b0

# Initial value for the MCMC
sd.gamma = 0.05
gamma = mle[9]
phi   = mle[1:3]
sig2  = mle[4]
beta  = mle[5:7]
tau2  = mle[8]

# MCMC setup
set.seed(31415)
burnin = 10000
M      = 10000
niter  = burnin+M
draws = matrix(0,niter,10)
runtime = system.time({
for (iter in 1:niter){
  # phi & sig2
  y1     = y[X[,2]<gamma]
  X1     = X[X[,2]<gamma,]
  B1     = solve(iB0+t(X1)%*%X1/sig2)
  b1     = B1%*%(iB0b0+t(X1)%*%y1/sig2)
  phi    = b1 + t(chol(B1))%*%rnorm(3)
  n1     = length(y1)
  par1   = (n1-2)/2
  par2   = sum((y1-X1%*%phi)^2)/2
  sig21  = 1/rgamma(1,par1,par2)
  accept = min(1,(1+sig2)/(1+sig21))
  if (runif(1)<accept){
    sig2 = sig21
  }
  
  # beta & tau2
  y2     = y[X[,2]>=gamma]
  X2     = X[X[,2]>=gamma,]
  B1     = solve(iB0+t(X2)%*%X2/tau2)
  b1     = B1%*%(iB0b0+t(X2)%*%y2/tau2)
  beta   = b1 + t(chol(B1))%*%rnorm(3)
  n1     = length(y2)
  par1   = (n1-2)/2
  par2   = sum((y2-X2%*%beta)^2)/2
  tau21  = 1/rgamma(1,par1,par2)
  accept = min(1,(1+tau2)/(1+tau21))
  if (runif(1)<accept){
    tau2 = tau21
  }
  
  logden = sum(dnorm(y1,X1%*%phi,sqrt(sig2),log=TRUE))+
           sum(dnorm(y2,X2%*%beta,sqrt(tau2),log=TRUE))

  # gamma
  gamma1 = rnorm(1,gamma,sd.gamma)
  if ((gamma1>ymin)&(gamma1<ymax)){
    y1     = y[X[,2]<gamma1]
    X1     = X[X[,2]<gamma1,]
    y2     = y[X[,2]>=gamma1]
    X2     = X[X[,2]>=gamma1,]
    lognum = sum(dnorm(y1,X1%*%phi,sqrt(sig2),log=TRUE))+
             sum(dnorm(y2,X2%*%beta,sqrt(tau2),log=TRUE))
    logaccept = min(0,lognum-logden)
    if (log(runif(1))<logaccept){
      gamma = gamma1
    }
  }
  n1 = sum(X[,2]<gamma)
  # Storing the draws
  draws[iter,] = c(phi,sig2,beta,tau2,gamma,n1)
}})

draws = draws[(burnin+1):niter,]

runtime

names = c("Intercepts","AR1 coeff","AR2 coeff","Residual variance")
par(mfrow=c(2,3))
for (i in 1:4){
  ts.plot(draws[,4+i],ylim=range(draws[,c(i,4+i)]),xlab="Iterations",ylab="")
  lines(draws[,i],col=2)
  abline(h=mle[i],col=3)
  title(names[i])
}
ts.plot(draws[,9],xlab="Iterations",ylab="",main=expression(gamma))
abline(h=mle[9],col=3,lwd=2)
hist(draws[,9],prob=TRUE,main=expression(gamma))

par(mfrow=c(2,2))
plot(density(draws[,1]),xlim=range(draws[,c(1,5)]),main="",xlab="")
lines(density(draws[,5]),col=2)
title(expression(paste(phi[0]," and ",beta[0])))
plot(density(draws[,2]),xlim=range(draws[,c(2,6)]),main="",xlab="")
lines(density(draws[,6]),col=2)
title(expression(paste(phi[1]," and ",beta[1])))
plot(density(draws[,3]),xlim=range(draws[,c(3,7)]),main="",xlab="")
lines(density(draws[,7]),col=2)
title(expression(paste(phi[2]," and ",beta[2])))
plot(density(draws[,4]),xlim=range(draws[,c(4,8)]),main="",xlab="")
lines(density(draws[,8]),col=2)
title(expression(paste(sigma^2," and ",tau^2)))

par(mfrow=c(1,1))
plot(draws[,9],draws[,10],xlab="gamma",ylab="Counts of y[t-1]<gamma")




#install.packages("tsDyn")
#library("tsDyn")
#mod.setar <- setar(y,d=1,m=2,thVar=lag(y,2),thDelay=1)
#mod.setar
#mle