import math
from pomegranate import *

"""
Explanation: Classic Monty Hall Bayesian Network
authors:<br>
Jacob Schreiber [<a href="mailto:jmschreiber91@gmail.com">jmschreiber91@gmail.com</a>]<br>
Nicholas Farn [<a href="mailto:nicholasfarn@gmail.com">nicholasfarn@gmail.com</a>]
Lets test out the Bayesian Network framework to produce the Monty Hall problem, but modified a little. The Monty Hall problem is basically a game show where a guest chooses one of three doors to open, with an unknown one having a prize behind it. Monty then opens another non-chosen door without a prize behind it, and asks the guest if they would like to change their answer. Many people were surprised to find that if the guest changed their answer, there was a 66% chance of success as opposed to a 50% as might be expected if there were two doors.
This can be modelled as a Bayesian network with three nodes-- guest, prize, and Monty, each over the domain of door 'A', 'B', 'C'. Monty is dependent on both guest and prize, in that it can't be either of them. Lets extend this a little bit to say the guest has an untrustworthy friend whose answer he will not go with.
End of explanation
"""


guest = DiscreteDistribution( { 'A': 1./3, 'B': 1./3, 'C': 1./3 } )
prize = DiscreteDistribution( { 'A': 1./3, 'B': 1./3, 'C': 1./3 } )

"""
Explanation: Let's create the distributions for the guest and the prize. Note that both distributions are independent of one another.
End of explanation
"""


monty = ConditionalProbabilityTable(
	[[ 'A', 'A', 'A', 0.0 ],
	 [ 'A', 'A', 'B', 0.5 ],
	 [ 'A', 'A', 'C', 0.5 ],
	 [ 'A', 'B', 'A', 0.0 ],
	 [ 'A', 'B', 'B', 0.0 ],
	 [ 'A', 'B', 'C', 1.0 ],
	 [ 'A', 'C', 'A', 0.0 ],
	 [ 'A', 'C', 'B', 1.0 ],
	 [ 'A', 'C', 'C', 0.0 ],
	 [ 'B', 'A', 'A', 0.0 ],
	 [ 'B', 'A', 'B', 0.0 ],
	 [ 'B', 'A', 'C', 1.0 ],
	 [ 'B', 'B', 'A', 0.5 ],
	 [ 'B', 'B', 'B', 0.0 ],
	 [ 'B', 'B', 'C', 0.5 ],
	 [ 'B', 'C', 'A', 1.0 ],
	 [ 'B', 'C', 'B', 0.0 ],
	 [ 'B', 'C', 'C', 0.0 ],
	 [ 'C', 'A', 'A', 0.0 ],
	 [ 'C', 'A', 'B', 1.0 ],
	 [ 'C', 'A', 'C', 0.0 ],
	 [ 'C', 'B', 'A', 1.0 ],
	 [ 'C', 'B', 'B', 0.0 ],
	 [ 'C', 'B', 'C', 0.0 ],
	 [ 'C', 'C', 'A', 0.5 ],
	 [ 'C', 'C', 'B', 0.5 ],
	 [ 'C', 'C', 'C', 0.0 ]], [guest, prize] )

"""
Explanation: Now let's create the conditional probability table for our Monty. The table is dependent on both the guest and the prize.
End of explanation
"""


s1 = State( guest, name="guest" )
s2 = State( prize, name="prize" )
s3 = State( monty, name="monty" )

"""
Explanation: Now lets create the states for the bayesian network.
End of explanation
"""


network = BayesianNetwork( "test" )
network.add_states( s1, s2, s3 )

"""
Explanation: Then the bayesian network itself, adding the states in after.
End of explanation
"""


network.add_transition( s1, s3 )
network.add_transition( s2, s3 )

"""
Explanation: Then the transitions.
End of explanation
"""


network.bake()

"""
Explanation: With a "bake" to finalize the structure of our network.
End of explanation
"""


print("\t".join([ state.name for state in network.states ]))

"""
Explanation: Now we can check the possible states in our network.
End of explanation
"""


observations = { 'guest' : 'A' }
beliefs = map( str, network.predict_proba( observations ) )
print("\n".join( "{}\t{}".format( state.name, belief ) for state, belief in zip( network.states, beliefs ) ))

"""
Explanation: Now we can see what happens to our network when our Guest chooses 'A'.
End of explanation
"""


observations = { 'guest' : 'A', 'monty' : 'B' }
beliefs = map( str, network.predict_proba( observations ) )
print("\n".join( "{}\t{}".format( state.name, belief ) for state, belief in zip( network.states, beliefs ) ))

"""
Explanation: Now our host chooses 'B'. (note that prize goes to 66% if you switch)
End of explanation
"""


observations = { 'monty' : 'B' }
beliefs = map( str, network.predict_proba( observations ) )
print("\n".join( "{}\t{}".format( state.name, belief ) for state, belief in zip( network.states, beliefs ) ))

"""
Explanation: We can also see what happens if our host simply chooses 'B'.
End of explanation
"""


data = [[ 'A', 'A', 'C' ],
		[ 'A', 'A', 'C' ],
		[ 'A', 'A', 'B' ],
		[ 'A', 'A', 'A' ],
		[ 'A', 'A', 'C' ],
		[ 'B', 'B', 'B' ],
		[ 'B', 'B', 'C' ],
		[ 'C', 'C', 'A' ],
		[ 'C', 'C', 'C' ],
		[ 'C', 'C', 'C' ],
		[ 'C', 'C', 'C' ],
		[ 'C', 'B', 'A' ]]

network.fit( data )

"""
Explanation: Now let's train our network on the following set of data.
End of explanation
"""


print(monty)

"""
Explanation: Let's see the results! Starting with the Monty.
End of explanation
"""


print(prize)

"""
Explanation: Then our Prize.
End of explanation
"""


print(guest)

"""
Explanation: Finally our Guest.
End of explanation
"""


