import gradio as gr
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib_venn import venn2, venn3

def make_venn(file):
    try:
        # Load CSV
        df = pd.read_csv(file.name)
        
        # Drop primary key (first column)
        categories = df.columns[1:]
        if len(categories) < 2:
            return "Error: Need at least 2 category columns!"
        if len(categories) > 3:
            return "Error: Only up to 3 categories supported!"
        
        # Convert each category column into a set of IDs
        sets = [set(df[df[col] == 1].index) for col in categories]
        
        # Plot Venn
        plt.figure(figsize=(6,6))
        if len(sets) == 2:
            venn2(sets, set_labels=categories)
        elif len(sets) == 3:
            venn3(sets, set_labels=categories)
        
        out_file = "venn.png"
        plt.savefig(out_file)
        plt.close()
        return out_file
    
    except Exception as e:
        return f"Error: {e}"

demo = gr.Interface(
    fn=make_venn,
    inputs=gr.File(file_types=[".csv"], label="Upload CSV"),
    outputs=gr.Image(type="filepath", label="Venn Diagram"),
    title="Venn Diagram Generator",
    description="Upload a CSV with a primary key column + up to 3 category columns (binary flags: 0/1)."
)

if __name__ == "__main__":
    demo.launch()