# Import the ArangoDB-NetworkX Adapter
from adbnx_adapter.adapter import ADBNX_Adapter

# Import a sample graph from NetworkX
from networkx import grid_2d_graph

# This is the connection information for your ArangoDB instance
# (Let's assume that the ArangoDB fraud-detection data dump is imported to this endpoint)
con = {
    "hostname": "localhost",
    "protocol": "http",
    "port": 8529,
    "username": "root",
    "password": "rootpassword",
    "dbName": "_system",
}

# This instantiates your ADBNX Adapter with your connection credentials
adbnx_adapter = ADBNX_Adapter(con)

# ArangoDB to NetworkX via Graph
nx_fraud_graph = adbnx_adapter.arangodb_graph_to_networkx("fraud-detection")

# ArangoDB to NetworkX via Collections
nx_fraud_graph_2 = adbnx_adapter.arangodb_collections_to_networkx(
        "fraud-detection", 
        {"account", "bank", "branch", "Class", "customer"}, # Specify vertex collections
        {"accountHolder", "Relationship", "transaction"} # Specify edge collections
)

# ArangoDB to NetworkX via Metagraph
metagraph = {
    "vertexCollections": {
        "account": {"Balance", "account_type", "customer_id", "rank"},
        "customer": {"Name", "rank"},
    },
    "edgeCollections": {
        "transaction": {"transaction_amt", "sender_bank_id", "receiver_bank_id"},
        "accountHolder": {},
    },
}
nx_fraud_graph_3 = adbnx_adapter.arangodb_to_networkx("fraud-detection", metagraph)

# NetworkX to ArangoDB
nx_grid_graph = grid_2d_graph(5, 5)
adb_grid_edge_definitions = [
    {
        "edge_collection": "to",
        "from_vertex_collections": ["Grid_Node"],
        "to_vertex_collections": ["Grid_Node"],
    }
]
adb_grid_graph = adbnx_adapter.networkx_to_arangodb("Grid", nx_grid_graph, adb_grid_edge_definitions)