"""
ionerdss.model.graph_based.complexes.examples

This module contains example graphs representing different types of complexes
for testing and demonstration purposes.

"""

import networkx as nx

def graph_5l93():
    """
    Graph representing HIV-1 CA-SP1, which is a hexagonal complex
    with 18 nodes and 24 edges.

    @note: download PDB file from https://www.rcsb.org/structure/5L93
    for more information.

    @warning: This graph is not generated by ionerdss, but a synthetic example.

    Returns:
        nx.Graph: The graph representing the complex.
    """
    G = nx.Graph()
    G.add_nodes_from([(i, {"type": "A"}) for i in range(18)])
    G.add_edges_from([
        (0, 1, {"type": "hex"}), (1, 2, {"type": "hex"}), (2, 3, {"type": "hex"}),
        (3, 4, {"type": "hex"}), (4, 5, {"type": "hex"}), (5, 0, {"type": "hex"}),
        (6, 7, {"type": "tri"}), (7, 0, {"type": "tri"}), (0, 6, {"type": "tri"}),
        (7, 1, {"type": "di"}), (8, 9, {"type": "tri"}), (9, 1, {"type": "tri"}),
        (1, 8, {"type": "tri"}), (9, 2, {"type": "di"}), (10, 11, {"type": "tri"}),
        (11, 2, {"type": "tri"}), (2, 10, {"type": "tri"}), (11, 3, {"type": "di"}),
        (12, 13, {"type": "tri"}), (13, 3, {"type": "tri"}), (3, 12, {"type": "tri"}),
        (13, 4, {"type": "di"}), (14, 15, {"type": "tri"}), (15, 4, {"type": "tri"}),
        (4, 14, {"type": "tri"}), (15, 5, {"type": "di"}), (16, 17, {"type": "tri"}),
        (17, 5, {"type": "tri"}), (5, 16, {"type": "tri"}), (17, 0, {"type": "di"}),
    ])
    return G

def graph_8y7s():
    """
    Graph representing a homotetramer with 4 nodes and 6 edges.

    @note: download PDB file from https://www.rcsb.org/structure/8Y7S
    for more information.

    @warning: This graph is not generated by ionerdss, but a synthetic example.

    Returns:
        nx.Graph: The graph representing the complex.
    """
    G = nx.Graph()
    G.add_nodes_from([(i, {"type": "X"}) for i in range(4)])
    G.add_edges_from([
        (0, 1, {"type": "a"}), (0, 2, {"type": "b"}), (0, 3, {"type": "c"}),
        (2, 3, {"type": "a"}), (1, 3, {"type": "b"}), (1, 2, {"type": "c"}),
    ])
    return G

def graph_asymmetry_4mer():
    """
    Graph representing an asymmetric tetramer with 4 nodes and 4 edges.

    @warning: This graph is a synthetic example, not representing any real complex.

    Returns:
        nx.Graph: The graph representing the complex.
    """
    G = nx.Graph()
    G.add_nodes_from([(0, {"type": "A"}), (1, {"type": "B"}), (2, {"type": "C"}), (3, {"type": "A"})])
    G.add_edges_from([
        (0, 1, {"type": "ab"}), (1, 2, {"type": "bc"}), (2, 3, {"type": "ca"}), (3, 1, {"type": "ab"})
    ])
    return G

################################
# Graph Registry
################################

def graph_ring_25():
    """
    Graph representing a 25-node ring (cycle) complex.
    
    @warning: This graph is a synthetic example to benchmark algorithmic scaling.
    
    Returns:
        nx.Graph: The graph representing the complex.
    """
    G = nx.Graph()
    G.add_nodes_from([(i, {"type": "A"}) for i in range(25)])
    edges = [(i, (i + 1) % 25, {"type": "bond"}) for i in range(25)]
    G.add_edges_from(edges)
    return G

def graph_complex_30mer():
    """
    Graph representing a complicated 30-mer complex.
    Nodes: 30
    Average connection degree: 6
    Node types: 5 different types
    Edge types: 3 different varying bond types
    
    @warning: This graph is a synthetic example to benchmark algorithmic scaling on very dense graphs.
    
    Returns:
        nx.Graph: The graph representing the complex.
    """
    G = nx.Graph()
    node_types = ["A", "B", "C", "D", "E"]
    G.add_nodes_from([(i, {"type": node_types[i % 5]}) for i in range(30)])
    
    edge_types = ["bond1", "bond2", "bond3"]
    edges = []
    # To achieve average degree 6 uniformly, connect each node to its 3 nearest neighbors in a ring topology
    for i in range(30):
        for step in [1, 2, 3]:
            j = (i + step) % 30
            edges.append((i, j, {"type": edge_types[(i + j) % 3]}))
            
    G.add_edges_from(edges)
    return G

# Registry of example graphs
graph_registry = {
    "5l93": graph_5l93,
    "8y7s": graph_8y7s,
    "asymmetry_4mer": graph_asymmetry_4mer,
    "ring_25": graph_ring_25,
    "complex_30mer": graph_complex_30mer,
}

def get_example(name: str):
    """
    Get example with a key from the registry.
    See `graph_registry` for available examples.

    Args:
        name (str): The key of the example.

    Returns:
        nx.Graph: The graph representing the complex.

    Raises:
        ValueError: If the key is not found in the registry.
    """
    try:
        return graph_registry[name]()
    except KeyError:
        raise ValueError(f"Unknown example: {name}. Available: {list(graph_registry)}")
    
def show_available_examples(do_print = True):
    """
    Return the registry of examples.
    
    Args:
        do_print (bool): Whether to print the registry.

    Returns:
        list: The registry of examples.
    """
    if do_print:
        print("Available examples:", list(graph_registry.keys()))
    return(list(graph_registry.keys()))
