Implementation of game of sticks from http://nifty.stanford.edu/2014/laaksonen-vihavainen-game-of-sticks/handout.html

Game of Sticks

#Implementation of 'Game of Sticks' (http://nifty.stanford.edu/2014/laaksonen-vihavainen-game-of-sticks/handout.html)

#Instructions
#In the game of sticks there is a heap of sticks on a board. 
#On their turn, each player picks up 1 to 3 sticks. 
#The one who has to pick the final stick will be the loser.

import random

hats=[0]

while 1:
	game_count=0

	print "Welcome to the game of sticks!"
	no_of_sticks=input("How many sticks are there on the table initially (10-100)? ")

	while no_of_sticks<10 or no_of_sticks>100:
		print "Please enter a number between 10 and 100"
		no_of_sticks=input("How many sticks are there on the table initially (10-100)? ")	

	print "Options:\n Play against a friend (1)\n Play against the computer (2)\n Play against the trained computer (only for first time. Press 2 next time\notherwise another set of simulations will take place.) (3)"
	option=input("Which option do you take (1-3)? ")
	while option<1 or option>3:
		option=input("Which option do you take (1-3)? ")


	#--------------------PLAYER VS PLAYER. Will merge the code common to both afterwards--------------------
	if option==1:
		player=0
		while no_of_sticks>2:
			print "\nThere are",no_of_sticks,"sticks on the board."
			sticks_taken=input("Player "+str(player)+": How many sticks do you take (1-3)? ")
			while sticks_taken<1 or sticks_taken>3:
				print "Please enter a number between 1 and 3"
				sticks_taken=input("Player",player,": How many sticks do you take (1-3)? ")
			no_of_sticks-=sticks_taken
			player+=1
			player%=2

		while no_of_sticks>=1: #To match the grammar of the expected output when no_of_sticks==1
			print "\nThere is/are",no_of_sticks,"stick(s) on the board."
			sticks_taken=input("Player "+str(player)+": How many sticks do you take (1-2)? ")
			while no_of_sticks-sticks_taken<0:
				print"Only",no_of_sticks,"stick(s) left. You have no option but to choose (1-2) stick(s)."
				sticks_taken=input("Player: How many sticks do you take (1-2)? ")
			no_of_sticks-=sticks_taken

		print "Player "+str(player)+", you lose."


	#------------------------------TRAIN AI------------------------------
	elif option==3:
		print "Training AI, please wait..."
		max_n_of_sticks=100 #Max. no. of sticks
						#Named n_of_sticks so as not to confuse with no_of_sticks
		#Initialize hats[]

		if game_count==0: #Only initialize at the start
			hats=[0]
			hats2=[0]
			for i in xrange(max_n_of_sticks):
				hats.append([1,1,1]) #hats is 1-indexed
				hats2.append([1,1,1])

		no_of_trials=10**5
		for i in xrange(no_of_trials): #Play no_of_trials trails to train AI
			
			n_of_sticks=random.randint(10,100) #Initialize n_of_sticks after every game
			print no_of_trials-i,"to go"
			hat_indices=[] #Also initialize hat_indices after every trial
			hat2_indices=[] #Also initialize hat_indices after every trial

			turn=0
			while n_of_sticks>2:
				turn+=1
				turn%=2 #To toggle turns
				#print "\nThere are "+str(n_of_sticks)+" sticks on the board."
				if turn==0: #If other AI's turn
					#Other AI chooses n_of_sticks
					hat_index=n_of_sticks
					sum1=hats2[n_of_sticks][0]
					sum2=sum1+hats2[n_of_sticks][1]
					sum3=sum(hats2[n_of_sticks])
					random_no=random.randint(1,sum3)
					if random_no>sum2: ball_no=3
					elif random_no>sum1: ball_no=2
					else: ball_no=1
					ball_no=random.randint(1,3)
					sticks_taken=ball_no
					#print "Other AI selects",sticks_taken
					hat2_indices.append([hat_index,sticks_taken-1]) #To add/remove from hats[] depending on whether the AI wins/loses
																	#sticks_taken-1 because indices starts from 0

				else: #Our AI's turn
					hat_index=n_of_sticks
					sum1=hats[n_of_sticks][0]
					sum2=sum1+hats[n_of_sticks][1]
					sum3=sum(hats[n_of_sticks])
					random_no=random.randint(1,sum3)
					if random_no>sum2: ball_no=3
					elif random_no>sum1: ball_no=2
					else: ball_no=1
					sticks_taken=ball_no
					#print "Our AI selects",sticks_taken
					hat_indices.append([hat_index,sticks_taken-1]) #To add/remove from hats[] depending on whether the AI wins/loses
				n_of_sticks-=sticks_taken

			while n_of_sticks>=1:
				turn+=1
				turn%=2 #To toggle turns
				#print "\nThere are "+str(n_of_sticks)+" sticks on the board."
				if turn==0: #If other AI's turn
					#Other AI chooses n_of_sticks
					hat_index=n_of_sticks
					sum1=hats2[n_of_sticks][0]
					sum2=sum1+hats2[n_of_sticks][1]
					sum3=sum(hats2[n_of_sticks])
					random_no=random.randint(1,sum3)
					if random_no>sum2: ball_no=3
					elif random_no>sum1: ball_no=2
					else: ball_no=1
					ball_no=random.randint(1,2) #Can only select a max. of 2 balls
					sticks_taken=ball_no
					while sticks_taken>n_of_sticks:
						#No need to calculate sum1, sum2, sum3 again
						random_no=random.randint(1,sum3)
						if random_no>sum2: ball_no=3
						elif random_no>sum1: ball_no=2
						else: ball_no=1
						ball_no=random.randint(1,2)
						sticks_taken=ball_no
					#print "Other AI selects",sticks_taken
					hat2_indices.append([hat_index,sticks_taken-1]) #To add/remove from hats[] depending on whether the AI wins/loses

				else: #Our AI's turn
					hat_index=n_of_sticks
					sum1=hats[n_of_sticks][0]
					sum2=sum1+hats[n_of_sticks][1]
					sum3=sum(hats[n_of_sticks])
					random_no=random.randint(1,sum3)
					if random_no>sum2: ball_no=3
					elif random_no>sum1: ball_no=2
					else: ball_no=1
					sticks_taken=ball_no
					while sticks_taken>n_of_sticks:
						#No need to calculate sum1, sum2, sum3 again
						random_no=random.randint(1,sum3)
						if random_no>sum2: ball_no=3
						elif random_no>sum1: ball_no=2
						else: ball_no=1
						sticks_taken=ball_no
					#print "Our AI selects",sticks_taken
					hat_indices.append([hat_index,sticks_taken-1]) #To add/remove from hats[] depending on whether the AI wins/loses
				n_of_sticks-=sticks_taken					

			if turn==0:
				#print "WIN!"
				#print "Other AI, you lose."
				for hat in hat_indices:
					hats[hat[0]][hat[1]]+=1

				for hat in hat2_indices:
					#print "hat_index",hat[0],"ball_no",hat[1] #Debug statement
					if hats2[hat[0]][hat[1]]!=1:
						hats2[hat[0]][hat[1]]-=1
			else:
				#print "LOSE"
				#print "Our AI loses."
				for hat in hat_indices:
					#print "hat_index",hat[0],"ball_no",hat[1] #Debug statement
					if hats[hat[0]][hat[1]]!=1:
						hats[hat[0]][hat[1]]-=1

				for hat in hat2_indices:
					hats2[hat[0]][hat[1]]+=1
			#print no_of_trials-i,"to go"

			#for hat in hats: #Print hat contents after every simulation
				#print hat

		#------------>print HAT CONTENTS
		for hat in hats: #Uncomment this to get contents of hats[] for tained AI after simulation
			print hat

		option=2 #So the player vs AI game can continue.
	#-----------------End of AI training------------------------


	#--------------------------------PLAYER VS AI----------------------------
	if option==2: #Only 'if' so that the condition is true when option=2 during AI training.

		#If no_of_sticks for current game are more than no of hats, increase no of hats
										# by an amount #so that len(hats)=no_of_sticks+1
		if no_of_sticks>len(hats)-1:	#Increment only if no_of_sticks > len(hats)-1, 
										#otherwise will throw exception as xrange(-ve value)
			for i in xrange(no_of_sticks+1-len(hats)): #+1 to make hats[] 1-indexed.
						# +1 creates one more element when adding more hats when not running for the first time.
				hats.append([1,1,1])
		hat_indices=[]

		turn=1
		while no_of_sticks>2:
			turn+=1
			turn%=2
			print "\nThere are",no_of_sticks,"sticks on the board."
			if turn==0: #If player's turn
				sticks_taken=input("Player 1: How many sticks do you take (1-3)? ")
				while sticks_taken<1 or sticks_taken>3:
					print"Please enter a number between 1 and 3"
					sticks_taken=input("Player 1: How many sticks do you take (1-3)? ")
			else: #AI's turn
				hat_index=no_of_sticks
				sum1=hats[no_of_sticks][0]
				sum2=sum1+hats[no_of_sticks][1]
				sum3=sum(hats[no_of_sticks])
				random_no=random.randint(sum1,sum3)
				if random_no>sum2: ball_no=3
				elif random_no>sum1: ball_no=2
				else: ball_no=1

				sticks_taken=ball_no
				print "AI selects",sticks_taken
				hat_indices.append([hat_index,sticks_taken-1]) #To add/remove from hats[] depending on whether the AI wins/loses
			no_of_sticks-=sticks_taken

		while no_of_sticks>=1:
			turn+=1
			turn%=2
			print "\nThere is",no_of_sticks,"stick(s) on the board."
			if turn==0:
				sticks_taken=input("Player 1: How many sticks do you take (1-2)? ")
				while sticks_taken>no_of_sticks:
					print"Only",no_of_sticks,"stick(s) left. "
					sticks_taken=input("Player: How many sticks do you take? ")

			else:
				hat_index=no_of_sticks
				sum1=hats[no_of_sticks][0]
				sum2=sum1+hats[no_of_sticks][1]
				sum3=sum(hats[no_of_sticks])
				random_no=random.randint(1,sum3)
				if random_no>sum2: ball_no=3
				elif random_no>sum1: ball_no=2
				else: ball_no=1

				sticks_taken=ball_no
				while sticks_taken>no_of_sticks:
					#No need to calculate sum1, sum2, sum3 again
					random_no=random.randint(1,sum3)
					if random_no>sum2: ball_no=3
					elif random_no>sum1: ball_no=2
					else: ball_no=1
					sticks_taken=ball_no
				print "AI selects",sticks_taken
				hat_indices.append([hat_index,sticks_taken-1]) #To add/remove from hats[] depending on whether the AI wins/loses

			no_of_sticks-=sticks_taken

		if turn==0:
			print "Player 1, you lose."
			for hat in hat_indices:
				hats[hat[0]][hat[1]]+=1
		else:
			print "AI loses."
			for hat in hat_indices:
				#print "hat_index",hat[0],"ball_no",hat[1] #Debug statement
				if hats[hat[0]][hat[1]]!=1:
					hats[hat[0]][hat[1]]-=1

	game_count+=1
	op=input("Play again (1 = yes, 0 = no)? ")
	if op==0: break