import numpy as np
import matplotlib.pyplot as plt

# %% Number of events (experiments; draws; trials):

n = 100000

# Symbols of samples:

# 'C1C2': Selection 1 = car. Selection 2 = car.
# 'C1G2': Selection 1 = car. Selection 2 = goat.
# 'G1C2': Selection 1 = goat. Selection 2 = car.
# 'G1G2': Selection 1 = goat. Selection 2 = goat.

# %% Params of choice():

population = ['C1C2', 'C1G2', 'G1C2', 'G1G2']
sample_size = 1
replace = True
prob = [1/6, 1/6, 2/6, 2/6]

# %% Variables holding results:

num_C1C2 = 0
num_C1G2 = 0
num_G1C2 = 0
num_G1G2 = 0
num_not_classified = 0

# %% For loop realizing the n trials (samplings):

for k in range(0, n):
    
    # Sample number k:
    sample_k = np.random.choice(population, sample_size, 
                            replace, prob)

    # Classification of sample_k:
    if   sample_k[0] == 'C1C2': num_C1C2 += 1
    elif sample_k[0] == 'C1G2': num_C1G2 += 1
    elif sample_k[0] == 'G1C2': num_G1C2 += 1
    elif sample_k[0] == 'G1G2': num_G1G2 += 1
    else: num_not_classified += 1


# %% Distribution of samples::

distr_samples = np.array([num_C1C2,
                          num_C1G2,
                          num_G1C2,
                          num_G1G2])

# Relative occurencies:

frac_distr_samples = distr_samples/n


# %% Distribution of events:

events = ['C2_after_change', 'C2_after_no_change']

distr_events = np.array([num_G1C2, num_C1C2])
frac_distr_events = distr_events/(num_G1C2 + num_C1C2)

ratio_num_G1C2_to_num_C1C2 = num_G1C2/num_C1C2
print('Increase of probability of GETTING CAR due to '
      + 'CHANGE =', f'{ratio_num_G1C2_to_num_C1C2:.1f}')

# %% Plotting:

plt.close('all')
fig1 = plt.figure(num=1, figsize=(30/2.54, 24/2.54))#Inch

plt.subplot(2,1,1)
plt.bar(population, frac_distr_samples, width=0.9, 
        color=['m', 'b', 'y', 'c'])
# plt.grid()
plt.title('Distribution of samples')
plt.xlabel('Samples')
plt.ylabel('Fraction = Probability')

plt.subplot(2,1,2)
plt.bar(events, frac_distr_events, width=0.9,
        color=['g', 'r'])
# plt.grid()
plt.title('Distribution of events')
plt.xlabel('Events')
plt.ylabel('Fraction = Probability')

# plt.savefig('plot_monty_hall.pdf')
plt.show()