############################################################################
#                                                                          #
# Imports & Constants                                                      #
#                                                                          #
############################################################################

import argparse
import gpudb
import sys
import json
from collections import OrderedDict
from tabulate import tabulate
from operator import itemgetter


OPTION_NO_DROP_ERROR = {"no_error_if_not_exists": "true"}
OPTION_NO_CREATE_ERROR = {"no_error_if_exists": "true"}

SCHEMA = "graph_q_social"
TABLE_P = SCHEMA + ".people"
TABLE_K = SCHEMA + ".knows"

GRAPH_S = SCHEMA + "." + "social_relationships"
TABLE_Q1 = GRAPH_S + "_queried_jane_to_chess"
TABLE_Q2 = GRAPH_S + "_queried_males"
TABLE_Q3 = GRAPH_S + "_queried_females_or_chess"
TABLE_Q4 = GRAPH_S + "_queried_females_to_chess"
TABLE_Q1_TARGETS = TABLE_Q1 + "_nodes"
TABLE_Q2_TARGETS = TABLE_Q2 + "_nodes"
TABLE_Q3_TARGETS = TABLE_Q3 + "_nodes"
TABLE_Q4_TARGETS = TABLE_Q4 + "_nodes"


############################################################################
#                                                                          #
# Table Setup Function                                                     #
#                                                                          #
############################################################################

def table_setup():
    """ Setup necessary source tables for graph solver examples. """

    print("\n===========")
    print("TABLE SETUP")
    print("===========\n")

    # Clear any related tables in case they already exist
    kinetica.clear_table(table_name=TABLE_P, options=OPTION_NO_DROP_ERROR)
    kinetica.clear_table(table_name=TABLE_K, options=OPTION_NO_DROP_ERROR)

    # Create the graph example schema, if it doesn't exist
    kinetica.create_schema(SCHEMA, options=OPTION_NO_CREATE_ERROR)

    # Create the People table
    try:
        table_p_obj = gpudb.GPUdbTable(
            _type = [
                ["name", "string", "char16"],
                ["age", "int"],
                ["interest", "string", "char16"],
                ["gender", "string", "char8"]
            ],
            name = TABLE_P,
            db = kinetica,
            options = {}
        )
        print("{} table object successfully created.".format(TABLE_P))
    except gpudb.GPUdbException as e:
        print("{} table object creation failure: {}".format(TABLE_P, str(e)))

    # Define records and insert them into the People table
    p_records = [
        ["Susan", 22, "dance", "female"],
        ["Bill", 60, "golf", "male"],
        ["Alex", 34, "chess", "male"],
        ["Jane", 40, "business", "female"],
        ["Tom", 29, "chess", "male"]
    ]
    table_p_obj.insert_records(p_records)
    print("{} records inserted into {} table.".format(table_p_obj.size(), TABLE_P))

    # Create the Relation table
    try:
        table_k_obj = gpudb.GPUdbTable(
            _type = [
                ["name1", "string", "char16"],
                ["name2", "string", "char16"],
                ["since", "long"],
                ["relation", "string", "char32"]
            ],
            name = TABLE_K,
            db = kinetica,
            options = {}
        )
        print("{} table object successfully created.".format(TABLE_K))
    except gpudb.GPUdbException as e:
        print("{} table object creation failure: {}".format(TABLE_K, str(e)))

    # Define records and insert them into the Relation table
    k_records = [
        ["Jane", "Bill", 2010, "friend"],
        ["Bill", "Susan", 1990, "friend"],
        ["Bill", "Alex", 2001, "family"],
        ["Alex", "Tom", 2001, "friend"],
        ["Susan", "Alex", 2002, "friend"]
    ]
    table_k_obj.insert_records(k_records)
    print("{} records inserted into {} table.".format(table_k_obj.size(), TABLE_K))


# end table_setup()


############################################################################
#                                                                          #
# Graph Setup Function                                                     #
#                                                                          #
############################################################################

def graph_setup():
    """ Setup graphs to be used in query examples. """

    print("\n===========")
    print("GRAPH SETUP")
    print("===========\n")

    # Create a graph from TABLE_P and TABLE_K
    print("Creating {}".format(GRAPH_S))
    create_s_graph_response = kinetica.create_graph(
        graph_name = GRAPH_S,
        directed_graph = False,
        nodes = [
            TABLE_P + ".name AS NAME",
            TABLE_P + ".interest AS LABEL",
            "",
            TABLE_P + ".name AS NAME",
            TABLE_P + ".gender AS LABEL"
        ],
        edges = [
            TABLE_K + ".name1 AS NODE1_NAME",
            TABLE_K + ".name2 AS NODE2_NAME",
            TABLE_K + ".relation AS LABEL"
        ],
        weights = [],
        restrictions = [],
        options = {
            "recreate": "true"
        }
    )
    if create_s_graph_response["status_info"]["status"] == "OK":
        print("{} creation success!".format(GRAPH_S))
        print("Number of nodes: {}".format(create_s_graph_response["num_nodes"]))
        print("Number of edges: {}".format(create_s_graph_response["num_edges"]))
    else:
        print("{} creation failure: \n\t{}".format(
            GRAPH_S, 
            create_s_graph_response["status_info"]["message"]
        ))
        print("Exiting...")
        sys.exit()

# end graph_setup()


############################################################################
#                                                                          #
# Query Graph Function                                                     #
#                                                                          #
############################################################################

def query_graph_example():
    """Demonstrate querying a graph using integer IDs."""

    print("\n=====================")
    print("QUERY GRAPH EXAMPLE 1")
    print("=====================\n")

    # Clear any related tables in case they already exist
    kinetica.clear_table(table_name=TABLE_Q1, options=OPTION_NO_DROP_ERROR)
    kinetica.clear_table(table_name=TABLE_Q1_TARGETS, options=OPTION_NO_DROP_ERROR)

    # Query GRAPH_S for nodes connected to 'Jane' where the node has an
    # interest in chess and they are not connected via 'family'
    print(
        "Querying {} for nodes connected to Jane where the node has an "
        "interest in chess and it is not connected via family".format(GRAPH_S)
    )

    query1_s_graph_response = kinetica.query_graph(
        graph_name = GRAPH_S,
        queries = [
            "{'Jane'} AS NODE_NAME",
            "",
            "{'chess'} AS TARGET_NODE_LABEL"
        ],
        restrictions = [
            "{'family'} AS EDGE_LABEL",
            "{0} AS ONOFFCOMPARED"
        ],
        adjacency_table = TABLE_Q1,
        rings = 4
    )
    if query1_s_graph_response["status_info"]["status"] == "OK":
        print("{} graph queried successfully.".format(GRAPH_S))

        # Pretty print query results and targets
        print("\nQuery results for adjacency table {}:".format(TABLE_Q1))
        tabulate_records(kinetica, TABLE_Q1)
      
        print("\nQuery results for target nodes table {}:".format(TABLE_Q1_TARGETS))
        tabulate_records(kinetica, TABLE_Q1_TARGETS)
    else:
        print("{} graph query failure: \n\t{}".format(
            GRAPH_S,
            query1_s_graph_response["status_info"]["message"]
        ))
        print("Exiting...")
        sys.exit()

    print("\n=====================")
    print("QUERY GRAPH EXAMPLE 2")
    print("=====================\n")

    # Clear any related tables in case they already exist
    kinetica.clear_table(table_name=TABLE_Q2, options=OPTION_NO_DROP_ERROR)
    kinetica.clear_table(table_name=TABLE_Q2_TARGETS, options=OPTION_NO_DROP_ERROR)

    # Query GRAPH_S for the nodes connected directly to 'males'
    print("Querying {} for nodes connected to directly to 'males'".format(GRAPH_S))

    query2_s_graph_response = kinetica.query_graph(
        graph_name = GRAPH_S,
        queries = [
            "{'male'} AS NODE_LABEL"
        ],
        restrictions = [],
        adjacency_table = TABLE_Q2,
        rings = 1
    )
    if query2_s_graph_response["status_info"]["status"] == "OK":
        print("{} graph queried successfully.".format(GRAPH_S))

        # Pretty print query results and targets
        print("\nQuery results for adjacency table {}:".format(TABLE_Q2))
        tabulate_records(kinetica, TABLE_Q2, ["QUERY_EDGE_ID"])

        print("\nQuery results for target nodes table {}:".format(TABLE_Q2_TARGETS))
        tabulate_records(kinetica, TABLE_Q2_TARGETS, ["QUERY_NODE_ID_TARGET"])
    else:
        print("{} graph query failure: \n\t{}".format(
            GRAPH_S,
            query2_s_graph_response["status_info"]["message"]
        ))
        print("Exiting...")
        sys.exit()

    print("\n=====================")
    print("QUERY GRAPH EXAMPLE 3")
    print("=====================\n")

    # Clear any related tables in case they already exist
    kinetica.clear_table(table_name=TABLE_Q3, options=OPTION_NO_DROP_ERROR)
    kinetica.clear_table(table_name=TABLE_Q3_TARGETS, options=OPTION_NO_DROP_ERROR)

    # Query GRAPH_S for 'female' nodes or nodes interested in 'chess'
    print("Querying {} for 'female' nodes or nodes interested in 'chess'.".format(GRAPH_S))

    query3_s_graph_response = kinetica.query_graph(
        graph_name = GRAPH_S,
        queries = [
            "{'female', 'chess'} AS NODE_LABEL",
        ],
        restrictions = [],
        adjacency_table = TABLE_Q3,
        rings = 0
    )
    if query3_s_graph_response["status_info"]["status"] == "OK":
        print("{} graph queried successfully.".format(GRAPH_S))
        
        # Pretty print query results and targets
        print("\nQuery results for adjacency table {}:".format(TABLE_Q3))
        tabulate_records(kinetica, TABLE_Q3)

        print("\nQuery results for target nodes table {}:".format(TABLE_Q3_TARGETS))
        tabulate_records(kinetica, TABLE_Q3_TARGETS, ["QUERY_NODE_ID_TARGET"])
    else:
        print("{} graph query failure: \n\t{}".format(
            GRAPH_S,
            query3_s_graph_response["status_info"]["message"]
        ))
        print("Exiting...")
        sys.exit()

    print("\n=====================")
    print("QUERY GRAPH EXAMPLE 4")
    print("=====================\n")

    # Clear any related tables in case they already exist
    kinetica.clear_table(table_name=TABLE_Q4, options=OPTION_NO_DROP_ERROR)
    kinetica.clear_table(table_name=TABLE_Q4_TARGETS, options=OPTION_NO_DROP_ERROR)

    # Query GRAPH_S for nodes connected to 'females' where the node has an
    # interest in chess
    print(
        "Querying {} for nodes connected to 'females' where the node has an "
        "interest in chess".format(GRAPH_S)
    )
    
    query4_s_graph_response = kinetica.query_graph(
        graph_name = GRAPH_S,
        queries = [
            "{'female'} AS NODE_LABEL",
            "",
            "{'chess'} AS TARGET_NODE_LABEL"
        ],
        restrictions = [],
        adjacency_table = TABLE_Q4,
        rings = 2
    )
    if query4_s_graph_response["status_info"]["status"] == "OK":
        print("{} graph queried successfully.".format(GRAPH_S))
        
        # Pretty print query results and targets
        print("\nQuery results for adjacency table {}:".format(TABLE_Q4))
        tabulate_records(kinetica, TABLE_Q4, ["PATH_ID", "RING_ID"])

        print("\nQuery results for target nodes table {}:".format(TABLE_Q4_TARGETS))
        tabulate_records(kinetica, TABLE_Q4_TARGETS, ["QUERY_NODE_ID_SOURCE", "QUERY_NODE_ID_TARGET"])
    else:
        print("{} graph query failure: \n\t{}".format(
            GRAPH_S,
            query4_s_graph_response["status_info"]["message"]
        ))
        print("Exiting...")
        sys.exit()
    print("")

# end query_graph_example()


def tabulate_records(database_object, table_name, sort_columns = None):
    
    # Get records JSON
    get_records_response = database_object.get_records(
        table_name = table_name,
        encoding = "json"
    )["records_json"]

    if not get_records_response:
        print("No adjacencies or targets found.")
    else:
        # For each record in the response, append the record's key-value pairs to 
        # a list 'rows' for printing. 
        rows = []
        for record in range(0, len(get_records_response)):
            record_dict = json.loads(get_records_response[record])
            rows.append(record_dict)

        if sort_columns:
            sort_columns.reverse()
            for sort_column in sort_columns:
                rows.sort(key=itemgetter(sort_column))

        # Using tabulate, pretty print the 'rows' list with the keys as headers
        print(tabulate(rows, headers="keys", tablefmt="grid"))

# end tabulate_records()


if __name__ == '__main__':

    # Set up args
    parser = argparse.ArgumentParser(description='Run query social graph examples.')
    parser.add_argument('--url', default='http://127.0.0.1:9191', help='Kinetica URL to run examples against')
    parser.add_argument('--username', default='', help='Username of user to run example with')
    parser.add_argument('--password', default='', help='Password of user')

    args, unknown = parser.parse_known_args()

    # Establish connection with a locally-running instance of Kinetica
    kinetica = gpudb.GPUdb(host = [args.url], username = args.username, password = args.password)

    # Execute defined functions
    table_setup()
    graph_setup()
    query_graph_example()