utils.py

import time, math
import numpy as np
import matplotlib.pyplot as plt
import inspect, ast
from collections import deque, namedtuple, defaultdict

red   = lambda text : f'\033[91m{text}\033[0m'
green = lambda text : f'\033[92m{text}\033[0m'

def test(actual, expected):
    try: # hacky way to get the actual expression:
        call_line = inspect.getframeinfo(inspect.currentframe().f_back).code_context[0].strip()
        actual_expr = ast.unparse(ast.parse(call_line).body[0].value.args[0])
    except Exception as e:
        actual_expr = ''

    if actual == expected:
        print(green('PASS:'), f'{actual_expr} -> {actual} == {expected} (expected)')
    else:
        print(red('FAIL:'), f'{actual_expr} -> {actual} != {expected} (expected)')



def plot_time_complexity(func, input_gen, input_sizes=None, input_size_descr='(n)', repeat=25):
    sizes = input_sizes if input_sizes else np.linspace(100, 1000, 10, dtype=int).tolist()

    # Measure
    print(f'Measuring time complexity of {func.__name__} ..')
    times = []
    for n in sizes:
        run_times = []
        for _ in range(repeat):
            args = input_gen(n)
            start = time.perf_counter_ns()
            func(*args) if isinstance(args, tuple) else func(args)  # to support multiple arguments
            run_times.append(time.perf_counter_ns() - start)
        avg_time = min(run_times) # min is more robust than mean for noisy runs
        times.append(avg_time)

    # Baselines (normalized at the first point)
    n = np.linspace(sizes[0], sizes[-1], num=sizes[-1]-sizes[0]+1)
    normalize = lambda curve: curve / curve[0] * times[0]
    baselines = {
        'O(2ⁿ)':      normalize(2**n),
        'O(n³)':      normalize(n ** 3),
        'O(n²)':      normalize(n ** 2),
        'O(n log n)': normalize(n * np.log(n)),
        'O(n)':       normalize(n),
        'O(log n)':   normalize(np.log(n)),
        'O(1)':       normalize(np.ones_like(n)),
    }

    # Plot
    plt.figure(figsize=(7, 5))
    plt.plot(sizes, times, 'o-', label=func.__name__)
    for label, curve in baselines.items():
        plt.plot(n, curve, '--', label=label, alpha=0.5)
    plt.xlabel(f"Input size {input_size_descr}")
    plt.ylabel("Time (s)")
    plt.title(f"Time Complexity of {func.__name__} (normalized at n={sizes[0]})")
    plt.ylim(0, max(times[-1], baselines['O(n)'][-1]))
    plt.yticks([])
    plt.legend(loc='upper left')
    plt.tight_layout()
    plt.show()



def print_binary_tree_array(arr, max_precision=2, empty_node='•', title=''): # expect all gaps in the array to be preserved
    if title:
        print(title)

    if not arr: return
    height = math.ceil(math.log2(len(arr) + 1))
    total_nodes = 2 ** height

    # Format values and compute max width
    value_width = 2
    nodes_formatted = []
    for i in range(total_nodes):
        v = arr[i] if i < len(arr) else None  # we will loop also through missing nodes
        v = f" {v:.{max_precision}f} " if isinstance(v, float) else f'{v} ' if v is not None else empty_node
        nodes_formatted.append(v)
        value_width = max(value_width, len(v))
    if value_width % 2 == 1:  # even number centers better
        value_width += 1

    for level in range(height):
        # Print node values
        n = 2 ** level  # total nodes on the level
        arr = nodes_formatted[n-1: 2*n-1]
        slot_width = 2 ** (height - level - 1) * value_width
        print("".join(v.center(slot_width) for v in arr))

        # Print the edges
        if level < height - 1:
            arm = "─" * (slot_width // 4 - 1)
            connector = f"┌{arm}┴{arm}┐".center(slot_width)
            edge_line = [connector for v in arr]
            print("".join(edge_line))


def print_binary_tree(root, max_precision=2, empty_node="•", title=''):
    arr = nodes_to_array_with_gaps(root)
    return print_binary_tree_array(arr, max_precision, empty_node, title)


def nodes_to_array_with_gaps(root):
    """Level-order to array, preserving gaps by adding elements with None."""
    if not root: return []
    arr = []
    q = deque([root])
    while q and any(node is not None for node in q): # preserve the gaps
        node = q.popleft()
        arr.append(node.value if node else None)
        q.append(node.left    if node else None)
        q.append(node.right   if node else None)
    return arr


def print_trie(root, compact=True):
    print('•') # root node
    if not root.children: return

    stack = [(root, '', '', '', True)] # (node, prefix, indent, char, is_last)
    while stack:
        node, prefix, indent, char, is_last = stack.pop()

        # Collapse chain of single children
        if compact:
            while len(node.children) == 1 and not node.is_end:
                (next_char, next_node), = node.children.items()
                node = next_node
                char += next_char

        # Format and print
        connector = "└─ " if is_last else "├─ "
        ending = f' ({prefix}{char})' if node.is_end else ''
        extra = f' cache={node.cache}' if node.cache is not None else ''
        print(f'{indent}{connector}{char}{ending}{extra}')

        # Recurse into children
        new_prefix = prefix + char
        new_indent = indent + ("   " if is_last else "│  ")
        for i, (char, child) in enumerate(reversed(node.children.items())):
            is_last = i == 0 # note the order of children is reversed
            stack.append((child, new_prefix, new_indent, char, is_last))


def print_union_find_trees(parents):
    children = defaultdict(list)
    roots = []
    for node, parent in parents.items():
        if parent == node:
            roots.append(node)
        else:
            children[parent].append(node)

    def collapse_chain(node):
        chain = [node]
        while len(children[node]) == 1:
            node = children[node][0]
            chain.append(node)
        return node, chain

    def walk_tree(root, prefix=""):
        if not prefix:
            print(f'[{root}]')
            prefix = ' '

        for i, child in enumerate(children[root]):
            is_last = i == len(children[root]) - 1
            branch = "└─ " if is_last else "├─ "
            continuation = " " if is_last else "│ "
            tail, chain = collapse_chain(child)
            print(prefix + branch + " <- ".join(map(str, chain)))
            if children[tail]:
                walk_tree(tail, prefix + continuation)

    for root in roots:
        walk_tree(root)

def draw_directed_graph(adj_map:list|dict, node_labels=None, sort=True):
    if isinstance(adj_map, list):
        adj_map = {i: edges for i, edges in enumerate(adj_map)}  # [{edges}] -> {node: {edges}}

    all_nodes = set(adj_map.keys()) | set(v for edges in adj_map.values() for v in edges)
    if sort:
        all_nodes = sorted(all_nodes)
    n = len(all_nodes)

    # Position vertices on a unit circle to be evenly spaced
    offset = 3 * np.pi / 4  # place first vertex on top left
    angles = np.linspace(0, 2 * np.pi, n, endpoint=False) + offset
    x = np.cos(angles)
    y = np.sin(angles)

    # Prepare a convenient data structure
    Node = namedtuple('Node', ['value', 'edges', 'label', 'x', 'y', 'angle'])
    nodes = {}
    for i, node in enumerate(all_nodes):
        label = node_labels[i] if node_labels else str(node)
        edges = adj_map[node] if node in adj_map else {}
        nodes[node] = Node(x=x[i], y=y[i], angle=angles[i], value=node, edges=edges, label=label)


    # Draw edges with weights
    fig, ax = plt.subplots(figsize=(8, 8))
    for u in nodes.values():
        for v_value in u.edges:
            v = nodes[v_value]
            is_self_connection = u is v
            is_bidirectional = (u.value in v.edges)

            # Draw arrow edges
            if not is_self_connection:
                ax.annotate(
                    text="", xy=(v.x, v.y), xytext=(u.x, u.y),
                    arrowprops=dict(
                        arrowstyle="->", color="black", shrinkA=15, shrinkB=15, lw=1,
                        connectionstyle=f"arc3,rad={-0.1 if is_bidirectional else 0}"
                    )
                )

            # Calculate the position for the weight label
            has_weights = isinstance(u.edges, dict) and u.edges[v_value] is not None
            if has_weights:
                weight = u.edges[v_value]
                label = f"{weight}"
                if is_self_connection:
                    mid_x = u.x * 1.1
                    mid_y = u.y * 1.1
                    label += " ⟳"
                elif is_bidirectional:
                    dx = v.x - u.x
                    dy = v.y - u.y
                    norm = np.sqrt(dx ** 2 + dy ** 2)
                    offset_x = -dy / norm * .1
                    offset_y = dx / norm * .1
                    mid_x = (u.x + v.x) / 2 + offset_x
                    mid_y = (u.y + v.y) / 2 + offset_y
                else:
                    mid_x = (u.x + v.x) / 2
                    mid_y = (u.y + v.y) / 2

                plt.text(
                    mid_x, mid_y, label, ha="center", va="center", fontsize=10,
                    bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="black", alpha=0.9), zorder=4
                )


    # Draw nodes
    ax.scatter(x, y, s=700, c="skyblue", edgecolors="black", zorder=3)
    for u in nodes.values():
        ax.text(u.x , u.y, u.label, ha='center', va='center', fontsize=13, zorder=5)

    ax.set_aspect('equal')
    plt.axis("off")
    plt.show()


def draw_lattice(grid, colors=('white', 'black')):
    rows = len(grid)
    cols = len(grid[0]) if rows > 0 else 0
    assert len(colors) >= len(set(v for row in grid for v in row)), 'Not enough colors for unique values'

    # Plot points with black/white colors
    for i in range(rows):
        for j in range(cols):
            v = grid[i][j]
            x, y = j, rows - 1 - i # flip the y axis so row 0 is at the top
            plt.scatter(x, y, c=colors[v], s=200, edgecolor="black", zorder=2)

    # Draw lattice edges
    plt.hlines(y=np.arange(rows), xmin=0, xmax=cols-1, colors="k", lw=1, zorder=1)
    plt.vlines(x=np.arange(cols), ymin=0, ymax=rows-1, colors="k", lw=1, zorder=1)

    plt.xticks(range(cols))
    plt.yticks(range(rows))
    plt.grid(False)
    plt.axis("off")
    plt.tight_layout()
    plt.show()



if __name__ == '__main__':
    test(sum([1, 2, 3]), 6)
    test(sum([1, 2, 3]), 5)
    plot_time_complexity(sorted, lambda n: list(reversed(range(n))))
    print_binary_tree_array(list(range(2 ** 5 - 4)))

    adj_list = [{1},{2},{3},{4},{5},{2},{},{},{1}]
    draw_directed_graph(adj_list)

    adj_list_weighted = [{1: 1.2, 4: 3.1}, {0: 7.2, 2: 2.5, 4: 0.5}, {1: 3.2, 4: 4.5}, {4: 5.5}, {}]
    draw_directed_graph(adj_list_weighted, node_labels=["A", "B", "C", "D", "F"])

    adj_map = {2: {3}, 3: {4}, 4: {5}, 5: {2}}
    draw_directed_graph(adj_map)

    adj_map_weighted = {'A': {'B': 100}, 'B': {'C': 100}, 'C': {'D': 100}, 'D': {'A': 100}}
    draw_directed_graph(adj_map_weighted)

    adj_map = {2: {1:100, 3:100}, 3: {4}, 5: {}}
    draw_directed_graph(adj_map)