#!/usr/bin/env python3

import math
import random
from pylab import bar
from pylab import show

def code_1():
    n = 10
    p = 0.5
    for k in range(0, n+1):
        prob_value = math.comb(n, k) * pow(p, n)
        print(f"P(X={k}) = {prob_value}")
    
def code_2():
    n, m, p = 5000, 5000, 0.2
    # compute number of heads found in n tosses of a fair unbiased
    # coin in m different trials.
    coin_tosses = [sum([random.random() < p for i in range(n)]) for j in range(m)]
    frequencies = [coin_tosses.count(k) for k in range(n+1)]
    
    # trials = [sum([random.randint(0,1) for i in range(n)]) for j in range(t)]
    bar(range(n+1), frequencies, align='center')
    show()

def code_3():
    n, m = 5000, 5000
    # compute number of heads found in n tosses of a fair unbiased
    # coin in m different trials.
    coin_tosses_02 = [sum([random.random() < 0.2 for i in range(n)]) for j in range(m)]
    frequencies_02 = [coin_tosses_02.count(k) for k in range(n+1)]

    coin_tosses_05 = [sum([random.random() < 0.5 for i in range(n)]) for j in range(m)]
    frequencies_05 = [coin_tosses_05.count(k) for k in range(n+1)]

    coin_tosses_07 = [sum([random.random() < 0.7 for i in range(n)]) for j in range(m)]
    frequencies_07 = [coin_tosses_07.count(k) for k in range(n+1)]

    bar(range(n+1), frequencies_02, align='center', color='red')
    bar(range(n+1), frequencies_05, align='center', color='blue')
    bar(range(n+1), frequencies_07, align='center', color='green')
    show()    

# -------------------------------
    
if __name__ == "__main__":
    # code_2()

    code_3()