import matplotlib.pyplot as plt
import numpy as np
from random import choices
from typing import List, Tuple

def plot(seq: List[Tuple[int, float]], proba: float, title: str):
    x = np.array([x[0] for x in seq])
    y = np.array([x[1] for x in seq])

    plt.title(title)
    plt.axhline(y=proba, color='r', linestyle='--')
    plt.plot(x, y)
    plt.xlabel("n")
    plt.ylabel("Frecuencia relativa")
    plt.yticks([0, proba, 1.0])
    plt.show()

def experimento_1(n: int):
    """
        Probabilidad total.

    :param n: Veces que se repetirá el experimento.

    :return: Frecuencia relativa del evento A: "Sale una pelota blanca".
    """
    C_1: List[str] = ["B" for _ in range(3)] + ["N" for _ in range(7)]
    C_2: List[str] = ["B" for _ in range(6)] + ["N" for _ in range(6)]
    succ: int = 0
    for i in range(1, n + 1):
        caja: str = choices(["C_1", "C_2"], weights=[0.5, 0.5], k=1)[0]
        if caja == "C_1":
            pelota: str = choices(C_1, k=1)[0]
            if pelota == "B":
                succ += 1
        else:
            pelota: str = choices(C_2, k=1)[0]
            if pelota == "B":
                succ += 1

    return succ/n

def experimento_2(n: int):
    """
        Teorema de Bayes.

    :param n: Veces que se repetirá el experimento.
    :return: Frecuencia relativa del evento 
					   B_1 | A: "La pelota blanca viene de la caja uno".
    """
    C_1: List[str] = ["B" for _ in range(3)] + ["N" for _ in range(7)]
    C_2: List[str] = ["B" for _ in range(6)] + ["N" for _ in range(6)]
    succ_caja_1_blanca: int = 0
    succ_blanca: int = 0
    for i in range(1, n + 1):
        caja: str = choices(["C_1", "C_2"], weights=[0.5, 0.5], k=1)[0]
        if caja == "C_1":
            pelota: str = choices(C_1, k=1)[0]
            if pelota == "B":
                succ_blanca += 1
                succ_caja_1_blanca += 1
        else:
            pelota: str = choices(C_2, k=1)[0]
            if pelota == "B":
                succ_blanca += 1
    try:
        freq = succ_caja_1_blanca / succ_blanca
    except ZeroDivisionError:
        return 0
    return freq

if __name__ == '__main__':
    to_plot = []
    for n in range(1, 10_000, 100):
        to_plot.append((n, experimento_1(n)))
    plot(to_plot, proba=2/5, title="Proba total")

    to_plot = []
    for n in range(1, 10_000, 100):
        to_plot.append((n, experimento_2(n)))
    plot(to_plot, proba=3/8, title="Bayes")