%%capture
!git clone -b oasis_connector --single-branch https://github.com/arangodb/interactive_tutorials.git
!git clone -b 1.0.2 --single-branch https://github.com/arangoml/dgl-adapter.git
!rsync -av interactive_tutorials/ ./ --exclude=.git
!pip3 install adbdgl_adapter==1.0.2
!pip3 install matplotlib
!pip3 install pyArango
!pip3 install networkx ## For drawing purposes 

import json
import oasis
import matplotlib.pyplot as plt

import dgl
import torch
import networkx as nx

from dgl import remove_self_loop
from dgl.data import KarateClubDataset
from dgl.data import MiniGCDataset

from adbdgl_adapter.adapter import ADBDGL_Adapter
from adbdgl_adapter.controller import ADBDGL_Controller
from adbdgl_adapter.typings import Json, ArangoMetagraph, DGLCanonicalEType, DGLDataDict

# A homogeneous graph with 6 nodes, and 5 edges
g = dgl.graph(([0, 0, 0, 0, 0], [1, 2, 3, 4, 5]))
print(g)

# Print the graph's canonical edge types
print("\nCanonical Edge Types: ", g.canonical_etypes)
# [('_N', '_E', '_N')]
# '_N' being the only Node type
# '_E' being the only Edge type

# A heterogeneous graph with 8 nodes, and 7 edges
g = dgl.heterograph({
    ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
    ('user', 'follows', 'game'): (torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])),
    ('user', 'plays', 'game'): (torch.tensor([1, 3]), torch.tensor([2, 3]))
})

print(g)
print("\nCanonical Edge Types: ", g.canonical_etypes)
print("\nNode Types: ", g.ntypes)
print("\nEdge Types: ", g.etypes)

# A homogeneous graph with 6 nodes, and 5 edges
g = dgl.graph(([0, 0, 0, 0, 0], [1, 2, 3, 4, 5]))

# Assign an integer value for each node.
g.ndata['x'] = torch.tensor([151, 124, 41, 89, 76, 55])
# Assign a 4-dimensional edge feature vector for each edge.
g.edata['a'] = torch.randn(5, 4)

print(g)
print("\nNode Data X attribute: ", g.ndata['x'])
print("\nEdge Data A attribute: ", g.edata['a'])


# NOTE: The following line ndata insertion will fail, since not all nodes have been assigned an attribute value
# g.ndata['bad_attribute'] = torch.tensor([0,10,20,30,40])

g = dgl.heterograph({
    ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
    ('user', 'follows', 'game'): (torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])),
    ('user', 'plays', 'game'): (torch.tensor([1, 3]), torch.tensor([2, 3]))
})

# Get the number of all nodes in the graph
print("All nodes: ", g.num_nodes())

# Get the number of user nodes
print("User nodes: ", g.num_nodes('user'))

# Nodes of different types have separate IDs,
# hence not well-defined without a type specified
# print(g.nodes())
#DGLError: Node type name must be specified if there are more than one node types.

print(g.nodes('user'))

g = dgl.heterograph({
    ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
    ('user', 'follows', 'game'): (torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])),
    ('user', 'plays', 'game'): (torch.tensor([1, 3]), torch.tensor([2, 3]))
})

g.nodes['user'].data['age'] = torch.tensor([21, 16, 38, 64])
# An alternative (yet equivalent) syntax:
# g.ndata['age'] = {'user': torch.tensor([21, 16, 38, 64])}

print(g.ndata)

# Request temporary instance from the managed ArangoDB Cloud Oasis.
con = oasis.getTempCredentials()

# Connect to the db via the python-arango driver
db = oasis.connect_python_arango(con)

print('\n--------------------')
print("https://{}:{}".format(con["hostname"], con["port"]))
print("Username: " + con["username"])
print("Password: " + con["password"])
print("Database: " + con["dbName"])
print('--------------------\n')

%%capture
!chmod -R 755 ./tools
!./tools/arangorestore -c none --server.endpoint http+ssl://{con["hostname"]}:{con["port"]} --server.username {con["username"]} --server.database {con["dbName"]} --server.password {con["password"]} --replication-factor 3  --input-directory "dgl-adapter/examples/data/fraud_dump" --include-system-collections true

adbdgl_adapter = ADBDGL_Adapter(con)

# Define graph name
graph_name = "fraud-detection"

# Create DGL graph from ArangoDB graph
dgl_g = adbdgl_adapter.arangodb_graph_to_dgl(graph_name)

# You can also provide valid Python-Arango AQL query options to the command above, like such:
# dgl_g = aadbdgl_adapter.arangodb_graph_to_dgl(graph_name, ttl=1000, stream=True)
# See more here: https://docs.python-arango.com/en/main/specs.html#arango.aql.AQL.execute

# Show graph data
print('\n--------------------')
print(dgl_g)
print(dgl_g.ntypes)
print(dgl_g.etypes)

# Define collection
vertex_collections = {"account", "Class", "customer"}
edge_collections = {"accountHolder", "Relationship", "transaction"}

# Create DGL from ArangoDB collections
dgl_g = adbdgl_adapter.arangodb_collections_to_dgl("fraud-detection", vertex_collections, edge_collections)

# You can also provide valid Python-Arango AQL query options to the command above, like such:
# dgl_g = adbdgl_adapter.arangodb_collections_to_dgl("fraud-detection", vertex_collections, edge_collections, ttl=1000, stream=True)
# See more here: https://docs.python-arango.com/en/main/specs.html#arango.aql.AQL.execute

# Show graph data
print('\n--------------------')
print(dgl_g)
print(dgl_g.ntypes)
print(dgl_g.etypes)

# Define Metagraph
fraud_detection_metagraph = {
    "vertexCollections": {
        "account": {"rank", "Balance", "customer_id"},
        "Class": {"concrete"},
        "customer": {"rank"},
    },
    "edgeCollections": {
        "accountHolder": {},
        "Relationship": {},
        "transaction": {"receiver_bank_id", "sender_bank_id", "transaction_amt"},
    },
}

# Create DGL Graph from attributes
dgl_g = adbdgl_adapter.arangodb_to_dgl('FraudDetection',  fraud_detection_metagraph)

# You can also provide valid Python-Arango AQL query options to the command above, like such:
# dgl_g = adbdgl_adapter.arangodb_to_dgl(graph_name = 'FraudDetection',  fraud_detection_metagraph, ttl=1000, stream=True)
# See more here: https://docs.python-arango.com/en/main/specs.html#arango.aql.AQL.execute

# Show graph data
print('\n--------------')
print(dgl_g)
print('\n--------------')
print(dgl_g.ndata)
print('--------------\n')
print(dgl_g.edata)

# Define Metagraph
fraud_detection_metagraph = {
    "vertexCollections": {
        "account": {"rank"},
        "Class": {"concrete", "name"},
        "customer": {"Sex", "Ssn", "rank"},
    },
    "edgeCollections": {
        "accountHolder": {},
        "Relationship": {},
        "transaction": {"receiver_bank_id", "sender_bank_id", "transaction_amt", "transaction_date", "trans_time"},
    },
}

# When converting to DGL via an ArangoDB Metagraph that contains non-numerical values, a user-defined 
# Controller class is required to specify how ArangoDB attributes should be converted to DGL features.
class FraudDetection_ADBDGL_Controller(ADBDGL_Controller):
    """ArangoDB-DGL controller.

    Responsible for controlling how ArangoDB attributes
    are converted into DGL features, and vice-versa.

    You can derive your own custom ADBDGL_Controller if you want to maintain
    consistency between your ArangoDB attributes & your DGL features.
    """

    def _adb_attribute_to_dgl_feature(self, key: str, col: str, val):
        """
        Given an ArangoDB attribute key, its assigned value (for an arbitrary document),
        and the collection it belongs to, convert it to a valid
        DGL feature: https://docs.dgl.ai/en/0.6.x/guide/graph-feature.html.

        NOTE: You must override this function if you want to transfer non-numerical
        ArangoDB attributes to DGL (DGL only accepts 'attributes' (a.k.a features)
        of numerical types). Read more about DGL features here:
        https://docs.dgl.ai/en/0.6.x/new-tutorial/2_dglgraph.html#assigning-node-and-edge-features-to-graph.

        :param key: The ArangoDB attribute key name
        :type key: str
        :param col: The ArangoDB collection of the ArangoDB document.
        :type col: str
        :param val: The assigned attribute value of the ArangoDB document.
        :type val: Any
        :return: The attribute's representation as a DGL Feature
        :rtype: Any
        """
        try:
          if col == "transaction":
            if key == "transaction_date":
              return int(str(val).replace("-", ""))
    
            if key == "trans_time":
              return int(str(val).replace(":", ""))
    
          if col == "customer":
            if key == "Sex":
              return 0 if val == "M" else 1

            if key == "Ssn":
              return int(str(val).replace("-", ""))

          if col == "Class":
            if key == "name":
              if val == "Bank":
                return 0
              elif val == "Branch":
                return 1
              elif val == "Account":
                return 2
              elif val == "Customer":
                return 3
              else:
                return -1
        except (ValueError, TypeError, SyntaxError):
          return 0

        return super()._adb_attribute_to_dgl_feature(key, col, val)

fraud_adbdgl_adapter = ADBDGL_Adapter(con, FraudDetection_ADBDGL_Controller())

# Create DGL Graph from attributes
dgl_g = fraud_adbdgl_adapter.arangodb_to_dgl('FraudDetection',  fraud_detection_metagraph)

# You can also provide valid Python-Arango AQL query options to the command above, like such:
# dgl_g = fraud_adbdgl_adapter.arangodb_to_dgl(graph_name = 'FraudDetection',  fraud_detection_metagraph, ttl=1000, stream=True)
# See more here: https://docs.python-arango.com/en/main/specs.html#arango.aql.AQL.execute

# Show graph data
print('\n--------------')
print(dgl_g)
print('\n--------------')
print(dgl_g.ndata)
print('--------------\n')
print(dgl_g.edata)

# Load the dgl graph & draw
dgl_karate_graph = KarateClubDataset()[0]
nx.draw(dgl_karate_graph.to_networkx(), with_labels=True)

# Create the ArangoDB graph
name = "Karate"
db.delete_graph(name, drop_collections=True, ignore_missing=True)
adb_karate_graph = adbdgl_adapter.dgl_to_arangodb(name, dgl_karate_graph)

print('\n--------------------')
print("https://{}:{}".format(con["hostname"], con["port"]))
print("Username: " + con["username"])
print("Password: " + con["password"])
print("Database: " + con["dbName"])
print('--------------------\n')
print(f"Inspect the graph here: https://tutorials.arangodb.cloud:8529/_db/{con['dbName']}/_admin/aardvark/index.html#graph/{name}\n")
print(f"View the original graph below:")

# Load the dgl graphs & draw
dgl_lollipop_graph = remove_self_loop(MiniGCDataset(8, 7, 8)[3][0])
plt.figure(1)
nx.draw(dgl_lollipop_graph.to_networkx(), with_labels=True)

dgl_hypercube_graph = remove_self_loop(MiniGCDataset(8, 8, 9)[4][0])
plt.figure(2)
nx.draw(dgl_hypercube_graph.to_networkx(), with_labels=True)

dgl_clique_graph = remove_self_loop(MiniGCDataset(8, 6, 7)[6][0])
plt.figure(3)
nx.draw(dgl_clique_graph.to_networkx(), with_labels=True)

# Create the ArangoDB graphs
lollipop = "Lollipop"
hypercube = "Hypercube"
clique = "Clique"

db.delete_graph(lollipop, drop_collections=True, ignore_missing=True)
db.delete_graph(hypercube, drop_collections=True, ignore_missing=True)
db.delete_graph(clique, drop_collections=True, ignore_missing=True)

adb_lollipop_graph = adbdgl_adapter.dgl_to_arangodb(lollipop, dgl_lollipop_graph)
adb_hypercube_graph = adbdgl_adapter.dgl_to_arangodb(hypercube, dgl_hypercube_graph)
adb_clique_graph = adbdgl_adapter.dgl_to_arangodb(clique, dgl_clique_graph)

print('\n--------------------')
print("https://{}:{}".format(con["hostname"], con["port"]))
print("Username: " + con["username"])
print("Password: " + con["password"])
print("Database: " + con["dbName"])
print('--------------------\n')
print("\nInspect the graphs here:\n")
print(f"1) https://tutorials.arangodb.cloud:8529/_db/{con['dbName']}/_admin/aardvark/index.html#graph/{lollipop}")
print(f"2) https://tutorials.arangodb.cloud:8529/_db/{con['dbName']}/_admin/aardvark/index.html#graph/{hypercube}")
print(f"3) https://tutorials.arangodb.cloud:8529/_db/{con['dbName']}/_admin/aardvark/index.html#graph/{clique}\n")
print(f"\nView the original graphs below:")

from torch.functional import Tensor

# Load the dgl graphs
dgl_lollipop_graph = remove_self_loop(MiniGCDataset(8, 7, 8)[3][0])
dgl_hypercube_graph = remove_self_loop(MiniGCDataset(8, 8, 9)[4][0])
dgl_clique_graph = remove_self_loop(MiniGCDataset(8, 6, 7)[6][0])

 # Add DGL Node & Edge Features to each graph
dgl_lollipop_graph.ndata["random_ndata"] = torch.tensor(
    [[i, i, i] for i in range(0, dgl_lollipop_graph.num_nodes())]
)
dgl_lollipop_graph.edata["random_edata"] = torch.rand(dgl_lollipop_graph.num_edges())

dgl_hypercube_graph.ndata["random_ndata"] = torch.rand(dgl_hypercube_graph.num_nodes())
dgl_hypercube_graph.edata["random_edata"] = torch.tensor(
    [[[i], [i], [i]] for i in range(0, dgl_hypercube_graph.num_edges())]
)

dgl_clique_graph.ndata['clique_ndata'] = torch.tensor([1,2,3,4,5,6])
dgl_clique_graph.edata['clique_edata'] = torch.tensor(
    [1 if i % 2 == 0 else 0 for i in range(0, dgl_clique_graph.num_edges())]
)


# When converting to ArangoDB from DGL, a user-defined Controller class
# is required to specify how DGL features (aka attributes) should be converted 
# into ArangoDB attributes. NOTE: A custom Controller is NOT needed you want to
# keep the numerical-based values of your DGL features.
class Clique_ADBDGL_Controller(ADBDGL_Controller):
    """ArangoDB-DGL controller.

    Responsible for controlling how ArangoDB attributes
    are converted into DGL features, and vice-versa.

    You can derive your own custom ADBDGL_Controller if you want to maintain
    consistency between your ArangoDB attributes & your DGL features.
    """

    def _dgl_feature_to_adb_attribute(self, key: str, col: str, val: Tensor):
        """
        Given a DGL feature key, its assigned value (for an arbitrary node or edge),
        and the collection it belongs to, convert it to a valid ArangoDB attribute
        (e.g string, list, number, ...).

        NOTE: No action is needed here if you want to keep the numerical-based values
        of your DGL features.

        :param key: The DGL attribute key name
        :type key: str
        :param col: The ArangoDB collection of the (soon-to-be) ArangoDB document.
        :type col: str
        :param val: The assigned attribute value of the DGL node.
        :type val: Tensor
        :return: The feature's representation as an ArangoDB Attribute
        :rtype: Any
        """
        if key == "clique_ndata":
          if val == 1:
            return "one is fun"
          elif val == 2:
            return "two is blue"
          elif val == 3:
            return "three is free"
          elif val == 4:
            return "four is more"
          else: # No special string for values 5 & 6
            return f"ERROR! Unrecognized value, got {val}"

        if key == "clique_edata":
          return bool(val)

        return super()._dgl_feature_to_adb_attribute(key, col, val)

# Re-instantiate a new adapter specifically for the Clique Graph Conversion
clique_adbgl_adapter = ADBDGL_Adapter(con, Clique_ADBDGL_Controller())

# Create the ArangoDB graphs
lollipop = "Lollipop_With_Attributes"
hypercube = "Hypercube_With_Attributes"
clique = "Clique_With_Attributes"

db.delete_graph(lollipop, drop_collections=True, ignore_missing=True)
db.delete_graph(hypercube, drop_collections=True, ignore_missing=True)
db.delete_graph(clique, drop_collections=True, ignore_missing=True)

adb_lollipop_graph = adbdgl_adapter.dgl_to_arangodb(lollipop, dgl_lollipop_graph)
adb_hypercube_graph = adbdgl_adapter.dgl_to_arangodb(hypercube, dgl_hypercube_graph)
adb_clique_graph = clique_adbgl_adapter.dgl_to_arangodb(clique, dgl_clique_graph) # Notice the new adapter here!

print('\n--------------------')
print("https://{}:{}".format(con["hostname"], con["port"]))
print("Username: " + con["username"])
print("Password: " + con["password"])
print("Database: " + con["dbName"])
print('--------------------\n')
print("\nInspect the graphs here:\n")
print(f"1) https://tutorials.arangodb.cloud:8529/_db/{con['dbName']}/_admin/aardvark/index.html#graph/{lollipop}")
print(f"2) https://tutorials.arangodb.cloud:8529/_db/{con['dbName']}/_admin/aardvark/index.html#graph/{hypercube}")
print(f"3) https://tutorials.arangodb.cloud:8529/_db/{con['dbName']}/_admin/aardvark/index.html#graph/{clique}\n")