daglib/tests/test_performance.py

"""Performance test scaffolding for generating large DAGs."""

from __future__ import annotations

import random
import time
from collections.abc import Callable

import pytest

from daglib.dag import DAG


class DAGGenerator:
    """Utility class for procedurally generating test DAGs."""

    @staticmethod
    def chain(n: int) -> DAG[int]:
        """
        Generate a simple chain: 0 -> 1 -> 2 -> ... -> n-1

        Time complexity: O(n)
        Edge count: n-1
        """
        g = DAG[int]()
        for i in range(n - 1):
            g.add_edge(i, i + 1)
        return g

    @staticmethod
    def tree(depth: int, branching_factor: int = 2) -> DAG[int]:
        """
        Generate a balanced tree with given depth and branching factor.

        Args:
            depth: Depth of the tree (root is depth 0)
            branching_factor: Number of children per node

        Time complexity: O(b^d) where b=branching_factor, d=depth
        Node count: (b^(d+1) - 1) / (b - 1)
        """
        g = DAG[int]()
        node_id = 0
        queue = [(node_id, 0)]  # (node_id, current_depth)
        node_id += 1

        while queue:
            parent, curr_depth = queue.pop(0)
            if curr_depth < depth:
                for _ in range(branching_factor):
                    child = node_id
                    g.add_edge(parent, child)
                    queue.append((child, curr_depth + 1))
                    node_id += 1

        return g

    @staticmethod
    def diamond(layers: int, nodes_per_layer: int = 10) -> DAG[int]:
        """
        Generate a diamond-shaped DAG with multiple layers.
        Each node in layer i connects to all nodes in layer i+1.

        Args:
            layers: Number of layers
            nodes_per_layer: Number of nodes in each middle layer

        Edge count: O(layers * nodes_per_layer^2)
        """
        g = DAG[int]()
        node_id = 0

        # Create layers
        layer_nodes = []
        for layer in range(layers):
            if layer == 0 or layer == layers - 1:
                # Source and sink are single nodes
                layer_nodes.append([node_id])
                node_id += 1
            else:
                # Middle layers have multiple nodes
                current_layer = list(range(node_id, node_id + nodes_per_layer))
                layer_nodes.append(current_layer)
                node_id += nodes_per_layer

        # Connect layers
        for i in range(len(layer_nodes) - 1):
            for u in layer_nodes[i]:
                for v in layer_nodes[i + 1]:
                    g.add_edge(u, v)

        return g

    @staticmethod
    def random_dag(n_nodes: int, edge_probability: float = 0.1, seed: int | None = None) -> DAG[int]:
        """
        Generate a random DAG by considering all possible edges i->j where i < j.

        Args:
            n_nodes: Number of nodes
            edge_probability: Probability of creating each possible edge
            seed: Random seed for reproducibility

        Expected edge count: n*(n-1)/2 * edge_probability
        """
        if seed is not None:
            random.seed(seed)

        g = DAG[int]()
        for i in range(n_nodes):
            for j in range(i + 1, n_nodes):
                if random.random() < edge_probability:
                    g.add_edge(i, j)

        return g

    @staticmethod
    def grid_dag(rows: int, cols: int) -> DAG[tuple[int, int]]:
        """
        Generate a grid DAG where edges flow right and down.
        Node (i,j) connects to (i+1,j) and (i,j+1).

        Args:
            rows: Number of rows
            cols: Number of columns

        Node count: rows * cols
        Edge count: ~2 * rows * cols
        """
        g = DAG[tuple[int, int]]()

        for i in range(rows):
            for j in range(cols):
                # Add edge going right
                if j < cols - 1:
                    g.add_edge((i, j), (i, j + 1))
                # Add edge going down
                if i < rows - 1:
                    g.add_edge((i, j), (i + 1, j))

        return g

    @staticmethod
    def layered_random(
        n_layers: int, nodes_per_layer: int, connection_probability: float = 0.3, seed: int | None = None
    ) -> DAG[int]:
        """
        Generate a layered DAG where nodes only connect to the next layer.

        Args:
            n_layers: Number of layers
            nodes_per_layer: Number of nodes in each layer
            connection_probability: Probability of edge between adjacent layers
            seed: Random seed

        Guarantees: No cycles (edges only go forward)
        """
        if seed is not None:
            random.seed(seed)

        g = DAG[int]()
        node_id = 0
        layers = []

        # Create layers
        for _ in range(n_layers):
            layer = list(range(node_id, node_id + nodes_per_layer))
            layers.append(layer)
            node_id += nodes_per_layer

        # Connect adjacent layers
        for i in range(len(layers) - 1):
            for u in layers[i]:
                for v in layers[i + 1]:
                    if random.random() < connection_probability:
                        g.add_edge(u, v)

        return g

    @staticmethod
    def complete_bipartite(n_left: int, n_right: int) -> DAG[int]:
        """
        Generate a complete bipartite DAG.
        All nodes in left partition connect to all nodes in right partition.

        Args:
            n_left: Number of nodes in left partition
            n_right: Number of nodes in right partition

        Edge count: n_left * n_right
        """
        g = DAG[int]()

        for i in range(n_left):
            for j in range(n_left, n_left + n_right):
                g.add_edge(i, j)

        return g


# Performance benchmarking utilities


def benchmark_operation(operation: Callable[[], None], n_runs: int = 3) -> dict[str, float]:
    """
    Benchmark an operation by running it multiple times.

    Returns:
        Dictionary with min, max, avg, and median times in seconds
    """
    times = []
    for _ in range(n_runs):
        start = time.perf_counter()
        operation()
        end = time.perf_counter()
        times.append(end - start)

    times.sort()
    return {
        "min": min(times),
        "max": max(times),
        "avg": sum(times) / len(times),
        "median": times[len(times) // 2],
    }


# Example performance tests


class TestPerformanceChain:
    """Performance tests for chain DAGs."""

    @pytest.mark.parametrize("n", [100, 1000, 10_000])
    def test_chain_creation(self, n: int) -> None:
        """Test creation time for chain DAGs."""
        result = benchmark_operation(lambda: DAGGenerator.chain(n))
        print(f"\nChain({n}) creation: {result['avg']:.4f}s")

    @pytest.mark.parametrize("n", [100, 1000, 5_000])
    def test_chain_traversal(self, n: int) -> None:
        """Test traversal time for chain DAGs."""
        g = DAGGenerator.chain(n)
        result = benchmark_operation(lambda: g.topo_sort())
        print(f"\nChain({n}) topo_sort: {result['avg']:.4f}s")


class TestPerformanceTree:
    """Performance tests for tree DAGs."""

    @pytest.mark.parametrize("depth,branching", [(5, 2), (4, 3), (3, 5)])
    def test_tree_creation(self, depth: int, branching: int) -> None:
        """Test creation time for tree DAGs."""
        result = benchmark_operation(lambda: DAGGenerator.tree(depth, branching))
        g = DAGGenerator.tree(depth, branching)
        print(f"\nTree(depth={depth}, branch={branching}, nodes={len(g._succ)}) creation: {result['avg']:.4f}s")


class TestPerformanceDiamond:
    """Performance tests for diamond DAGs."""

    @pytest.mark.parametrize("layers,nodes", [(5, 10), (5, 50), (10, 20)])
    def test_diamond_creation(self, layers: int, nodes: int) -> None:
        """Test creation time for diamond DAGs."""
        result = benchmark_operation(lambda: DAGGenerator.diamond(layers, nodes))
        g = DAGGenerator.diamond(layers, nodes)
        print(
            f"\nDiamond(layers={layers}, nodes_per_layer={nodes}, total_edges={len(g)}) creation: {result['avg']:.4f}s"
        )


class TestPerformanceRandom:
    """Performance tests for random DAGs."""

    @pytest.mark.parametrize("n,prob", [(100, 0.1), (500, 0.05), (1000, 0.01)])
    def test_random_creation(self, n: int, prob: float) -> None:
        """Test creation time for random DAGs."""
        result = benchmark_operation(lambda: DAGGenerator.random_dag(n, prob, seed=42))
        g = DAGGenerator.random_dag(n, prob, seed=42)
        print(f"\nRandomDAG(n={n}, p={prob}, edges={len(g)}) creation: {result['avg']:.4f}s")


# Usage examples in docstring

"""
Example usage in your own tests:

```python
from tests.test_performance import DAGGenerator, benchmark_operation

# Generate a large chain for testing
g = DAGGenerator.chain(10000)

# Generate a random DAG
g = DAGGenerator.random_dag(1000, edge_probability=0.05, seed=42)

# Generate a layered DAG (useful for realistic pipeline tests)
g = DAGGenerator.layered_random(10, 100, connection_probability=0.2)

# Benchmark a custom operation
def my_operation():
    g = DAGGenerator.tree(5, 3)
    return g.topo_sort()

stats = benchmark_operation(my_operation, n_runs=10)
print(f"Operation took {stats['avg']:.4f}s on average")

# Generate different graph types for stress testing
graphs = {
    "small_chain": DAGGenerator.chain(100),
    "large_chain": DAGGenerator.chain(10000),
    "balanced_tree": DAGGenerator.tree(depth=7, branching_factor=2),
    "wide_tree": DAGGenerator.tree(depth=3, branching_factor=10),
    "dense_diamond": DAGGenerator.diamond(5, 100),
    "sparse_random": DAGGenerator.random_dag(1000, 0.01),
    "dense_random": DAGGenerator.random_dag(100, 0.3),
    "grid_10x10": DAGGenerator.grid_dag(10, 10),
    "grid_100x100": DAGGenerator.grid_dag(100, 100),
}
```

To run performance tests:
```bash
pytest tests/test_performance.py -v -s
```

To run specific test patterns:
```bash
# Run only chain tests
pytest tests/test_performance.py::TestPerformanceChain -v -s

# Run with specific parameters
pytest tests/test_performance.py::TestPerformanceTree::test_tree_creation -v -s
```
"""