import gradio as gr
import matplotlib.pyplot as plt
from matplotlib_venn import venn3
from io import BytesIO
from PIL import Image

def calculate_probabilities(A, B, C, AB, AC, BC, ABC):
    total = A + B + C - AB - AC - BC + ABC
    if total == 0:
        return {
            "P(A)": 0,
            "P(B)": 0,
            "P(C)": 0,
            "P(A ∩ B)": 0,
            "P(A ∩ C)": 0,
            "P(B ∩ C)": 0,
            "P(A ∩ B ∩ C)": 0,
        }

    P_A = A / total if total > 0 else 0
    P_B = B / total if total > 0 else 0
    P_C = C / total if total > 0 else 0
    P_AB = AB / total if total > 0 else 0
    P_AC = AC / total if total > 0 else 0
    P_BC = BC / total if total > 0 else 0
    P_ABC = ABC / total if total > 0 else 0
    
    return {
        "P(A)": P_A,
        "P(B)": P_B,
        "P(C)": P_C,
        "P(A ∩ B)": P_AB,
        "P(A ∩ C)": P_AC,
        "P(B ∩ C)": P_BC,
        "P(A ∩ B ∩ C)": P_ABC,
    }

def plot_venn(A, B, C, AB, AC, BC, ABC):
    plt.figure(figsize=(8, 8))
    venn = venn3(subsets=(A, B, AB, C, AC, BC, ABC), set_labels=('A', 'B', 'C'))
    plt.title("Venn Diagram")
    
    buf = BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    img = Image.open(buf)
    return img

def main(A, B, C, AB, AC, BC, ABC):
    venn_diagram = plot_venn(A, B, C, AB, AC, BC, ABC)
    probabilities = calculate_probabilities(A, B, C, AB, AC, BC, ABC)
    return venn_diagram, probabilities

iface = gr.Interface(
    fn=main,
    inputs=[
        gr.Number(label="Cantidad en A"),
        gr.Number(label="Cantidad en B"),
        gr.Number(label="Cantidad en C"),
        gr.Number(label="Cantidad en A ∩ B"),
        gr.Number(label="Cantidad en A ∩ C"),
        gr.Number(label="Cantidad en B ∩ C"),
        gr.Number(label="Cantidad en A ∩ B ∩ C")
    ],
    outputs=[
        gr.Image(type="pil", label="Diagrama de Venn"),
        gr.JSON(label="Probabilidades Calculadas")
    ]
)

iface.launch()