#!/usr/bin/python
##########################################################################
#
# Mean tosses to coin-toss patterns
#
# Scott Hendrickson
# 2008-Apr-03
#
# This work is licensed under a
# Creative Commons Attribution-Noncommercial 3.0 United States License
# See http://creativecommons.org/licenses/by-nc/3.0/us/
#
##########################################################################
import random, re

class coinTossSimulation:
	def __init__(this, pat):
		this.pattern = pat
		this.plen = len(pat)
		this.record = []
		this.iter = 20000
		this.simulated = False

	def simulate(this):
		for i in range(0,this.iter):
			count = 0
			cont = True
			this.tosses = ''
			while(cont):
				if random.random() >= 0.50:
					this.tosses += 'H'
				else:
					this.tosses += 'T'
				count += 1
				if this.pattern == this.tosses[-this.plen:]:
					cont = False
			this.record.append(count)
		this.simulated = True

	def meanTosses(this):
		if not this.simulated:
			this.simulate()
		return sum(this.record)/float(len(this.record))

	def printMeanTosses(this):
		print "E<N> for pattern '"+this.pattern+"': %.2f" % this.meanTosses()

	def setIter(this,n):
		this.iter = n

	def histogram(this):
		hist = {}
		for i in this.record:
			if i in hist:
				hist[i] += 1
			else:
				hist[i] = 1

		xmax = max(hist.keys())
		xmin = min(hist.keys())
		filledHist = []
		for i in range(xmin, xmax+1):
			if i in hist:
				filledHist.append((i,hist[i]))
			else:
				filledHist.append((i,0))
		return filledHist

class coinPermutations:
	def __init__(this, n):
		this.length = n
		this.items = 2**n
		this.generated = False
		this.plist = []

	def getPerms(this):
		if not this.generated:
			this.generate(this.items)
		return this.plist

	def generate(this, n):
		sl = ['' for i in range(0,this.items)]
		for i in range(0,this.length):
			bit = 'H'
			for j in range(0,this.items,2**(i+1)):
				bit = this.flipBit(bit)
				for k in range(0,2**(i+1)):
					sl[j+k] += bit
					bit = this.flipBit(bit)
		this.generated = True
		this.plist = sl

	def flipBit(this,b):
		if b == 'H':
			bit = 'T'
		else:
			bit = 'H'
		return bit

class coinPermutationsWOPattern(coinPermutations):
	"""Returns the number of permutations for n-1 tosses that do not include
	the target pattern"""
	def __init__(this, n, p):
		this.pat = p
		this.l = len(this.pat)
		this.patRE = re.compile(this.pat)
		# Setup base class to generate all permutations without the desired pattern
		# at the end.
		assert(n - this.l > 0)
		coinPermutations.__init__(this, n - this.l)

	def getPerms(this):
		if not this.generated:
			this.filter()
		return this.permsWOPat

	def filter(this):
		this.generate(this.items)
		this.permsWOPat = []
		while len(this.plist) > 0:
			i = this.plist.pop()
			# Permutation (n-l) + pattern (l) ( = n )
			full = i + this.pat[0:-1]
			# Exclude all permutations of ( n-1 ) that contain pattern
			if not this.patRE.search(full):
				this.permsWOPat.append(full)
		return this.permsWOPat

	def numPermWithoutPattern(this):
		if not this.generated:
			this.filter()
		return len(this.permsWOPat)


if __name__ == '__main__':
	######################################
	# Examples
	######################################
	# Coin toss permutation generator
	p = coinPermutations(2)
	print "\nAll permutations of 2 coin tosses:"
	print p.getPerms()
	p = coinPermutations(3)
	print "\nAll permutations of 3 coin tosses:"
	print p.getPerms()

	# Permutations without specified pattern
	p = coinPermutationsWOPattern(3, 'TH')
	print "\nPermutations of 3 coin tosses without pattern "
	print "'TH' in first 2 tosses (perms, num of perms): "
	print p.getPerms(), p.numPermWithoutPattern()
	p = coinPermutationsWOPattern(4, 'HTH')
	print "\nPermutations of 4 coin tosses without pattern "
	print "'HTH' in first 3 tosses (perms, num of perms): "
	print p.getPerms(), p.numPermWithoutPattern()

	# Mean tosses to a particular pattern
	b = coinTossSimulation('HHHHH')
	print "\nExpected number of coin tosses to get pattern 'HHHHH': "
	b.printMeanTosses()

	# Compute <n> for all permutations for patterns
	# of length 1 through 4
	print "\nExpected number of coin tosses for all patterns\nfor pattern lengths = 1, 2, 3, 4:"
	for r in range(1,5):
		print "Patterns of length %d --" % r
		p = coinPermutations(r)
		for g in p.getPerms():
			a = coinTossSimulation(g)
			a.printMeanTosses()
			tmp = a.histogram()
			for x in tmp:
				print x[0],"\t",x[1]

	del p
	del a
	#####################################################
	# Counting permutations without the target pattern
	# First term nonzero at l with probability 1/2^l
	n = 26
	n = 12
	l = 2
	expectedN = l * (2**-l)

	print "\n\nTosses\tcoinPermutationsWOPattern(TT)"
	for r in range(l+1,n):
		p = coinPermutationsWOPattern(r, 'TT')
		# Term = n P(n) = n (N(n-1)/2^-(n-1)) (1/2)
		# Last 1/2 is the probability of pattern match at n.
		# N is the number of permutations at n-1 where the last
		# l-1 tosses are the first l-1 tosses that match pattern
		expectedN += r * p.numPermWithoutPattern() * 2**(-r)
		print "%3d %6d %.3f" % (r,p.numPermWithoutPattern(), expectedN)
		del p


	#####################################################
	# Counting permutations without the target pattern
	# First term nonzero at l with probability 1/2^l
	l = 2
	expectedN = l * (2**-l)

	print "\n\nTosses\tcoinPermutationsWOPattern(TH)"
	for r in range(l+1,n):
		p = coinPermutationsWOPattern(r, 'TH')
		# Term = n P(n) = n (N(n-1)/2^-(n-1)) (1/2)
		# Last 1/2 is the probability of pattern match at n.
		# N is the number of permutations at n-1 where the last
		# l-1 tosses are the first l-1 tosses that match pattern
		expectedN += r * p.numPermWithoutPattern() * 2**(-r)
		print "%3d %6d %.3f" % (r,p.numPermWithoutPattern(), expectedN)
		del p

	#####################################################
	# Counting permutations without the target pattern
	# First term nonzero at l with probability 1/2^l
	l = 3
	expectedN = l * (2**-l)

	print "\n\nTosses\tcoinPermutationsWOPattern(TTH)"
	for r in range(l+1,n):
		p = coinPermutationsWOPattern(r, 'TTH')
		# Term = n P(n) = n (N(n-1)/2^-(n-1)) (1/2)
		# Last 1/2 is the probability of pattern match at n.
		# N is the number of permutations at n-1 where the last
		# l-1 tosses are the first l-1 tosses that match pattern
		expectedN += r * p.numPermWithoutPattern() * 2**(-r)
		print "%3d %6d %.3f" % (r,p.numPermWithoutPattern(), expectedN)
		del p

	#####################################################
	# Counting permutations without the target pattern
	# First term nonzero at l with probability 1/2^l
	l = 3
	expectedN = l * (2**-l)

	print "\n\nTosses\tcoinPermutationsWOPattern(HTH)"
	for r in range(l+1,n):
		p = coinPermutationsWOPattern(r, 'HTH')
		# Term = n P(n) = n (N(n-1)/2^-(n-1)) (1/2)
		# Last 1/2 is the probability of pattern match at n.
		# N is the number of permutations at n-1 where the last
		# l-1 tosses are the first l-1 tosses that match pattern
		expectedN += r * p.numPermWithoutPattern() * 2**(-r)
		print "%3d %6d %.3f" % (r,p.numPermWithoutPattern(), expectedN)
		del p