# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_synthetic_data.ipynb.

# %% auto 0
__all__ = ['obfuscate_nodes', 'create_random_graph', 'create_and_save_random_graph', 'generate_and_save_graphs']

# %% ../nbs/01_synthetic_data.ipynb 5
import networkx as nx
import csv
from pathlib import Path
import random
import string
from tqdm import tqdm
from pympler.asizeof import asizeof

# %% ../nbs/01_synthetic_data.ipynb 9
def obfuscate_nodes(n_nodes:int,
                    node_labels_pool:list[str]=None,
                    str_len:int=10) -> dict:
    """
    Maps nodes from range [0, n_nodes) to unique random strings of length str_len.
    All upper case letters and digits are used, starts with a letter.
    If node_labels_pool is provided, nodes selected from it if len(node_labels_pool) >= n_nodes,
    otherwise, additional unique nodes are generated.
    """
    if node_labels_pool is not None and len(node_labels_pool) > 0:
        assert len(node_labels_pool[0]) == str_len
    
    def encode(str_len):
        return ''.join(random.choices(string.ascii_uppercase, k=1)) +\
               ''.join(random.choices(string.ascii_uppercase + string.digits, k=(str_len-1)))
    
    if node_labels_pool is not None:
        if len(node_labels_pool) >= n_nodes:
            # select random nodes from node_labels_pool
            random.shuffle(node_labels_pool)
            vals = node_labels_pool[:n_nodes]
            d = dict(zip(range(n_nodes), vals))
        elif len(node_labels_pool) < n_nodes:
            # use all nodes from node_labels_pool
            random.shuffle(node_labels_pool)
            vals = node_labels_pool[:]
            # generate additional unique nodes
            for i in range(n_nodes - len(node_labels_pool)):
                x = encode(str_len)
                while x in vals: # guarantee uniqueness
                    x = encode(str_len)
                vals.append(x)
            # shuffle pre-defined and newly generated nodes
            random.shuffle(vals)
            d = dict(zip(range(n_nodes), vals))
    else:   # no node_labels_pool provided
        d = {}
        for i in range(n_nodes):
            x = encode(str_len)
            # gaurantee uniqueness
            while x in d.values():
                x = encode(str_len)
            
            d[i] = x
    
    return d

# %% ../nbs/01_synthetic_data.ipynb 14
def create_random_graph(generator:callable, obfuscate:bool=True, node_labels_pool:list[str]=None, **kwargs):
    """Thin wrapper around networkx's random graph generator

    Args:
        generator (callable): networkx's random graph generator
        **kwargs: passed to generator
    Example:
        >>> G = create_random_graph(nx.erdos_renyi_graph, n=10, p=0.6, directed=False)
    """
    G = generator(**kwargs)
    if obfuscate:
        G = nx.relabel_nodes(G, obfuscate_nodes(G.number_of_nodes(), node_labels_pool=node_labels_pool))
    return G

# %% ../nbs/01_synthetic_data.ipynb 19
def create_and_save_random_graph(generator:callable,
                                 label_edges:str,
                                 path_to:str|Path,
                                 obfuscate:bool=True,
                                 node_labels_pool:list[str]=None,
                                 **kwargs) -> nx.classes.graph.Graph:
    """
    Creates a random graph and saves it to path_to.

    Args:
        generator (callable): networkx's random graph generator
        obfuscate (bool, optional): if true, obfuscates node names. Defaults to True.
        path_to (str|Path, optional): path to save the graph. Defaults to None.
    
    Example:
        >>> create_random_graph_and_save(nx.erdos_renyi_graph, n=10, p=0.6, directed=False, path_to=path_data / 'synthetic' / '1.csv')
    """
    path_to = Path(path_to)

    G = create_random_graph(generator=generator, obfuscate=obfuscate, node_labels_pool=node_labels_pool, **kwargs)
    edges = G.edges()
    with open(path_to, 'w', newline='') as csvfile:
        writer = csv.writer(csvfile, delimiter=' ')
        for edge in edges:
            writer.writerow([edge[0], edge[1], label_edges])
    
    return G

# %% ../nbs/01_synthetic_data.ipynb 26
def generate_and_save_graphs(graph_params, path_dir_to):
    """
    Generates and saves multiple random graphs based on provided parameters.

    Parameters:
    graph_params (list of tuples): A list of tuples where each tuple contains:
        - n (int): Number of nodes in the graph.
        - p (float): Probability of edge creation.
        - label_edges (str): Label for the edges.
    path_dir_to (Path): Directory path where the graph files will be saved.
    
    The method iterates over the provided graph parameters, generates an Erdős–Rényi graph,
    saves it to a CSV file, and maintains a pool of node labels to ensure uniqueness.
    A progress bar is displayed to indicate the number of layers being processed.
    """
    path_dir_to = Path(path_dir_to)
    node_labels_pool = []

    graphs = []
    
    for i, (n, p, label_edges) in enumerate(tqdm(graph_params, desc='Generating layers'), start=1):
        G = create_and_save_random_graph(nx.erdos_renyi_graph, label_edges, path_dir_to / f'{i}.csv',
                                         obfuscate=True, node_labels_pool=node_labels_pool, n=n, p=p)
        node_labels_pool.extend(list(G.nodes()))

        graphs.append(G)
    
    # collect metrics of graphs
    total_nodes = sum(G.number_of_nodes() for G in graphs)
    total_edges = sum(G.number_of_edges() for G in graphs)
    total_memory = sum(asizeof(G) for G in graphs) / (1024 ** 2)  # Approximate memory usage, MBs

    return (total_nodes, total_edges, total_memory)
