import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
from matplotlib_venn import venn2, venn3

# Función para calcular el teorema de Bayes con dos eventos
def bayes_two_events(prior_A, prob_B_given_A, prob_B_given_notA):
    prob_B = prob_B_given_A * prior_A + prob_B_given_notA * (1 - prior_A)  # Calculamos P(B)
    posterior_A_given_B = (prob_B_given_A * prior_A) / prob_B  # Calculamos P(A|B)
    return posterior_A_given_B, prob_B

# Función para crear el diagrama de Venn para dos eventos
def plot_venn(prior_A, prob_B_given_A, prob_B_given_notA):
    area_A = prior_A
    area_B = prob_B_given_A * prior_A + prob_B_given_notA * (1 - prior_A)
    area_AB = prob_B_given_A * prior_A

    fig, ax = plt.subplots()
    venn = venn2(subsets=(area_A, area_B - area_AB, area_AB), set_labels=('Evento A', 'Evento B'), ax=ax)
    venn.get_label_by_id('10').set_text(f'P(A) = {prior_A:.2f}')
    venn.get_label_by_id('01').set_text(f'P(B | ¬A) = {prob_B_given_notA:.2f}')
    venn.get_label_by_id('11').set_text(f'P(B | A) = {prob_B_given_A:.2f}')
    ax.set_title('Diagrama de Venn: Relaciones entre A y B')
    return fig

# Función principal para dos eventos
def bayes_app_two_events(prior_A, prob_B_given_A, prob_B_given_notA):
    posterior_A_given_B, prob_B = bayes_two_events(prior_A, prob_B_given_A, prob_B_given_notA)
    fig_venn = plot_venn(prior_A, prob_B_given_A, prob_B_given_notA)
    return fig_venn, f"P(B): {prob_B:.2f}", f"P(A|B): {posterior_A_given_B:.2f}"

# Función para múltiples eventos
def bayes_multiple_events(priors, likelihoods):
    priors = np.array(priors)
    likelihoods = np.array(likelihoods)
    prob_B = np.sum(priors * likelihoods)  # Calculamos P(B)
    posteriors = (priors * likelihoods) / prob_B  # Calculamos las probabilidades posteriores P(Ai|B) para cada evento Ai
    return posteriors, prob_B

# Función para crear diagramas de Venn o gráficos de partición para múltiples eventos
def plot_sets_multiple_events(priors, likelihoods, posteriors):
    num_events = len(priors)
    
    if num_events == 2:
        # Diagrama de Venn para 2 eventos
        fig, ax = plt.subplots()
        venn = venn2(subsets=(priors[0], priors[1], likelihoods[0] * priors[0]), set_labels=('Evento A1', 'Evento A2'), ax=ax)
        venn.get_label_by_id('10').set_text(f'P(A1) = {priors[0]:.2f}')
        venn.get_label_by_id('01').set_text(f'P(A2) = {priors[1]:.2f}')
        venn.get_label_by_id('11').set_text(f'P(B|A1) = {likelihoods[0]:.2f}')
        ax.set_title('Diagrama de Venn para 2 Eventos')
        return fig

    elif num_events == 3:
        # Diagrama de Venn para 3 eventos
        fig, ax = plt.subplots()
        venn = venn3(subsets=(priors[0], priors[1], likelihoods[0] * priors[0], priors[2], likelihoods[1] * priors[1], likelihoods[2] * priors[2], likelihoods[0] * likelihoods[1] * priors[0]), set_labels=('Evento A1', 'Evento A2', 'Evento A3'), ax=ax)
        venn.get_label_by_id('100').set_text(f'P(A1) = {priors[0]:.2f}')
        venn.get_label_by_id('010').set_text(f'P(A2) = {priors[1]:.2f}')
        venn.get_label_by_id('001').set_text(f'P(A3) = {priors[2]:.2f}')
        venn.get_label_by_id('110').set_text(f'P(B|A1 ∩ A2) = {likelihoods[0] * likelihoods[1]:.2f}')
        venn.get_label_by_id('011').set_text(f'P(B|A2 ∩ A3) = {likelihoods[1] * likelihoods[2]:.2f}')
        venn.get_label_by_id('101').set_text(f'P(B|A1 ∩ A3) = {likelihoods[0] * likelihoods[2]:.2f}')
        venn.get_label_by_id('111').set_text(f'P(B|A1 ∩ A2 ∩ A3) = {likelihoods[0] * likelihoods[1] * likelihoods[2]:.2f}')
        ax.set_title('Diagrama de Venn para 3 Eventos')
        return fig

    else:
        # Gráfico de partición para más de 3 eventos
        events = [f'A{i+1}' for i in range(num_events)]
        fig, ax = plt.subplots()
        ax.bar(events, posteriors, color='purple')
        for i, posterior in enumerate(posteriors):
            ax.text(i, posterior / 2, f'P(A{i+1}|B) = {posterior:.2f}', ha='center', va='center', color='white')
        ax.set_ylabel('Probabilidad Posterior')
        ax.set_title('Resultados del Teorema de Bayes para Múltiples Eventos')
        return fig

# Función principal para múltiples eventos
def bayes_app_multiple_events(priors, likelihoods):
    priors = list(map(float, priors.split(',')))
    likelihoods = list(map(float, likelihoods.split(',')))
    
    posteriors, prob_B = bayes_multiple_events(priors, likelihoods)
    fig = plot_sets_multiple_events(priors, likelihoods, posteriors)
    
    resultados = "\n".join([f"P(A{i+1}|B): {posterior:.2f}" for i, posterior in enumerate(posteriors)])
    return fig, f"P(B): {prob_B:.2f}", resultados

# Crear la interfaz de Gradio con dos pestañas
with gr.Blocks() as demo:
    with gr.Tab("Dos Eventos"):
        gr.Markdown("## Calculadora del Teorema de Bayes con Diagrama de Venn")
        gr.Markdown("Ajusta los valores de probabilidad para calcular la probabilidad de A dado B (P(A|B)).")

        prior_A = gr.Slider(0, 1, step=0.01, label="Probabilidad Inicial de A (P(A))", value=0.5)
        prob_B_given_A = gr.Slider(0, 1, step=0.01, label="Probabilidad de B dado A (P(B|A))", value=0.7)
        prob_B_given_notA = gr.Slider(0, 1, step=0.01, label="Probabilidad de B dado que no ocurre A (P(B|¬A))", value=0.2)

        venn_plot = gr.Plot()
        prob_B = gr.Textbox(label="Probabilidad de B (P(B))")
        prob_A_given_B = gr.Textbox(label="Probabilidad de A dado B (P(A|B))")

        btn = gr.Button("Calcular")
        btn.click(fn=bayes_app_two_events, inputs=[prior_A, prob_B_given_A, prob_B_given_notA], outputs=[venn_plot, prob_B, prob_A_given_B])


    with gr.Tab("Múltiples Eventos"):
        gr.Markdown("## Calculadora del Teorema de Bayes para Múltiples Eventos")
        gr.Markdown("Introduce las probabilidades de los eventos iniciales (Priors) y las probabilidades condicionales de B dado cada evento (Likelihoods) separadas por comas. Ejemplo: '0.2, 0.3, 0.5' para tres eventos.")
        
        priors_input = gr.Textbox(label="Priors (Probabilidades Iniciales)", placeholder="Ejemplo: 0.2, 0.3, 0.5")
        likelihoods_input = gr.Textbox(label="Likelihoods (Probabilidades Condicionales de B dado cada evento)", placeholder="Ejemplo: 0.6, 0.7, 0.4")
        
        plot_multi = gr.Plot()
        prob_B_multi = gr.Textbox(label="Probabilidad de B (P(B))")
        posteriors = gr.Textbox(label="Probabilidades Posteriores")

        btn_multi = gr.Button("Calcular")
        btn_multi.click(fn=bayes_app_multiple_events, inputs=[priors_input, likelihoods_input], outputs=[plot_multi, prob_B_multi, posteriors])

demo.launch()