#!/usr/bin/env python
#
# Copyright years ago by Thomas Strathmann <thomas@pdp7.org>
# released under GNU GPL v2
#
import operator
from random import random
from sys import argv

class MarkovChain:
	"""A simple markov chain class"""
	
	def __init__(self, P, x):
		"""P (transition matrix), x (state vector)"""
		self.P = P
		self.x = x

	def matrixmult(self, m, v):
		"""multiplicate the vector v with the matrix m"""
		return [reduce(operator.add, map(operator.mul, r, v)) for r in m]

	def next(self):
		"""advance to next state vector"""
		self.x = self.matrixmult(self.P, self.x)	# compute new state vector
		return self.state()

	def state(self):
		"""choose one likely state to be in right now based on the
		   current state vector and a random number)"""
		s1 = map(lambda x: x * random(), self.x)
		s2 = s1[:]
		s1.sort()
		s1.reverse()
		return s2.index(s1[0])
		
# Small demo that iterates the number of steps given as first argument
if __name__ == "__main__":
	P = [	[.5, 0, .5, 0, 0],
			[0, .25, .25, .25, .25],
			[.25, .25, .25, .25, 0],
			[.5, 0, .5, 0, 0],
			[.33, 0, 0, .33, .33]]
	x = [1, 0, 0, 0, 0]
	m = MarkovChain(P, x)

	if(len(argv) > 1):
		print m.state()
		for i in range(int(argv[1])):
			print m.next()