import numpy as np
import matplotlib.pyplot as plt

bridges = ['TTGGG', 'TGGGG', 'GGGGG']
bridge_pvals = [0.2, 0.2, 0.6]

def capture_creature(creatures):
    """Captures and returns a random creature and removes it from the list."""
    index = np.random.randint(len(creatures))
    return creatures.pop(index)

results = []

num_safe = 0
num_trolls_picked = 0
while num_trolls_picked < 100000:
    bridge = np.random.choice(bridges, 1, bridge_pvals)[0]
    creatures = list(bridge)
    captured_creature = capture_creature(creatures)
    if captured_creature == 'T':
        num_trolls_picked += 1
        if 'T' not in creatures:
            num_safe += 1
        results.append(num_safe / float(num_trolls_picked))

plt.plot(results)
plt.ylim(.3, .36)

print results[-1]

# priors[i] is the probability of there being i boxes with 90% valuable items
priors = np.array([.0, .05, .20, .25, .20, .20, .10, 0, 0, 0, 0, 0, 0, 0])
print priors.sum()


# P(Data | H_i)
p_data_given_hyp = []
for hyp in xrange(14):
    p_data_given_hyp.append((9 / 10.0) * (hyp / 13.0) + (2 / 10.0) * (13 - hyp) / 13.0)
print p_data_given_hyp

# P(Data)
p_data = np.dot(priors, p_data_given_hyp)

posteriors = (p_data_given_hyp * priors) / p_data

plt.plot(priors, label='prior')
plt.plot(posteriors, label='posterior')
plt.xlim(0, 13)
plt.legend()

p_valuable = 0.0
for hyp in xrange(14):
    p_valuable += posteriors[hyp] * ((8 / 9.0) * (hyp / 13.0) + (1 / 9.0) * (13 - hyp) / 13.0)
    
print p_valuable