Venn Diagram

venn.py

from reportlab.pdfgen import canvas
from reportlab.lib.colors import red, blue, black
from reportlab.lib.units import inch

def wrapper(func, output_file="output.pdf"):
    def inner(*args, **kwargs):
        c = canvas.Canvas(output_file)
        func(c, *args, **kwargs)
        c.save()
    return inner

def venn(c, set1: set, set2: set):
    radius = 2 * inch
    diameter = 2 * radius
    margin = 1 * inch
    page_width, page_height = margin + diameter + margin + diameter + margin, margin + diameter + margin

    c.setPageSize((page_width, page_height))

    def draw_circle(x, y, radius, color):
        c.setFillColor(color)
        c.setStrokeColor(black)
        c.setLineWidth(10)
        c.setFillAlpha(0.5)  # Set fill transparency to 50%
        c.circle(x, y, radius, fill=1, stroke=1)

    set1_size = len(set1)
    set2_size = len(set2)

    if set1_size == 0 and set2_size == 0:
        return
    
    if set1_size == 0 and set2_size != 0:
        draw_circle(margin + diameter + margin + radius, margin + radius, radius, blue)
        return

    if set1_size != 0 and set2_size == 0:
        draw_circle(margin + radius, margin + radius, radius, red)
        return

    if set1_size != 0 and set2_size != 0:
        intersection_size = len(set1.intersection(set2))
        if intersection_size == 0:
            # Use the larger set's size as the baseline for radius calculation
            if set1_size >= set2_size:
                set1_radius = radius  # Larger set gets full radius
                set2_radius = radius * (set2_size / set1_size)**0.5  # Scale smaller set proportionally
            else:
                set2_radius = radius  # Larger set gets full radius
                set1_radius = radius * (set1_size / set2_size)**0.5  # Scale smaller set proportionally
            
            draw_circle(margin + radius, margin + radius, set1_radius, red)
            draw_circle(margin + diameter + margin + radius, margin + radius, set2_radius, blue)
        else:
            # Calculate circle sizes based on set sizes
            if set1_size >= set2_size:
                set1_radius = radius
                set2_radius = radius * (set2_size / set1_size)**0.5
            else:
                set2_radius = radius
                set1_radius = radius * (set1_size / set2_size)**0.5
            
            # Calculate intersection ratio relative to the smaller set
            min_set_size = min(set1_size, set2_size)
            intersection_ratio = intersection_size / min_set_size
            
            # Calculate distance between centers based on intersection ratio
            min_distance = abs(set1_radius - set2_radius)  # Maximum overlap
            max_distance = set1_radius + set2_radius       # No overlap
            distance = max_distance - intersection_ratio * (max_distance - min_distance)
            
            # Set positions - fix right circle position
            set2_x = margin + diameter + margin + radius  # Fixed position for right circle
            set2_y = margin + radius
            set1_x = set2_x - distance  # Position left circle relative to right one
            set1_y = set2_y
            
            draw_circle(set1_x, set1_y, set1_radius, red)
            draw_circle(set2_x, set2_y, set2_radius, blue)

def test():
    testcases = [
        (set([1, 2, 3]), set([4, 5, 6])),
        (set([]), set([4, 5, 6])),
        (set([1]), set([])),
        (set([1, 2, 3]), set([4, 5, 6, 7, 8, 9])),
        (set([1, 2, 3]), set([4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15])),
        (set([1, 2, 3, 4]), set([4, 5, 6, 7, 8, 9])),
        (set([1, 2, 3, 4]), set([2, 3, 4, 5]))
    ]
    for index, testcase in enumerate(testcases):
        wrapper(venn, f"{index}.pdf")(*testcase)

if __name__ == "__main__":
    test()