# https://www.youtube.com/watch?v=rd-yH8hzTLs

'''
In this simulation the astrological signs are represented with the numbers
0 - 11. The position in the array where the signs are randomly stored define
each person and a sign is considered corretly assigned when the number of the
astrological sign is the same as the position in the array.

Example:

position:  0  1  2  3  4  5  6  7  8  9  10  11
array:    [3, 6, 4, 2, 1, 5, 7, 0, 8, 9, 11, 10]
                          ^        ^  ^
There are three correct assignments.
'''

from matplotlib import pyplot as plt
from random import random, shuffle

# The number of times to perform the experiment
iterations = 100000

# This is the probability that a given sign will swtich with another to end up
# in the correct place. This simulates the astrologer's intuition when he ask
# the participants questions.
human_factor_probability = 0.28

number_of_signs = 12

# Array for counting number of correct guesses (0 - 12)
correctness = [0] * (number_of_signs + 1)

for i in range(0, iterations):
    # Assign the zodiac signs randomly to the persons
    signs = [i for i in range(number_of_signs)]
    shuffle(signs)

    # Allow the signs to switch back to the correct position according to the
    # psycological probability
    for i, s in enumerate(signs):
        if s != i and random() <= human_factor_probability:
            signs[i], signs[s] = signs[s], signs[i]

    # Count the correct assignments and note it in the correctness array
    correct_guesses = sum(1 for z in signs if z == signs[z])
    correctness[correct_guesses] += 1


plt.plot(correctness, '-')
plt.title('Human factor probability: {}'.format(human_factor_probability))
plt.xlabel('Correct number of guesses')
plt.ylabel('Occurrences out of {} iterations'.format(iterations))
plt.show()

# Print the probablities for getting the different number of correct guesses
for i, count in enumerate(correctness):
    print('Probablity for {} correct guesses: {}'.format(i, count / iterations))

# Average number of correct guesses per round calculated as the centre of mass
# of the correctness array
print('Average correct guesses: {}'.format(
    sum(i * count for i, count in enumerate(correctness)) / sum(correctness)))