%%capture
!pip install torch
!pip install adbpyg-adapter==1.0.0
!pip install adb-cloud-connector
!git clone -b 1.0.0 --single-branch https://github.com/arangoml/pyg-adapter.git

## For drawing purposes 
!pip install matplotlib
!pip install networkx

# All imports

import pandas
import torch
from torch_geometric.data import Data, HeteroData
from torch_geometric.datasets import FakeDataset, FakeHeteroDataset, KarateClub
from torch_geometric.utils import to_networkx
from torch_geometric.typing import EdgeType

from adbpyg_adapter import ADBPyG_Adapter, ADBPyG_Controller
from adbpyg_adapter.encoders import IdentityEncoder, CategoricalEncoder
from adbpyg_adapter.typings import Json, ADBMetagraph, PyGMetagraph

from arango import ArangoClient
from adb_cloud_connector import get_temp_credentials

import json
import logging

import matplotlib.pyplot as plt
import networkx as nx

edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype=torch.long)
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)

data = Data(x=x, edge_index=edge_index)
print(data)

print(data.keys)

print(data['x'])


for key, item in data:
    print(f'{key} found in data')

print('edge_attr' in data)
print(data.num_nodes)
print(data.num_edges)
print(data.num_node_features)
print(data.has_isolated_nodes())
print(data.has_self_loops())
print(data.is_directed())

# Transfer data object to GPU (requires Tesla T4 GPU if running in Colab)
# device = torch.device('cuda')
# data = data.to(device)

data = HeteroData()

data["user"].x = torch.tensor([[21], [16], [38], [64]])
data[("user", "follows", "user")].edge_index = torch.tensor([[0, 1], [1, 2]])
data[("user", "follows", "game")].edge_index = torch.tensor([[0, 1, 2], [0, 1, 2]])
data[("user", "plays", "game")].edge_index = torch.tensor([[3, 3], [1, 2]])
data[("user", "plays", "game")].edge_attr = torch.tensor([[3], [5]])

print(data)
print(data.node_types)
print(data.edge_types)

# Request temporary instance from the managed ArangoDB Cloud Service.
con = get_temp_credentials()
print(json.dumps(con, indent=2))

# Connect to the db via the python-arango driver
db = ArangoClient(hosts=con["url"]).db(con["dbName"], con["username"], con["password"], verify=True)

!chmod -R 755 pyg-adapter/
!./pyg-adapter/tests/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 "pyg-adapter/tests/data/adb/imdb_dump" --include-system-collections true

# Create the IMDB graph
db.create_graph(
    "imdb",
    edge_definitions=[
        {
            "edge_collection": "Ratings",
            "from_vertex_collections": ["Users"],
            "to_vertex_collections": ["Movies"],
        },
    ],
)

adbpyg_adapter = ADBPyG_Adapter(db)

# Create the PyG graph & draw it
pyg_karate_graph = KarateClub()[0]
print(pyg_karate_graph)
nx.draw(to_networkx(pyg_karate_graph), with_labels=True)

name = "Karate"

# Delete the graph if it already exists
db.delete_graph(name, drop_collections=True, ignore_missing=True)

# Create the ArangoDB graph
adb_karate_graph = adbpyg_adapter.pyg_to_arangodb(name, pyg_karate_graph)

# You can also provide valid Python-Arango Import Bulk options to the command above, like such:
# adb_karate_graph = adbpyg_adapter.pyg_to_arangodb(name, pyg_karate_graph, batch_size=5, on_duplicate="replace")
# See the full parameter list at https://docs.python-arango.com/en/main/specs.html#arango.collection.Collection.import_bulk

print('\n--------------------')
print("URL: " + con["url"])
print("Username: " + con["username"])
print("Password: " + con["password"])
print("Database: " + con["dbName"])
print('--------------------\n')
print(f"View the created graph here: {con['url']}/_db/{con['dbName']}/_admin/aardvark/index.html#graph/{name}\n")
print(f"View the original graph below:\n")

# Create the PyG graph & draw it
pyg_homo_graph = FakeDataset(avg_num_nodes=30, edge_dim=1)[0] # 'edge_weight' property
print(pyg_homo_graph)
nx.draw(to_networkx(pyg_homo_graph), with_labels=True)

name = "FakeHomo"

# Delete the graph if it already exists
db.delete_graph(name, drop_collections=True, ignore_missing=True)

# Create the ArangoDB graph
adb_homo_graph = adbpyg_adapter.pyg_to_arangodb(name, pyg_homo_graph)

print('\n--------------------')
print("URL: " + con["url"])
print("Username: " + con["username"])
print("Password: " + con["password"])
print("Database: " + con["dbName"])
print('--------------------\n')
print(f"View the created graph here: {con['url']}/_db/{con['dbName']}/_admin/aardvark/index.html#graph/{name}\n")
print(f"View the original graph below:\n")

# Create the PyG graph
pyg_hetero_graph = FakeHeteroDataset(avg_num_nodes=30, edge_dim=2)[0] # 'edge_attr' property

name = "FakeHetero"

# Delete the graph if it already exists
db.delete_graph(name, drop_collections=True, ignore_missing=True)

# Create the ArangoDB graphs
adb_hetero_graph = adbpyg_adapter.pyg_to_arangodb(name, pyg_hetero_graph)

print('\n--------------------')
print("URL: " + con["url"])
print("Username: " + con["username"])
print("Password: " + con["password"])
print("Database: " + con["dbName"])
print('--------------------\n')
print(f"View the created graph here: {con['url']}/_db/{con['dbName']}/_admin/aardvark/index.html#graph/{name}\n")
print(f"View the original graph below:\n")

# Create the PyG graph
pyg_hetero_graph = FakeHeteroDataset(
    num_node_types=2,
    num_edge_types=3,
    avg_num_nodes=20,
    avg_num_channels=3,  # avg number of features per node
    edge_dim=2,  # number of features per edge
    num_classes=3,  # number of unique label values
)[0]
print(pyg_hetero_graph)
nx.draw(to_networkx(pyg_hetero_graph.to_homogeneous()), with_labels=True)

# Define the metagraph
def y_tensor_to_2_column_dataframe(pyg_tensor):
    label_map = {0: "Kiwi", 1: "Blueberry", 2: "Avocado"}

    df = pandas.DataFrame(columns=["label_num", "label_str"])
    df["label_num"] = pyg_tensor.tolist()
    df["label_str"] = df["label_num"].map(label_map)

    return df

metagraph = {
    "nodeTypes": {
        "v0": {
            "x": "features",  # 1) you can specify a string value for attribute renaming
            "y": y_tensor_to_2_column_dataframe,  # 2) you can specify a function for user-defined handling, as long as the function returns a Pandas DataFrame
        },
    },
    "edgeTypes": {
        ("v0", "e0", "v0"): {
            # 3) you can specify a list of strings for tensor dissasembly (if you know the number of node/edge features in advance)
            "edge_attr": [ "a", "b"]  
        },
    },
}

name = "FakeHetero"

db.delete_graph(name, drop_collections=True, ignore_missing=True)

# Create the ArangoDB graph with `explicit_metagraph=False`
adb_hetero_graph = adbpyg_adapter.pyg_to_arangodb(name, pyg_hetero_graph, metagraph, explicit_metagraph=False)

# Create the ArangoDB graph with `explicit_metagraph=True`
# With `explicit_metagraph=True`, the node & edge types omitted from the metagraph will NOT be converted to ArangoDB.
# Only 'v0' and ('v0', 'e0', 'v0') will be brought over (i.e 'v1', ('v0', 'e0', 'v1'), ... are ignored):
## adb_hetero_graph_explicit = adbpyg_adapter.pyg_to_arangodb(name, pyg_hetero_graph, metagraph, explicit_metagraph=True)

print('\n--------------------')
print("URL: " + con["url"])
print("Username: " + con["username"])
print("Password: " + con["password"])
print("Database: " + con["dbName"])
print('--------------------\n')
print(f"View the created graph here: {con['url']}/_db/{con['dbName']}/_admin/aardvark/index.html#graph/{name}\n")
print(f"View the original graph below:\n")

# Create the PyG graph
pyg_hetero_graph = FakeHeteroDataset(avg_num_nodes=30, edge_dim=2)[0] # 'edge_attr' property

name = "FakeHetero"

db.delete_graph(name, drop_collections=True, ignore_missing=True)

# Create a custom ADBPyG_Controller
class Custom_ADBPyG_Controller(ADBPyG_Controller):
    """ArangoDB-PyG controller.

    Responsible for controlling how nodes & edges are handled when
    transitioning from PyG to ArangoDB.

    You can derive your own custom ADBPyG_Controller.
    """

    def _prepare_pyg_node(self, pyg_node: dict, col: str) -> dict:
        """Optionally modify a PyG node object before it gets inserted into its designated ArangoDB collection.

        :param pyg_node: The PyG node object to (optionally) modify.
        :param col: The ArangoDB collection the PyG node belongs to.
        :return: The PyG Node object
        """
        pyg_node["foo"] = "bar"
        return pyg_node

    def _prepare_pyg_edge(self, pyg_edge: dict, edge_type: tuple) -> dict:
        """Optionally modify a PyG edge object before it gets inserted into its designated ArangoDB collection.

        :param pyg_edge: The PyG edge object to (optionally) modify.
        :param edge_type: The Edge Type of the PyG edge. Formatted
            as (from_collection, edge_collection, to_collection)
        :return: The PyG Edge object
        """
        pyg_edge["bar"] = "foo"
        return pyg_edge

# Instantiate new adapter & create the ArangoDB graph
adb_hetero_graph = ADBPyG_Adapter(db, Custom_ADBPyG_Controller()).pyg_to_arangodb(name, pyg_hetero_graph)

print('\n--------------------')
print("URL: " + con["url"])
print("Username: " + con["username"])
print("Password: " + con["password"])
print("Database: " + con["dbName"])
print('--------------------\n')
print(f"View the created graph here: {con['url']}/_db/{con['dbName']}/_admin/aardvark/index.html#graph/{name}\n")

# Start from scratch! (with a smaller graph)
data = FakeHeteroDataset(
    num_node_types=2,
    num_edge_types=3,
    avg_num_nodes=20,
    avg_num_channels=3,  # avg number of features per node
    edge_dim=2,  # number of features per edge
    num_classes=3,  # number of unique label values
)[0]

adbpyg_adapter.pyg_to_arangodb("FakeHetero", data, overwrite_graph=True, overwrite=True)

# Define graph name
graph_name = "FakeHetero"

# Create PyG graph from the ArangoDB graph
pyg_hetero_graph = adbpyg_adapter.arangodb_graph_to_pyg(graph_name)

# You can also provide valid Python-Arango AQL query options to the command above, like such:
# pyg_hetero_graph = adbpyg_adapter.arangodb_graph_to_pyg(graph_name, ttl=1000, stream=True)
# See the full parameter list at https://docs.python-arango.com/en/main/specs.html#arango.aql.AQL.execute

# Show graph data
print('\n--------------------')
print(pyg_hetero_graph)

# Define collection names
v_cols = {"v0", "v1"}
e_cols = {"e0"}

# Create PyG graph from the ArangoDB collections
pyg_hetero_graph = adbpyg_adapter.arangodb_collections_to_pyg("FakeHetero", v_cols, e_cols)

# Show graph data
print('\n--------------------')
print(pyg_hetero_graph)

# Define the Metagraph that transfers ArangoDB attributes "as is",
# meaning the data is already formatted to PyG data standards
metagraph_v1 = {
    "vertexCollections": {
        # we instruct the adapter to create the "x" and "y" tensor data from the "x" and "y" ArangoDB attributes
        "v0": { "x": "x", "y": "y"},  
        "v1": {"x": "x"},
    },
    "edgeCollections": {
        "e0": {"edge_attr": "edge_attr"},
    },
}

# Create PyG Graph
pyg_hetero_graph = adbpyg_adapter.arangodb_to_pyg("FakeHetero", metagraph_v1)

# Show graph data
print('\n--------------------')
print(pyg_hetero_graph)

# Define the Metagraph that transfers attributes via user-defined encoders
metagraph_v2 = {
    "vertexCollections": {
        "Movies": {
            "x": {  # Build a feature matrix from the "Action" & "Drama" document attributes
                "Action": IdentityEncoder(dtype=torch.long),
                "Drama": IdentityEncoder(dtype=torch.long),
            },
            "y": "Comedy",
        },
        "Users": {
            "x": {
                "Gender": CategoricalEncoder(), # CategoricalEncoder(mapping={"M": 0, "F": 1}),
                "Age": IdentityEncoder(dtype=torch.long),
            }
        },
    },
    "edgeCollections": {
        "Ratings": {
            "edge_weight": "Rating"
        }
    },
}

# Create PyG Graph
pyg_imdb_graph = adbpyg_adapter.arangodb_to_pyg("IMDB", metagraph_v2)

# Show graph data
print('\n--------------------')
print(pyg_imdb_graph)

# Define the metagraph that transfers attributes via user-defined functions
def udf_v0_x(v0_df):
    # process v0_df here to return v0 "x" feature matrix
    # v0_df["x"] = ...
    return torch.tensor(v0_df["x"].to_list())


def udf_v1_x(v1_df):
    # process v1_df here to return v1 "x" feature matrix
    # v1_df["x"] = ...
    return torch.tensor(v1_df["x"].to_list())


metagraph_v3 = {
    "vertexCollections": {
        "v0": {
            "x": udf_v0_x,  # supports named functions
            "y": lambda df: torch.tensor(df["y"].to_list()),  # also supports lambda functions
        },
        "v1": {"x": udf_v1_x},
    },
    "edgeCollections": {
        "e0": {"edge_attr": (lambda df: torch.tensor(df["edge_attr"].to_list()))},
    },
}

# Create PyG Graph
pyg_hetero_graph = adbpyg_adapter.arangodb_to_pyg("FakeHetero", metagraph_v3)

# Show graph data
print('\n--------------------')
print(pyg_hetero_graph)