# A simple implementation of the Baum Welch algorithm (written for 
# illustration only)

# INITIAL VALUES: 
# Random starting probabilities for x_1 and transprobs: 
M <- 3
library(LearnBayes)
startprobs <- rdirichlet(1, rep(1, M))
transprobs <- matrix(rdirichlet(M, rep(1, M)), M, M, byrow=TRUE)

# For the data y_i we assume
# y_i | x_i ~ Normal(x_i, 0.3)
std <- 0.5
y <-c(0.96,0.95,2.21,1.11,0.42,1.86,0.73,0.37,0.78,-0.18,2.11,1.97,1.32,1.69,
      2.42,1.62,2.55,2.12,1.45,2.45,3.72,2.48,2.81,2.48,2.62,3.90,3.08,3.82,
      2.91,3.46,1.59,0.96,1.20,1.06,1.35,1.35,0.64,1.85,0.95,0.82,2.63,2.32,
      1.53,2.88,1.83,0.70,2.11,1.93,2.00,1.91,3.18,3.90,2.74,3.43,3.19,2.64,
      3.12,3.89,2.22,2.79,1.24,0.77,1.46,0.40,1.58,0.58,-0.04,1.98,1.00,1.50,
      1.82,1.31,2.36,1.88,2.33,2.00,2.00,1.80,2.20,1.48,2.40,3.68,3.70,3.40,
      3.67,2.94,2.93,2.59,2.46,2.15,1.25,1.61,0.31,0.46,1.71,0.64,0.29,1.24,
      2.13,0.74)
N <- length(y)

Niter <- 500
startpResults <- matrix(startprobs, Niter, M, byrow=T)
transpResults <- array(0, c(Niter, M, M))
transpResults[1,,] <- transprobs

for (it in 2:Niter) {
  
  # Run forward-backward with the parameters startpResults[i,]
  # and transpResults[i,,]
  
  # Compute and store forward probabilities: 
  start <- startpResults[it-1,]*dnorm(y[1], 1:M, std)
  probsForward <- matrix(start/sum(start), N, M, byrow=T)
  for (i in 2:N) {
    res <- dnorm(y[i], 1:M, std)*probsForward[i-1,]%*%transpResults[it-1,,]
    probsForward[i,] <- res/sum(res) # only to improve numerics
  }
  
  # Compute and store backward probabilities: 
  probsBackward <- matrix(1, N, M, byrow=T)
  for (i in (N-1):1) {
    res <- transpResults[it-1,,]%*%probsBackward[i+1,]*dnorm(y[i+1], 1:M, std)
    probsBackward[i,] <- res/sum(res) # only to improve numerics
  }
  
  # Compute the updated parameters: 
  
  start <- probsForward[1,]*probsBackward[1,]
  startpResults[it,] <- start/sum(start)
  
  for (ct in 1:(N-1)) {
    mat <- transpResults[it-1,,]*matrix(probsForward[ct,],M,M)*
      t(matrix(dnorm(y[ct+1],1:M, std), M, M))*
      t(matrix(probsBackward[ct,], M, M))
    mat <- mat/sum(mat)
    transpResults[it,,] <- transpResults[it,,] + mat
  }
  ss <- apply(transpResults[it,,],1, sum)
  transpResults[it,,] <- transpResults[it,,]/ss

  if (it==Niter) {
   # print(round(probsForward[1:20,],3))
  #  print(round(probsBackward[1:20,],3))
    result <- probsForward*probsBackward
    sumresult <- apply(result, 1, sum)
    result <- result/sumresult 
  #  print(round(result[1:20,],3))
    
    ## Expected value for x_i: 
    expected <- apply(t(result)*(1:M), 2, sum)
    standd   <- sqrt(abs(apply(t(result)*(1:M)^2, 2, sum)-expected^2))
    par(mfrow=c(1,1))
    plot(y, type="l")
    lines(expected, col="red")
    lines(expected + standd, col="blue")
    lines(expected - standd, col="blue")

    print(round(transpResults[it,,],3))
  }
}