# (C) May 2009, Mathieu Blondel

import sys
sys.path.append('/home/mathieu/Desktop/sources/jahmm-0.6.1.jar')

from be.ac.ulg.montefiore.run.jahmm import ObservationVector
from be.ac.ulg.montefiore.run.jahmm import OpdfMultiGaussian
from be.ac.ulg.montefiore.run.jahmm import Hmm
from be.ac.ulg.montefiore.run.jahmm.toolbox import MarkovGenerator
from be.ac.ulg.montefiore.run.jahmm.learn import BaumWelchLearner

from java.util import ArrayList

def vectors_to_observations(vectors):
    arr = [ObservationVector(v) for v in vectors]
    return ArrayList(arr)

# 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]
cov_matrix1 = [[15.0, 0.0], [0.0, 30.0]]
opdf1 = OpdfMultiGaussian(means1, cov_matrix1)

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

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

hmm = Hmm(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()
generator = MarkovGenerator(hmm)

from random import randint

for i in range(100):
    # we generate 100 sequences of 20 + k observations where k is variable
    # to show that sequences can be of variable length
    obs_set.add(generator.observationSequence(20 + randint(0, 10)))

# We train our HMM with our fake data
print "HMM, before training"
print str(hmm) + "\n"

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

# Note that it shouldn't change that much since we trained the HMM
# with observations generated from itself
print "HMM, after training"
print str(hmm) + "\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]]
print "Log likelihood for ", obs
print hmm.lnProbability(vectors_to_observations(obs))
 
print "Viterbi sequence"
seq = hmm.mostLikelyStateSequence(vectors_to_observations(obs))
print list(seq)