import random
from collections import defaultdict
from graphviz import Digraph
from IPython.display import Image, display

# Step 1: Generate Random Subset Sum Problem
def generate_random_subset_sum(num_elements, target_sum_range):
    elements = [random.randint(1, 100) for _ in range(num_elements)]
    target_sum = random.randint(1, target_sum_range)
    return elements, target_sum

# Step 2: Build Graph Representation
def build_graph(elements, target_sum):
    graph = defaultdict(list)
    for u in range(target_sum + 1):
        for e in elements:
            if u + e <= target_sum:
                graph[u].append(u + e)
    return graph

# Step 3: Construct Boolean Circuits
def build_boolean_circuit(graph, target_sum):
    dot = Digraph(comment='Boolean Circuit for Subset Sum')

    # Create nodes for each sum from 0 to target_sum
    for i in range(target_sum + 1):
        dot.node(f's{i}', f'Sum {i}')

    # Create edges based on the graph transitions
    gate_counter = 1
    for u in graph:
        for v in graph[u]:
            element = v - u
            gate_label = f'AND{gate_counter}'
            gate_counter += 1
            dot.node(gate_label, 'AND', shape='rectangle')
            dot.edge(f's{u}', gate_label, label=str(element))
            dot.edge(gate_label, f's{v}')

    # Add the final OR gate to check if target_sum is achievable
    dot.node('OR', 'OR', shape='circle')
    dot.edge(f's{target_sum}', 'OR')

    return dot

# Step 4: Analyze Circuit Complexity
def analyze_circuit_complexity(graph):
    gate_count = 0
    for u in graph:
        gate_count += len(graph[u])
    return gate_count

# Main Function
def main():
    # Generate random subset sum problem
    elements, target_sum = generate_random_subset_sum(5, 50)
    print(f"Elements: {elements}")
    print(f"Target Sum: {target_sum}")

    # Build graph representation
    graph = build_graph(elements, target_sum)

    # Construct Boolean circuit
    dot = build_boolean_circuit(graph, target_sum)

    # Save and render the graph
    output_path = 'boolean_circuit'
    dot.render(output_path, format='png', view=False)
    print(f"Circuit saved to {output_path}.png")

    # Analyze circuit complexity
    complexity = analyze_circuit_complexity(graph)
    print(f"Number of gates: {complexity}")

    # Display the image in Jupyter notebook
    return Image(filename=f'{output_path}.png')

# Run the main function
if __name__ == "__main__":
    img = main()
    display(img)