# (C) May 2009, Mathieu Blondel

require 'java'
require '/home/mathieu/Desktop/sources/jahmm-0.6.1.jar'

include_class 'be.ac.ulg.montefiore.run.jahmm.ObservationVector'
include_class 'be.ac.ulg.montefiore.run.jahmm.OpdfMultiGaussian'
include_class 'be.ac.ulg.montefiore.run.jahmm.Hmm'
include_class 'be.ac.ulg.montefiore.run.jahmm.toolbox.MarkovGenerator'
include_class 'be.ac.ulg.montefiore.run.jahmm.learn.BaumWelchLearner'

import java.util.ArrayList

class Array
    # Converts a Ruby array of floats to a Java double array.
    # Supports one-dimension and two-dimension arrays.
    def to_double
        if JRUBY_VERSION =~ /^1.0/
            # This works for both one and two-dimensonal arrays.
            self.to_java(:double)
        else
            if self.length > 0 and self[0].is_a? Array
                self.to_java(Java::double[])
            else
                self.to_java(:double)
            end
        end
    end
end

def vectors_to_observations(vectors)
    arr = vectors.map { |v| ObservationVector.new(v.to_double) }
    ArrayList.new(arr)
end

# We have 2 states and each observation has 2 dimensions x and y.
# We consider x and y to be independent.

means1 = [100.0, 150.0].to_double
cov_matrix1 = [[15.0, 0.0], [0.0, 30.0]].to_double
opdf1 = OpdfMultiGaussian.new(means1, cov_matrix1)

means2 = [200.0, 300.0].to_double
cov_matrix2 = [[20.0, 0.0], [0.0, 40.0]].to_double
opdf2 = OpdfMultiGaussian.new(means2, cov_matrix2)

pi = [0.6, 0.4].to_double
a = [[0.7, 0.3], [0.4, 0.6]].to_double
b = ArrayList.new([opdf1, opdf2])

hmm = Hmm.new(pi, a, b)

# We generate fake data with the Markov generator.
# In the real world, we would use real training data.
# Don't forget that HMM are *generative* models ;)
obs_set = ArrayList.new
generator = MarkovGenerator.new(hmm)

100.times do
    # we generate 100 sequences of 20 + k observations where k is variable
    # to show that sequences can be of variable length
    obs_set << generator.observation_sequence(20 + rand(10))
end

# We train our HMM with our fake data
puts "HMM, before training"
puts hmm.to_s + "\n"

learner = BaumWelchLearner.new
hmm = learner.learn(hmm, obs_set)

# Note that it shouldn't change that much since we trained the HMM
# with observations generated from itself
puts "HMM, after training"
puts hmm.to_s + "\n"

# # Now let's find the log likelihood of a sample sequence
obs = [[85.0, 30.0], [60.0, 150.0], [150.0, 245.0]]
puts "Log likelihood for #{obs.join(", ")}"
puts hmm.ln_probability(vectors_to_observations(obs))
 
puts "Viterbi sequence"
seq = hmm.most_likely_state_sequence(vectors_to_observations(obs))
puts seq.to_a.join(", ")