# A Ruby implementation of
# the Viterbi algorithm based on the hidden Markov model (HMM)
# 
# An original Python code: a Wikipedia page "Viterbi algorithm" at
# http://en.wikipedia.org/wiki/Viterbi_algorithm
#
# Author: MISHIMA, Hiroyuki
#

require 'pp'

class Viterbi
  attr_accessor :observations
  attr_accessor :states
  attr_accessor :start_pr
  attr_accessor :transition_pr
  attr_accessor :emission_pr

  def forward_viterbi
    v = Array.new
    v << Hash.new
    path = Hash.new

    # Initialize base cases (t==0)
    states.each do |state|
      v[0][state] = start_pr[state] * emission_pr[state][observations[0]]
      path[state] = [state]
    end

    # Run Viterbi for t>0
    (1 ... (observations.length)).each do |t|
      v << Hash.new
      newpath = Hash.new
      
      states.each do |y|
        prob, state =
          states.map{|y0| [ v[t-1][y0] *
                            transition_pr[y0][y] *
                            emission_pr[y][observations[t]],
                            y0 ]}.max_by{|pr| pr[0]}
        v[t][y] = prob
        newpath[y] = path[state] + [y]
      end
      # Don't need to remeber tge old paths
      path = newpath.dup
    end
    prob, state =
      states.map{|y| [v[observations.length - 1][y], y]}.max_by{|pr| pr[0]}
    [v, prob, path[state]]
  end # def forward_viterbi 
end # class Viterbi

v = Viterbi.new
v.states  = [:rainy,  :sunny]
v.observations = [:walk, :shop, :clean]
v.start_pr = {
  :rainy => 0.6,
  :sunny => 0.4,
}
v.transition_pr = {
  :rainy => { :rainy => 0.7, :sunny => 0.3 },
  :sunny => { :rainy => 0.4, :sunny => 0.6 },
}
v.emission_pr = {
  :rainy => { :walk => 0.1, :shop => 0.4, :clean => 0.5 },
  :sunny => { :walk => 0.6, :shop => 0.3, :clean => 0.1 },
}

pp v.forward_viterbi

# Output ---
# [[{:rainy=>0.06, :sunny=>0.24},
#  {:rainy=>0.038400000000000004, :sunny=>0.043199999999999995},
#  {:rainy=>0.01344, :sunny=>0.0025919999999999997}],
# 0.01344,
# [:sunny, :rainy, :rainy]]