# Takes an AttributedGraph Object  as input and creates BERT input by generating
# context per node as a list[edges]
import json
import os
import random

import networkx as nx

"""
Given a graph G, where each vertex has a set of attributes 'a'
(belonging to set 'A' containing all attributes for graph G)
and {v[i], v[j], r} denotes a relation 'r' between vertex v[i] and v[j];

Generate , for each node v[i]:

    The multi-set sample, that captures node structure properties :
        Let N[i] be the neighbourhood of v[i] denotes as
        (r[j], v[j])
        Given p and q be the parameters denoting sample selection strategy (as
        described in node2vec)"
        Input = (G, A, v[i], p, q, k)
        Output = (Sample(N[i]), max_nbrs = k)
"""


class ContextGenerator:
    """
    Context Subgraph: includes a set of nodes selected with certain criteria (e.g., k-hop neighbors for v
    or k-hop neighbors connecting (u,v)) along with their related edges
    """
    def __init__(
        self,
        attributed_graph_obj,
        num_walks_per_node,  # walks per node
        pretrained_embeddings=None,
        optimal_context_per_relation=None,
    ):
        self.attr_graph = attributed_graph_obj
        self.num_walks_per_node = num_walks_per_node  # default=1
        self.pretrained_embeddings = pretrained_embeddings
        random.seed(10)

    def get_random_k_nbrs(self, nodeid, exclude_list, k):
        """
        returns randomly sampled 'k' nbrs for a given node as
        list[node_ids]
        if k == -1 :  returns all neighbours
        """
        all_nbrs = list(set(list(nx.all_neighbors(self.attr_graph.G, nodeid))))
        valid_nbrs = all_nbrs
        if exclude_list is None or len(exclude_list) == 0:
            valid_nbrs = all_nbrs
        else:
            valid_nbrs = [x for x in all_nbrs if x not in exclude_list]
        if k >= 0 and len(valid_nbrs) >= k:
            return random.sample(valid_nbrs, k)
        else:
            return all_nbrs

    def gen_random_walk_context(
        self, beam_width, max_num_edges, walk_type, valid_patterns=None
    ):  # FIXME - removed dangerous valid_patterns list init
        """
        Given as input an attributed networkx graph, generates
        random walk based node_context for  each node.

        Returns : list[node_context] == list[list[(source, dest, relation)]]
        Node context is generated by performing random walk around the
        node with beam_width = beam_width and depth = 2
        """
        G = self.attr_graph.get_graph()
        all_contexts = []
        for source in list(G.nodes()):
            # for u,v in nx.bfs_beam_edges(G, source, centrality.get, beam_width):
            # for u,v in nx.bfs_beam_edges(G, source, random_neighbour_goodness, beam_width):
            #    node_context.append((u,v, G.edges[u,v]['label']))
            for _ in range(self.num_walks_per_node):
                if walk_type == "bfs":
                    node_context = self.bfs_beam(source, beam_width)
                elif walk_type == "dfs":
                    node_context = self.random_chain(source, max_num_edges)
                else:
                    print("Unknown context generation strategy, select bfs/dfs")
                all_contexts.append(node_context)
        # valid_patterns = set(['1_1_1', '2_2_2', '1_2_1', '2_1_2'])
        # all_contexts = self.filter_walks_by_patterns(all_contexts, valid_patterns)
        return all_contexts

    def filter_walks_by_patterns(self, walks, valid_patterns):
        # valid_walks = [walk for walk in walks  if self.get_pattern(walk) in
        #               valid_patterns]
        valid_walks = []
        for walk in walks:
            walk_pattern = self.get_pattern(walk)
            if walk_pattern in valid_patterns:
                valid_walks.append(walk)
        return valid_walks

    def get_pattern(self, walk, first_last_same=False):
        """
        given a path = list[(source, target, relation)]
        returns a pattern = [source[0].type, source[1].type, source[2].type...]
        """
        pattern = []
        for edge in walk:
            # source = edge[0]
            relation = edge[2]
            # source_type = self.attr_graph.get_node_type(source)
            # pattern.append(source_type)
            pattern.append(relation)
        return "_".join(pattern)
        # return pattern

    def random_chain(self, source, max_num_edges):
        # generates a random walk of length=max_num_edges, without cycles
        # starting from source
        path = []
        nodes_so_far = []
        G = self.attr_graph.get_graph()
        for _ in range(max_num_edges):
            nbr = self.get_random_k_nbrs(source, exclude_list=nodes_so_far, k=1)[0]
            try:
                relation = G.edges[source, nbr]["label"]
            except KeyError:  # FIXME - replaced bare except
                relation = G.edges[nbr, source]["label"]
            path.append((source, nbr, relation))
            nodes_so_far.append(source)
            source = nbr
        return path

    def bfs_beam(self, source, beam_width):
        subgraph = list()
        G = self.attr_graph.get_graph()
        source_nbrs = self.get_random_k_nbrs(source, exclude_list=[], k=beam_width)
        for x in source_nbrs:
            try:
                relation = G.edges[source, x]["label"]
            except KeyError:  # FIXME - replaced bare except
                relation = G.edges[x, source]["label"]
            subgraph.append((source, x, relation))
        # for src_nbr in source_nbrs:
        #     src_nbr_hop2_cands = self.get_random_k_nbrs(
        #         src_nbr, exclude_list=[source], k=beam_width
        #     )
        #     for x in src_nbr_hop2_cands:
        #         try:
        #             relation = G.edges[src_nbr, x]["label"]
        #         except KeyError:  # FIXME - replaced bare except
        #             relation = G.edges[x, src_nbr]["label"]
        #         subgraph.append((src_nbr, x, relation))
        # print(len(subgraph))
        return subgraph

    def get_pretrain_subgraphs(
        self,
        data_path,
        data_name,
        beam_width,
        max_num_edges,
        walk_type,
        mode="pretrain",
    ):
        walk_file = os.path.join(data_path, data_name + "_walks.txt")
        if os.path.exists(walk_file):
            print("\n load walks ...")
            fin = open(walk_file, "r")
            all_walks = []
            for line in fin:
                line = json.loads(line)
                all_walks.append(line)
            fin.close()
        else:
            print("\n generate walks ...")
            # context_gen = ContextGenerator(attr_graph)
            # all_walks = self.gen_random_walk_context(beam_width)
            # for pretraining should we geerate walks without taking out test
            # and validation
            all_walks = self.gen_random_walk_context(
                beam_width, max_num_edges, walk_type
            )
            walk_file = os.path.join(data_path, data_name + "_walks.txt")
            fout = open(walk_file, "w")
            for walk in all_walks:
                json.dump(walk, fout)
                fout.write("\n")
            fout.close()
        print("no. of walks", len(all_walks))
        return all_walks

    def get_finetune_subgraphs(
        self, train_edges, valid_edges, test_edges, beam_width, max_seq_len, option
    ):
        """
        1) take 'test_edges+valid_edges' edges out from graph G
        2) Generate random context for all (u,v) belonging to train_edges(),
        valid_edges, test_edges
        3) where each random context is a list[(u,v,relation, 0/1) , (u1, v1, r1), (u2,v2,r2)...]
        Note : The first edge corresponds to a false or true edge and has 4 values in
        tuple. This edge is same as provided by the (u,v,r, 0/1) for the input
        train/test/valid set
        The rest of edges are the context, added here .
        Its formatted to be(source, target, relation) and contains true edges in graph
        Note: random context can be a bfs or dfs
        """
        print("getting finetune subgraphs...")
        # print("size of train, valid, test set", len(train_edges),
        #      len(valid_edges), len(test_edges))
        true_test_edges = [x for x in test_edges if x[3] == "1"]
        true_valid_edges = [x for x in valid_edges if x[3] == "1"]
        to_be_removed_edges = true_test_edges + true_valid_edges
        to_be_removed_source_target_pairs = [(x[1], x[2]) for x in to_be_removed_edges]
        self.attr_graph.G.remove_edges_from(to_be_removed_source_target_pairs)
        walks_by_task_dict = dict()
        for task, task_edges in zip(
            ["train", "valid", "test"], [train_edges, valid_edges, test_edges]
        ):
            task_subgraphs = []
            for edge in task_edges:
                source = edge[1]
                target = edge[2]
                relation = edge[0]
                is_true = edge[3]
                # source_target_contexts = self.get_path_with_edge_label(source, target,
                #                                      max_seq_len-1)
                source_target_contexts = self.get_selective_context(
                    source, target, beam_width, max_seq_len - 1, option
                )
                if len(source_target_contexts) == 0:
                    first_edge = (source, target, relation, is_true)
                    task_subgraphs.append([first_edge])
                else:
                    for subgraph in source_target_contexts:
                        first_edge = (source, target, relation, is_true)
                        subgraph.insert(0, first_edge)
                        task_subgraphs.append(subgraph)
                walks_by_task_dict[task] = task_subgraphs
        # print(walks_by_task_dict)
        # add edges back to graph, in case attributed graph is used by rest of
        # code
        return walks_by_task_dict

    def get_context(self, source, target, beam_width, max_seq_len):
        all_contexts = []
        # for i in range(self.num_walks_per_node):
        source_context = self.bfs_beam(source, beam_width)[0:max_seq_len]
        # target_context = self.bfs_beam(target, beam_width)[0:max_seq_len]
        all_contexts.append(source_context)
        # all_contexts.append(target_context)
        return all_contexts

    def get_selective_context(
        self, source, target, beam_width, max_seq_len, option, valid_patterns=None
    ):  # FIXME - removed dangerous default list init in valid_patterns
        """
        option = 1) all  2) pattern  3)shortest  4)random
        1) Return all paths connecting source and target of length <=max_seq_len
        2) Return only paths that match given pattern set
            #valid_patterns = set(['1_1_1', '2_2_2', '1_2_1', '2_1_2'])
            (1.1: generates paths connecting source and target of length <= max_seq_len
             1.2 filter paths by pattern
        3) Return shortest path between source and target
        4) Return random path between source and target
        """
        G = self.attr_graph.G
        paths = []
        if option == "shortest":
            try:
                path_node_list = nx.bidirectional_shortest_path(G, source, target)
                path_with_edge_label = []
                for i in range(len(path_node_list) - 1):
                    u = path_node_list[i]
                    v = path_node_list[i + 1]
                    path_with_edge_label.append((u, v, G.edges[u, v]["label"]))
                paths = [path_with_edge_label[0:max_seq_len]]
            except nx.exception.NetworkXNoPath:
                return []

        elif option == "all":
            paths = self.get_path_with_edge_label(source, target, max_seq_len)
            # print("found all paths", source, target, paths)
        elif option == "pattern":
            paths = self.get_path_with_edge_label(source, target, max_seq_len)
            if len(valid_patterns) > 0:
                paths = self.filter_walks_by_patterns(paths, valid_patterns)
            print("found pattern paths", source, target, paths)
        elif option == "random":
            node_list_paths = self.get_random_paths_between_nodes(
                source, target, beam_width, max_seq_len, currpath=[]
            )
            paths = []
            for node_list_path in node_list_paths:
                path_with_edge_label = []
                node_list_path.insert(0, source)
                for i in range(len(node_list_path) - 1):
                    u = node_list_path[i]
                    v = node_list_path[i + 1]
                    path_with_edge_label.append((u, v, G.edges[u, v]["label"]))
                paths.append(path_with_edge_label)
        elif option == "default":
            paths = self.get_context(source, target, beam_width, max_seq_len)
        return paths

    def get_random_paths_between_nodes(
        self, source, target, beam_width, max_seq_len, currpath
    ):
        """
        Generates random path between source and target , by selecting
        'beam_width' number of neighbors and upto length >= max_seq_len
        """
        if len(currpath) > max_seq_len or source == target:
            return []
        all_paths = []
        exclude_list = []

        # print("In generate paths", source, target, currpath)

        source_nbrs = self.get_random_k_nbrs(source, exclude_list, beam_width)
        # print("selected source nbrs", source_nbrs)
        if target in source_nbrs:
            # print("Found path ending in target", source_nbrs, target)
            path = [target]
            all_paths.append(path)

        for n in source_nbrs:
            if n != target and n not in currpath:
                new_source = n
                new_path = list(currpath)
                new_path.append(new_source)
                nbr_paths = self.get_random_paths_between_nodes(
                    new_source, target, beam_width, max_seq_len, new_path
                )
                new_path = new_path[0:-1]
                for p in nbr_paths:
                    p.insert(0, new_source)
                all_paths.extend(nbr_paths)
        return all_paths

    def get_path_metrics(self, G, source, target, paths):
        """
        Metrics per node pair:
            1) Number of paths connecting the nodes
            2) Number of metapaths
            5) Polysemy behavior : Variation in number of unique relations each node participates in
            6)
        """
        return

    def get_path_with_edge_label(self, source, target, max_seq_len):
        G = self.attr_graph.G
        paths = nx.all_simple_paths(G, source, target, cutoff=max_seq_len)
        paths_with_edge_labels = []
        for path in map(nx.utils.pairwise, paths):
            path_with_edge_label = []
            for u, v in path:
                path_with_edge_label.append((u, v, G.edges[u, v]["label"]))
            paths_with_edge_labels.append(path_with_edge_label)
        return paths_with_edge_labels
