# FORWARD FILTERING, BACKWARD SAMPLING
######################################
ffbs = function(y,x,alphas,beta,sig2,pi,xi,pi0){
  n   = length(y)
  sig = sqrt(sig2)
  # Forward filtering
  pit = rep(0,n)
  A   = dnorm(y[1],alphas[1]+beta*x[1],sig)*(xi*(1-pi0)+(1-pi)*pi0)
  B   = dnorm(y[1],alphas[2]+beta*x[1],sig)*((1-xi)*(1-pi0)+pi*pi0)
  pit[1] = B/(A+B)
  for (t in 2:n){
    A = dnorm(y[t],alphas[1]+beta*x[t],sig)*(xi*(1-pit[t-1])+(1-pi)*pit[t-1])
    B = dnorm(y[t],alphas[2]+beta*x[t],sig)*((1-xi)*(1-pit[t-1])+pi*pit[t-1])
    pit[t] = B/(A+B)	
  }
  # Backward sampling
  st    = rep(0,n)
  pibs  = pit[n]
  st[n] = rbinom(1,1,pibs)
  for (t in (n-1):1){
    if (st[t+1]==1){
      pibs = pit[t]*pi/((1-pit[t])*(1-xi)+pit[t]*pi)
    }else{
      pibs = pit[t]*(1-pi)/(pit[t]*(1-pi)+(1-pit[t])*xi)
    }
    st[t] = rbinom(1,1,pibs)
  }
  return(st)
}

pdf(file="markovswitching.pdf",width=12,height=6)
# SIMULATING SOME DATA
######################
set.seed(31415)
n      = 400
pi     = 0.995  # Pr(S(t)=1|S(t-1)=1)
xi     = 0.995  # Pr(S(t)=0|S(t-1)=0)
alphas = c(0.5,2.5)
beta   = 2
sig2   = 1
sig    = sqrt(sig2)
s      = rep(0,n)
s[1]   = 1
for (t in 2:n){
  prob = pi^s[t-1]*(1-xi)^(1-s[t-1])
  s[t] = rbinom(1,1,prob)
}
x = rnorm(n)
y = rnorm(n,alphas[1+s]+beta*x,sig)
true = c(alphas,beta,sig2,pi,xi,s)

par(mfrow=c(1,1))
plot(alphas[1+s],xlab="Time",main="",pch=16,cex=0.25,col=s+1,ylab="alpha")
legend("topright",legend=c("alpha0 (S=0)","alpha1 (S=1)"),col=1:2,lty=1)

par(mfrow=c(1,1))
plot(x,y,col=s+1,pch=16)


# Forward filtering
pi0 = 0.5
pit = rep(0,n)
A = dnorm(y[1],alphas[1]+beta*x[1],sig)*(xi*(1-pi0)+(1-pi)*pi0)
B = dnorm(y[1],alphas[2]+beta*x[1],sig)*((1-xi)*(1-pi0)+pi*pi0)
pit[1] = B/(A+B)
for (t in 2:n){
  A = dnorm(y[t],alphas[1]+beta*x[t],sig)*(xi*(1-pit[t-1])+(1-pi)*pit[t-1])
  B = dnorm(y[t],alphas[2]+beta*x[t],sig)*((1-xi)*(1-pit[t-1])+pi*pit[t-1])
  pit[t] = B/(A+B)	
}

par(mfrow=c(1,1))
ts.plot(pit,xlab="Time",ylab="Probability",main="Prob(S(t)=1|D(t))")

# Backward smoothing
####################
pib    = rep(0,n)
pib[t] = pit[n]
for (t in (n-1):1){
  pib[t] = pit[t]*pi/((1-pit[t])*(1-xi)+pit[t]*pi)*pib[t+1]+
           pit[t]*(1-pi)/(pit[t]*(1-pi)+(1-pit[t])*xi)*(1-pib[t+1])
}

par(mfrow=c(1,1))
ts.plot(pit,xlab="Time",ylab="Probability")
lines(pib,col=2)
lines(s,col=4)
legend("topright",legend=c("Prob(S(t)=1|D(t))","Prob(S(t)=1|D(n))","True"),col=c(1,2,4),lty=1)
title("Forward filtering\nbackward smoothing")

# Backward sampling
###################
st    = rep(0,n)
pibs   = pit[n]
st[n] = rbinom(1,1,pibs)
for (t in (n-1):1){
  if (st[t+1]==1){
    pibs = pit[t]*pi/((1-pit[t])*(1-xi)+pit[t]*pi)
  }else{
    pibs = pit[t]*(1-pi)/(pit[t]*(1-pi)+(1-pit[t])*xi)
  }
  st[t] = rbinom(1,1,pibs)
}

par(mfrow=c(1,1))
ts.plot(pit,xlab="Time",ylab="Probability",main="")
lines(st,col=2)
lines(s,col=4)
legend("topright",legend=c("Prob(S(t)=1|D(t))","S(t) from P(S(t)|D(n))","True"),col=c(1,2,4),lty=1)
title("Forward filtering\n and backward sampling")


# Learning s1,...,sn,alpha0,alpha1,beta,sig2 via MCMC
#####################################################
ma0=0
Va0=10
ma1=0
Va1=10
mb = 0
Vb = 10
nu0 = 3
sig20 = sig2

# initial values
sd = s
alpha1d = alphas[1]
alpha2d = alphas[2]
betad = beta
sig2d = sig2
pid = pi
xid = xi

# MCMC set up
set.seed(2325)
M0 = 0
M = 1000
L = 1
niter = M0+L*M
draws = matrix(0,niter,6+n)

for (iter in 1:niter){
  # sampling alpha0
  y0 = y[sd==0]
  x0 = x[sd==0]
  n0 = sum(sd==0)
  z0 = y0-betad*x0
  var = 1/(1/Va0+n0/sig2d)
  mean = var*(ma0/Va0+sum(z0)/sig2d)
  alpha0d = rnorm(1,mean,sqrt(var))

  # sampling alpha1
  y1 = y[sd==1]
  x1 = x[sd==1]
  n1 = sum(sd==1)
  z1 = y1-betad*x1
  var = 1/(1/Va1+n1/sig2d)
  mean = var*(ma1/Va1+sum(z1)/sig2d)
  alpha1d = rnorm(1,mean,sqrt(var))

  alphasd = c(alpha0d,alpha1d)

  # sampling beta
  z = y-alphasd[1+sd]
  var = 1/(1/Vb+sum(x^2)/sig2d)
  mean = var*(mb/Vb+sum(z*x)/sig2d)
  betad = rnorm(1,mean,sqrt(var))

  # sampling sig2
  par1 = (nu0+n)/2
  par2 = (nu0*sig20+sum((y-alphasd[sd+1]-betad*x)^2))/2
  sig2d = 1/rgamma(1,par1,par2)

  # sampling s1,...,sn via FFBS
  sd = ffbs(y,x,alphasd,betad,sig2d,pid,xid,pi0)

  # sampling pi and xi
  n0 = sum((sd[2:n]==0)&(sd[1:(n-1)]==0)==TRUE)
  n1 = sum((sd[2:n]==1)&(sd[1:(n-1)]==1)==TRUE)
  pid = rbeta(1,n1,sum(sd)-n1)
  xid = rbeta(1,n1,sum(sd==0)-n0)
  draws[iter,] = c(alphasd,betad,sig2d,pid,xid,sd)
}
ms = apply(draws[,7:(n+6)],2,mean)

names = c("alpha0","alpha1","beta","sig2","pi","xi")
par(mfrow=c(3,6))
for (i in 1:6){
  ts.plot(draws[,i],main=names[i],xlab="iteration",ylab="")
  abline(h=true[i],col=2,lwd=3)
}
for (i in 1:6)
  acf(draws[,i],main=names[i])
for (i in 1:6){
  hist(draws[,i],main=names[i],prob=TRUE,xlab="")
  abline(v=true[i],col=2,lwd=3)
}

par(mfrow=c(1,1))
ts.plot(s,xlab="Time",ylab="Probability",main="")
lines(round(ms),col=3)
points(ms,col=2,pch=16,cex=0.5)
abline(h=0.5,lty=3)
legend("topright",legend=c("True S(t)","P(S(t)=1|D(n))","P(S(t)=1|D(n))>0.5"),col=1:3,lty=1)

dev.off()