"""
Playing around with Hidden Markov Models in Natural Language Processing

Elias & Jim Oct 25 2011


1. Generate a long (1000 or more) sequence of hidden states.
   From that data alone, write code that finds
   a) the probability of each state,
   b) the conditional probabilities P(i|j), and
   c) compare those with what you calculate should be
      the theoretical values from the markov definitions.
2. Do the same for a long (1000 or more) sequence of obversables.
   How do you calculate the theoretical P(i|j) in this case?
   Do the values estimated from the data agree?
3. What does the Verterbi algorithm do?
   a) what is the input?
   b) what is the output?


"""

import random
import doctest
import string

#List tuples of States and Outputs
states = ('1', '2', '3')
outputs = ('a', 'b', 'c', 'd', 'e')

#Probabilities for first 'visible' node
start_prob = {'1' : 0.4, '2' : 0.5, '3' : 0.1}

#Probabilities for node transitions
transition_prob = {
   '1' : {'1' : 0.2, '2' : 0.5, '3' : 0.3},
   '2' : {'1' : 0.4, '2' : 0, '3': 0.6},
   '3' : {'1' : 0.5, '2' : 0.2, '3' : 0.3}
}

#Output probabilities from each node
emission_prob = {
   '1' : {'a': 0.5, 'b' : 0.1, 'c' : 0.3, 'd' : 0.1, 'e' : 0},
   '2' : {'a': 0.1, 'b' : 0.2, 'c' : 0.3, 'd' : 0.2, 'e' : 0.2},
   '3' : {'a': 0.2, 'b' : 0.3, 'c' : 0.1, 'd' : 0, 'e' : 0.4},
}

def weighted_random_choice(probabilities, x=False):
    """Given a dictionary of probabilities i.e. {'a':0.8, 'b':0.2},
       return one of the keys with probability given by the values.
       >>> weighted_random_choice({'a':0.5, 'b':0.1, 'c':0.4}, x=0.1)
       'a'
       >>> weighted_random_choice({'a':0.5, 'b':0.1, 'c':0.4}, x=0.55)
       'b'
       >>> weighted_random_choice({'a':0.5, 'b':0.1, 'c':0.4}, x=0.8)
       'c'
    """
    threshold = 0.0
    if not x:
        x = random.random()  # 0 <= x < 1.0
    choices = probabilities.keys()
    choices.sort()
    for choice in choices:
        threshold += probabilities[choice]
        #print "choice, prob, threshold, x = %s, %f, %f, %f" % \
        #    (choice, probabilities[choice], threshold, x)
        if x < threshold:
            return choice
    return None

def make_markov_states(how_many=1000):
    """Return one sequence of hidden markov states"""
    states = [weighted_random_choice(start_prob)]
    for i in range(how_many-1):
        prob_next_state = transition_prob[states[-1]]
        states.append(weighted_random_choice(prob_next_state))
    return states

def make_observables(hidden_states):
    """Return list of observables given hidden states"""
    return [weighted_random_choice(emission_prob[s]) for s in hidden_states]

def prob_raw(states):
	"""Return raw probabilities for each hidden or observable state.
	>>> prob_raw(['a', 'a', 'b', 'c', 'd', 'd', 'c', 'b', 'c', 'd'])
	Probability of a : 2/10
	Probability of c : 3/10
	Probability of b : 2/10
	Probability of d : 3/10
	>>> prob_raw(['hi', 'there', 'what', 'have', 'you', 'got', 'there'])
	Probability of what : 1/7
	Probability of there : 2/7
	Probability of hi : 1/7
	Probability of have : 1/7
	Probability of got : 1/7
	Probability of you : 1/7
	"""
	counts = {}
	for state in states:
		counts[state] = counts.get(state, 0) + 1
	
	for key in counts:
		print "Probability of %s : %d/%d" % \
			(key, counts[key], len(states))
			
def cond_prob(states):
	"""Return probability of i given j for a set of hidden or visible states"""
	count_pairs = {}
	counts = {}
	for i in range(len(states)-1):
		pair = (states[i], states[i+1])
		count_pairs[pair] = count_pairs.get(pair, 0) + 1
    
	for state in states:
		counts[state] = counts.get(state, 0) + 1
    
	for key in count_pairs:
		print "P(%s | %s) : %d/%d" % \
			(key[1], key[0], count_pairs[key], counts[key[1]])

def main():
	print "not working yet"

if __name__ == "__main__":
    doctest.testmod()
    main()