"""
Module implements a class to handle ASTs based on the Graph Tool library
"""

import graph_tool.all as gt
from graph_tool.draw import graph_draw, radial_tree_layout
from tqdm import tqdm
import pickle


class GT_AST(object):
    def __init__(self, ast_file, ast_string=None):
        """
        :param ast_file: string denotes the rep_contents path to the txt file representing the AST
        which follows the structure:
            ??
            node-type-x id-x
            -----
            node-type-x id-x
            node-type-y id-y
            -----
        That is a list of edges separated by -----
        Due to some artifact the first edge is spurious and can be safely ignored.

        This class has two main structures:
        node_map : mapping AST node-label (node-type-x id-x) to a vertex number
        graph: captures the AST representing each vertex with the corresponding vertex number in `node_map`
                it also maintains the node type through the vertex_property["type"]
        """
        self.node_map = dict()
        self.graph = gt.Graph()
        self.graph.vertex_properties["type"] = self.graph.new_vertex_property("string")
        self.graph.vertex_properties["state"] = self.graph.new_vertex_property("object")
        if ast_string != None:
            ast_string = ast_string.split("\n-----\n")
            ast_string = ast_string[1:-1]
            for double_line in ast_string:
                double_line = double_line.split("\n")
                edge = [0, 0]
                for i in range(2):
                    line = double_line[i]
                    node_label = line[:-1]
                    if node_label not in self.node_map:
                        self.node_map[node_label] = len(self.node_map)
                    edge[i] = self.node_map[node_label]
                self.graph.add_edge(*edge)
        else:
            try:

                with open(ast_file, "r", encoding='utf-8') as f:
                    # skipping the first 3 lines
                    for _ in range(3): next(f)
                    for line in f:
                        edge = [0, 0]
                        for i in range(2):
                            node_label = line[:-1]
                            if node_label not in self.node_map:
                                self.node_map[node_label] = len(self.node_map)
                            edge[i] = self.node_map[node_label]
                            line = next(f)
                        self.graph.add_edge(*edge)
            except UnicodeError:
                with open(ast_file, "r", encoding='utf-16') as f:
                    # skipping the first 3 lines
                    for _ in range(3): next(f)
                    for line in f:
                        edge = [0, 0]
                        for i in range(2):
                            node_label = line[:-1]
                            if node_label not in self.node_map:
                                self.node_map[node_label] = len(self.node_map)
                            edge[i] = self.node_map[node_label]
                            line = next(f)
                        self.graph.add_edge(*edge)
            except EOFError:
                print("EOF Error ", ast_file)

            # except KeyboardInterrupt:
            #     print("Keyboard interrupt!")
            #     exit(1)
            except StopIteration as e:
                print("Stop Iteration exception! " + e + " Exiting...")
                # print(ast_file, "inconsistent file format for reading, please check the file")
                exit(1)

        for k, v in self.node_map.items():
            self.graph.vp.type[v] = k.split()[0]

    def show(self, output_file="ast_viz.pdf"):
        """
        Draw the graph on a pdf
        :param output_file: name of the pdf file (must end with .pdf)
        :return:
        """
        pos = radial_tree_layout(self.graph, self.graph.vertex(0))
        scale = self.graph.num_vertices()

        graph_draw(self.graph, vertex_text=self.graph.vp.type,  # self.graph.vertex_index, #
                   pos=pos, vertex_font_size=scale,
                   output=output_file, output_size=(scale * 200, scale * 200))

    def get_nodes_context(self, node_type_map):
        """
        Parse out feasible node contexts from the AST.
        By context we mean the neighborhood of a node: which a set of nodes that satisfy a certain property with respect
        to the node at hand.
        e.g., a typical context of a node is the immediate children of the node.
        This means for a tree with n vertices (with m < n leaves) we have n - m contexts.
        At the moment, we define the context to be immediate child nodes (similar to Mou 2014) but with li = 1
        :param node_type_map: dict mapping node types to integers, this is useful when learning embeddings of these types.
        :return: a list of lists

        """

        # first let's get the number of nodes of subgraphs rooted at each node, we make use of the special vertex indexing our representation has
        # by traversing in a decreasing order of vertex label, we can compute the number of leaves under each node in one pass
        # this is helpful if we want to compute a special kind of lis

        # num_nodes = {} # table of number of nodes of subgraphs rooted at the given node
        # n = self.graph.num_vertices()
        # for v_idx in reversed(range(n)):
        #     v = self.graph.vertex(v_idx)
        #     if v.out_degree() == 0: # leaf
        #         num_nodes[v_idx] = 1
        #     else:
        #         print(v)
        #         num_nodes[v_idx] = 1
        #         for u in v.out_neighbors():
        #             num_nodes[v_idx] += num_nodes[int(u)]

        # print("ratios:")
        # for v in self.graph.vertices():
        #     if v == 0: continue
        #     print("{}:{}".format(v, num_nodes[v] / (num_nodes[next(v.in_neighbors())] - 1)))

        # dump contexts in the form of:
        # [num_children, parent_type, child_1type, child_2type, ..., l_parent, l_child1, l_child2 ...]
        # where l_s are set to 1 for now
        get_type_idx = lambda _: node_type_map[self.graph.vp['type'][_]]
        contexts = []
        for v in self.graph.vertices():
            n = v.out_degree()  # number of children
            if n != 0:  # not a leaf
                contexts.append([n, get_type_idx(v)])
                for u in v.out_neighbors():
                    contexts[-1].append(get_type_idx(u))
                contexts[-1].extend([1] * (n + 1))
        # print(contexts)
        return contexts

    def get_parents_children_type(self):
        """
        Get the possible child node types for every parent node type
        :return: a dictionary mapping strings to sets
        """
        result_dict = dict()

        for v in self.graph.vertices():
            parent_type = self.graph.vp.type[v]
            if parent_type not in result_dict:
                result_dict[parent_type] = set()
            for u in v.out_neighbors():
                child_type = self.graph.vp.type[u]
                result_dict[parent_type].add(child_type)

        return result_dict

    def get_repeated_children_type(self):
        """
        Get possible node types of repeated children for a parent node type
        :return: a dictionary mapping strings to sets of strings
        """
        result_dict = dict()
        for v in self.graph.vertices():
            parent_type = self.graph.vp.type[v]
            if parent_type not in result_dict:
                result_dict[parent_type] = set()

            children_types = set()
            for u in v.out_neighbors():
                child_type = self.graph.vp.type[u]
                if child_type not in children_types:
                    children_types.add(child_type)
                else:
                    result_dict[parent_type].add(child_type)
        return result_dict

    def get_singular_children_type(self, all_repeated):
        """
        Get the possible configurations with singular children for a parent node type
        The output is a dictionary of node types mapping to a set of tuples which contain
        encountered configurations of singular children. When a contiguous interval of
        repeated children is encountered (one or more repeated children) we insert the
        special symbol "REPEATED".
        Example: {"typeA":{("typeB", "REPEATED", "typeC"), ("REPEATED"),
                  "typeB":{("typeD", "typeB"),("typeC", "typeB")}}}
        :param all_repeated: the dictionary mapping node types to possible
        types of repeated children given the information from ALL ASTs
        :return: a dict mapping strings to sets of tuples of strings
        """
        REP = "REPEATED"
        repeated = all_repeated
        result_dict = dict()
        for v in self.graph.vertices():
            parent_type = self.graph.vp.type[v]
            if parent_type not in result_dict:
                result_dict[parent_type] = set()

            pattern = []
            rep_contents = []
            for u in v.out_neighbors():
                child_type = self.graph.vp.type[u]
                if child_type in repeated[parent_type]:

                    if len(pattern) == 0 or pattern[-1] != REP:
                        pattern.append(REP)
                        rep_contents.append([child_type])
                    elif pattern[-1] == REP:
                        rep_contents[-1].append(child_type)
                        # if child_type != rep_contents[-1][-2]:
                        #     print("WARNING! ", rep_contents[-1])
                else:
                    pattern.append(child_type)
            result_dict[parent_type].add(tuple(pattern))
        return result_dict


if __name__ == "__main__":

    import os

    all_asts = []
    # ast_dir = os.path.join("..", "..", "..", "NFS", "asts", "test_data")
    ast_dir = "data/parsed_obf"
    for ast_file in tqdm(os.listdir(ast_dir)):
        short_name = str(ast_file.split("/")[-1].split(".")[0])
        # if os.path.isfile(os.path.join("/home/sanjas/Desktop/pwsh_ast_sanja-master/src/utils/ps/parsed", name+".txt")):

        ast_file = os.path.join(ast_dir, ast_file)
        gt_ast = GT_AST(ast_file=ast_file)
        # name = "data/pdf_ast_obf/" + name + ".pdf"

        name = "data/pdf_ast_obf/" + short_name + ".pdf"
        gt_ast.show(name)
        ast = gt_ast.graph
        #all_asts.append(ast)
        with open('data/asts_obf/{}.p'.format(short_name), 'wb') as f:
            pickle.dump(ast, f)


    # with open('test_asts.p', 'wb') as f:
    #     pickle.dump(all_asts, f)

    # ast_file = os.path.join("..", "..", "..", "NFS", "asts", "MALICIOUS", "16.txt")
    # gt_ast = GT_AST(ast_file=ast_file)
    # gt_ast.show()
    # print(gt_ast.graph.num_edges())
    # print(gt_ast.graph.num_vertices())
    # for u in gt_ast.graph.vertex(3).out_neighbors():
    #     print(u, gt_ast.graph.vp['type'][u])
    #
    # from helper.py.helper_functions import lines2dict
    # node_type_file= os.path.join("..", "..", "..", "NFS", "mal_dict.txt")
    # node_type_map = lines2dict(node_type_file)
af9d14ed6f5f44eddd1a98ca7c28796193e4f07a
<mask token>


class GT_AST(object):
    <mask token>
    <mask token>
    <mask token>
    <mask token>
    <mask token>

    def get_singular_children_type(self, all_repeated):
        """
        Get the possible configurations with singular children for a parent node type
        The output is a dictionary of node types mapping to a set of tuples which contain
        encountered configurations of singular children. When a contiguous interval of
        repeated children is encountered (one or more repeated children) we insert the
        special symbol "REPEATED".
        Example: {"typeA":{("typeB", "REPEATED", "typeC"), ("REPEATED"),
                  "typeB":{("typeD", "typeB"),("typeC", "typeB")}}}
        :param all_repeated: the dictionary mapping node types to possible
        types of repeated children given the information from ALL ASTs
        :return: a dict mapping strings to sets of tuples of strings
        """
        REP = 'REPEATED'
        repeated = all_repeated
        result_dict = dict()
        for v in self.graph.vertices():
            parent_type = self.graph.vp.type[v]
            if parent_type not in result_dict:
                result_dict[parent_type] = set()
            pattern = []
            rep_contents = []
            for u in v.out_neighbors():
                child_type = self.graph.vp.type[u]
                if child_type in repeated[parent_type]:
                    if len(pattern) == 0 or pattern[-1] != REP:
                        pattern.append(REP)
                        rep_contents.append([child_type])
                    elif pattern[-1] == REP:
                        rep_contents[-1].append(child_type)
                else:
                    pattern.append(child_type)
            result_dict[parent_type].add(tuple(pattern))
        return result_dict


<mask token>

<mask token>


class GT_AST(object):

    def __init__(self, ast_file, ast_string=None):
        """
        :param ast_file: string denotes the rep_contents path to the txt file representing the AST
        which follows the structure:
            ??
            node-type-x id-x
            -----
            node-type-x id-x
            node-type-y id-y
            -----
        That is a list of edges separated by -----
        Due to some artifact the first edge is spurious and can be safely ignored.

        This class has two main structures:
        node_map : mapping AST node-label (node-type-x id-x) to a vertex number
        graph: captures the AST representing each vertex with the corresponding vertex number in `node_map`
                it also maintains the node type through the vertex_property["type"]
        """
        self.node_map = dict()
        self.graph = gt.Graph()
        self.graph.vertex_properties['type'] = self.graph.new_vertex_property(
            'string')
        self.graph.vertex_properties['state'] = self.graph.new_vertex_property(
            'object')
        if ast_string != None:
            ast_string = ast_string.split('\n-----\n')
            ast_string = ast_string[1:-1]
            for double_line in ast_string:
                double_line = double_line.split('\n')
                edge = [0, 0]
                for i in range(2):
                    line = double_line[i]
                    node_label = line[:-1]
                    if node_label not in self.node_map:
                        self.node_map[node_label] = len(self.node_map)
                    edge[i] = self.node_map[node_label]
                self.graph.add_edge(*edge)
        else:
            try:
                with open(ast_file, 'r', encoding='utf-8') as f:
                    for _ in range(3):
                        next(f)
                    for line in f:
                        edge = [0, 0]
                        for i in range(2):
                            node_label = line[:-1]
                            if node_label not in self.node_map:
                                self.node_map[node_label] = len(self.node_map)
                            edge[i] = self.node_map[node_label]
                            line = next(f)
                        self.graph.add_edge(*edge)
            except UnicodeError:
                with open(ast_file, 'r', encoding='utf-16') as f:
                    for _ in range(3):
                        next(f)
                    for line in f:
                        edge = [0, 0]
                        for i in range(2):
                            node_label = line[:-1]
                            if node_label not in self.node_map:
                                self.node_map[node_label] = len(self.node_map)
                            edge[i] = self.node_map[node_label]
                            line = next(f)
                        self.graph.add_edge(*edge)
            except EOFError:
                print('EOF Error ', ast_file)
            except StopIteration as e:
                print('Stop Iteration exception! ' + e + ' Exiting...')
                exit(1)
        for k, v in self.node_map.items():
            self.graph.vp.type[v] = k.split()[0]
    <mask token>

    def get_nodes_context(self, node_type_map):
        """
        Parse out feasible node contexts from the AST.
        By context we mean the neighborhood of a node: which a set of nodes that satisfy a certain property with respect
        to the node at hand.
        e.g., a typical context of a node is the immediate children of the node.
        This means for a tree with n vertices (with m < n leaves) we have n - m contexts.
        At the moment, we define the context to be immediate child nodes (similar to Mou 2014) but with li = 1
        :param node_type_map: dict mapping node types to integers, this is useful when learning embeddings of these types.
        :return: a list of lists

        """
        get_type_idx = lambda _: node_type_map[self.graph.vp['type'][_]]
        contexts = []
        for v in self.graph.vertices():
            n = v.out_degree()
            if n != 0:
                contexts.append([n, get_type_idx(v)])
                for u in v.out_neighbors():
                    contexts[-1].append(get_type_idx(u))
                contexts[-1].extend([1] * (n + 1))
        return contexts

    def get_parents_children_type(self):
        """
        Get the possible child node types for every parent node type
        :return: a dictionary mapping strings to sets
        """
        result_dict = dict()
        for v in self.graph.vertices():
            parent_type = self.graph.vp.type[v]
            if parent_type not in result_dict:
                result_dict[parent_type] = set()
            for u in v.out_neighbors():
                child_type = self.graph.vp.type[u]
                result_dict[parent_type].add(child_type)
        return result_dict

    def get_repeated_children_type(self):
        """
        Get possible node types of repeated children for a parent node type
        :return: a dictionary mapping strings to sets of strings
        """
        result_dict = dict()
        for v in self.graph.vertices():
            parent_type = self.graph.vp.type[v]
            if parent_type not in result_dict:
                result_dict[parent_type] = set()
            children_types = set()
            for u in v.out_neighbors():
                child_type = self.graph.vp.type[u]
                if child_type not in children_types:
                    children_types.add(child_type)
                else:
                    result_dict[parent_type].add(child_type)
        return result_dict

    def get_singular_children_type(self, all_repeated):
        """
        Get the possible configurations with singular children for a parent node type
        The output is a dictionary of node types mapping to a set of tuples which contain
        encountered configurations of singular children. When a contiguous interval of
        repeated children is encountered (one or more repeated children) we insert the
        special symbol "REPEATED".
        Example: {"typeA":{("typeB", "REPEATED", "typeC"), ("REPEATED"),
                  "typeB":{("typeD", "typeB"),("typeC", "typeB")}}}
        :param all_repeated: the dictionary mapping node types to possible
        types of repeated children given the information from ALL ASTs
        :return: a dict mapping strings to sets of tuples of strings
        """
        REP = 'REPEATED'
        repeated = all_repeated
        result_dict = dict()
        for v in self.graph.vertices():
            parent_type = self.graph.vp.type[v]
            if parent_type not in result_dict:
                result_dict[parent_type] = set()
            pattern = []
            rep_contents = []
            for u in v.out_neighbors():
                child_type = self.graph.vp.type[u]
                if child_type in repeated[parent_type]:
                    if len(pattern) == 0 or pattern[-1] != REP:
                        pattern.append(REP)
                        rep_contents.append([child_type])
                    elif pattern[-1] == REP:
                        rep_contents[-1].append(child_type)
                else:
                    pattern.append(child_type)
            result_dict[parent_type].add(tuple(pattern))
        return result_dict


<mask token>

<mask token>


class GT_AST(object):

    def __init__(self, ast_file, ast_string=None):
        """
        :param ast_file: string denotes the rep_contents path to the txt file representing the AST
        which follows the structure:
            ??
            node-type-x id-x
            -----
            node-type-x id-x
            node-type-y id-y
            -----
        That is a list of edges separated by -----
        Due to some artifact the first edge is spurious and can be safely ignored.

        This class has two main structures:
        node_map : mapping AST node-label (node-type-x id-x) to a vertex number
        graph: captures the AST representing each vertex with the corresponding vertex number in `node_map`
                it also maintains the node type through the vertex_property["type"]
        """
        self.node_map = dict()
        self.graph = gt.Graph()
        self.graph.vertex_properties['type'] = self.graph.new_vertex_property(
            'string')
        self.graph.vertex_properties['state'] = self.graph.new_vertex_property(
            'object')
        if ast_string != None:
            ast_string = ast_string.split('\n-----\n')
            ast_string = ast_string[1:-1]
            for double_line in ast_string:
                double_line = double_line.split('\n')
                edge = [0, 0]
                for i in range(2):
                    line = double_line[i]
                    node_label = line[:-1]
                    if node_label not in self.node_map:
                        self.node_map[node_label] = len(self.node_map)
                    edge[i] = self.node_map[node_label]
                self.graph.add_edge(*edge)
        else:
            try:
                with open(ast_file, 'r', encoding='utf-8') as f:
                    for _ in range(3):
                        next(f)
                    for line in f:
                        edge = [0, 0]
                        for i in range(2):
                            node_label = line[:-1]
                            if node_label not in self.node_map:
                                self.node_map[node_label] = len(self.node_map)
                            edge[i] = self.node_map[node_label]
                            line = next(f)
                        self.graph.add_edge(*edge)
            except UnicodeError:
                with open(ast_file, 'r', encoding='utf-16') as f:
                    for _ in range(3):
                        next(f)
                    for line in f:
                        edge = [0, 0]
                        for i in range(2):
                            node_label = line[:-1]
                            if node_label not in self.node_map:
                                self.node_map[node_label] = len(self.node_map)
                            edge[i] = self.node_map[node_label]
                            line = next(f)
                        self.graph.add_edge(*edge)
            except EOFError:
                print('EOF Error ', ast_file)
            except StopIteration as e:
                print('Stop Iteration exception! ' + e + ' Exiting...')
                exit(1)
        for k, v in self.node_map.items():
            self.graph.vp.type[v] = k.split()[0]

    def show(self, output_file='ast_viz.pdf'):
        """
        Draw the graph on a pdf
        :param output_file: name of the pdf file (must end with .pdf)
        :return:
        """
        pos = radial_tree_layout(self.graph, self.graph.vertex(0))
        scale = self.graph.num_vertices()
        graph_draw(self.graph, vertex_text=self.graph.vp.type, pos=pos,
            vertex_font_size=scale, output=output_file, output_size=(scale *
            200, scale * 200))

    def get_nodes_context(self, node_type_map):
        """
        Parse out feasible node contexts from the AST.
        By context we mean the neighborhood of a node: which a set of nodes that satisfy a certain property with respect
        to the node at hand.
        e.g., a typical context of a node is the immediate children of the node.
        This means for a tree with n vertices (with m < n leaves) we have n - m contexts.
        At the moment, we define the context to be immediate child nodes (similar to Mou 2014) but with li = 1
        :param node_type_map: dict mapping node types to integers, this is useful when learning embeddings of these types.
        :return: a list of lists

        """
        get_type_idx = lambda _: node_type_map[self.graph.vp['type'][_]]
        contexts = []
        for v in self.graph.vertices():
            n = v.out_degree()
            if n != 0:
                contexts.append([n, get_type_idx(v)])
                for u in v.out_neighbors():
                    contexts[-1].append(get_type_idx(u))
                contexts[-1].extend([1] * (n + 1))
        return contexts

    def get_parents_children_type(self):
        """
        Get the possible child node types for every parent node type
        :return: a dictionary mapping strings to sets
        """
        result_dict = dict()
        for v in self.graph.vertices():
            parent_type = self.graph.vp.type[v]
            if parent_type not in result_dict:
                result_dict[parent_type] = set()
            for u in v.out_neighbors():
                child_type = self.graph.vp.type[u]
                result_dict[parent_type].add(child_type)
        return result_dict

    def get_repeated_children_type(self):
        """
        Get possible node types of repeated children for a parent node type
        :return: a dictionary mapping strings to sets of strings
        """
        result_dict = dict()
        for v in self.graph.vertices():
            parent_type = self.graph.vp.type[v]
            if parent_type not in result_dict:
                result_dict[parent_type] = set()
            children_types = set()
            for u in v.out_neighbors():
                child_type = self.graph.vp.type[u]
                if child_type not in children_types:
                    children_types.add(child_type)
                else:
                    result_dict[parent_type].add(child_type)
        return result_dict

    def get_singular_children_type(self, all_repeated):
        """
        Get the possible configurations with singular children for a parent node type
        The output is a dictionary of node types mapping to a set of tuples which contain
        encountered configurations of singular children. When a contiguous interval of
        repeated children is encountered (one or more repeated children) we insert the
        special symbol "REPEATED".
        Example: {"typeA":{("typeB", "REPEATED", "typeC"), ("REPEATED"),
                  "typeB":{("typeD", "typeB"),("typeC", "typeB")}}}
        :param all_repeated: the dictionary mapping node types to possible
        types of repeated children given the information from ALL ASTs
        :return: a dict mapping strings to sets of tuples of strings
        """
        REP = 'REPEATED'
        repeated = all_repeated
        result_dict = dict()
        for v in self.graph.vertices():
            parent_type = self.graph.vp.type[v]
            if parent_type not in result_dict:
                result_dict[parent_type] = set()
            pattern = []
            rep_contents = []
            for u in v.out_neighbors():
                child_type = self.graph.vp.type[u]
                if child_type in repeated[parent_type]:
                    if len(pattern) == 0 or pattern[-1] != REP:
                        pattern.append(REP)
                        rep_contents.append([child_type])
                    elif pattern[-1] == REP:
                        rep_contents[-1].append(child_type)
                else:
                    pattern.append(child_type)
            result_dict[parent_type].add(tuple(pattern))
        return result_dict


<mask token>

<mask token>


class GT_AST(object):

    def __init__(self, ast_file, ast_string=None):
        """
        :param ast_file: string denotes the rep_contents path to the txt file representing the AST
        which follows the structure:
            ??
            node-type-x id-x
            -----
            node-type-x id-x
            node-type-y id-y
            -----
        That is a list of edges separated by -----
        Due to some artifact the first edge is spurious and can be safely ignored.

        This class has two main structures:
        node_map : mapping AST node-label (node-type-x id-x) to a vertex number
        graph: captures the AST representing each vertex with the corresponding vertex number in `node_map`
                it also maintains the node type through the vertex_property["type"]
        """
        self.node_map = dict()
        self.graph = gt.Graph()
        self.graph.vertex_properties['type'] = self.graph.new_vertex_property(
            'string')
        self.graph.vertex_properties['state'] = self.graph.new_vertex_property(
            'object')
        if ast_string != None:
            ast_string = ast_string.split('\n-----\n')
            ast_string = ast_string[1:-1]
            for double_line in ast_string:
                double_line = double_line.split('\n')
                edge = [0, 0]
                for i in range(2):
                    line = double_line[i]
                    node_label = line[:-1]
                    if node_label not in self.node_map:
                        self.node_map[node_label] = len(self.node_map)
                    edge[i] = self.node_map[node_label]
                self.graph.add_edge(*edge)
        else:
            try:
                with open(ast_file, 'r', encoding='utf-8') as f:
                    for _ in range(3):
                        next(f)
                    for line in f:
                        edge = [0, 0]
                        for i in range(2):
                            node_label = line[:-1]
                            if node_label not in self.node_map:
                                self.node_map[node_label] = len(self.node_map)
                            edge[i] = self.node_map[node_label]
                            line = next(f)
                        self.graph.add_edge(*edge)
            except UnicodeError:
                with open(ast_file, 'r', encoding='utf-16') as f:
                    for _ in range(3):
                        next(f)
                    for line in f:
                        edge = [0, 0]
                        for i in range(2):
                            node_label = line[:-1]
                            if node_label not in self.node_map:
                                self.node_map[node_label] = len(self.node_map)
                            edge[i] = self.node_map[node_label]
                            line = next(f)
                        self.graph.add_edge(*edge)
            except EOFError:
                print('EOF Error ', ast_file)
            except StopIteration as e:
                print('Stop Iteration exception! ' + e + ' Exiting...')
                exit(1)
        for k, v in self.node_map.items():
            self.graph.vp.type[v] = k.split()[0]

    def show(self, output_file='ast_viz.pdf'):
        """
        Draw the graph on a pdf
        :param output_file: name of the pdf file (must end with .pdf)
        :return:
        """
        pos = radial_tree_layout(self.graph, self.graph.vertex(0))
        scale = self.graph.num_vertices()
        graph_draw(self.graph, vertex_text=self.graph.vp.type, pos=pos,
            vertex_font_size=scale, output=output_file, output_size=(scale *
            200, scale * 200))

    def get_nodes_context(self, node_type_map):
        """
        Parse out feasible node contexts from the AST.
        By context we mean the neighborhood of a node: which a set of nodes that satisfy a certain property with respect
        to the node at hand.
        e.g., a typical context of a node is the immediate children of the node.
        This means for a tree with n vertices (with m < n leaves) we have n - m contexts.
        At the moment, we define the context to be immediate child nodes (similar to Mou 2014) but with li = 1
        :param node_type_map: dict mapping node types to integers, this is useful when learning embeddings of these types.
        :return: a list of lists

        """
        get_type_idx = lambda _: node_type_map[self.graph.vp['type'][_]]
        contexts = []
        for v in self.graph.vertices():
            n = v.out_degree()
            if n != 0:
                contexts.append([n, get_type_idx(v)])
                for u in v.out_neighbors():
                    contexts[-1].append(get_type_idx(u))
                contexts[-1].extend([1] * (n + 1))
        return contexts

    def get_parents_children_type(self):
        """
        Get the possible child node types for every parent node type
        :return: a dictionary mapping strings to sets
        """
        result_dict = dict()
        for v in self.graph.vertices():
            parent_type = self.graph.vp.type[v]
            if parent_type not in result_dict:
                result_dict[parent_type] = set()
            for u in v.out_neighbors():
                child_type = self.graph.vp.type[u]
                result_dict[parent_type].add(child_type)
        return result_dict

    def get_repeated_children_type(self):
        """
        Get possible node types of repeated children for a parent node type
        :return: a dictionary mapping strings to sets of strings
        """
        result_dict = dict()
        for v in self.graph.vertices():
            parent_type = self.graph.vp.type[v]
            if parent_type not in result_dict:
                result_dict[parent_type] = set()
            children_types = set()
            for u in v.out_neighbors():
                child_type = self.graph.vp.type[u]
                if child_type not in children_types:
                    children_types.add(child_type)
                else:
                    result_dict[parent_type].add(child_type)
        return result_dict

    def get_singular_children_type(self, all_repeated):
        """
        Get the possible configurations with singular children for a parent node type
        The output is a dictionary of node types mapping to a set of tuples which contain
        encountered configurations of singular children. When a contiguous interval of
        repeated children is encountered (one or more repeated children) we insert the
        special symbol "REPEATED".
        Example: {"typeA":{("typeB", "REPEATED", "typeC"), ("REPEATED"),
                  "typeB":{("typeD", "typeB"),("typeC", "typeB")}}}
        :param all_repeated: the dictionary mapping node types to possible
        types of repeated children given the information from ALL ASTs
        :return: a dict mapping strings to sets of tuples of strings
        """
        REP = 'REPEATED'
        repeated = all_repeated
        result_dict = dict()
        for v in self.graph.vertices():
            parent_type = self.graph.vp.type[v]
            if parent_type not in result_dict:
                result_dict[parent_type] = set()
            pattern = []
            rep_contents = []
            for u in v.out_neighbors():
                child_type = self.graph.vp.type[u]
                if child_type in repeated[parent_type]:
                    if len(pattern) == 0 or pattern[-1] != REP:
                        pattern.append(REP)
                        rep_contents.append([child_type])
                    elif pattern[-1] == REP:
                        rep_contents[-1].append(child_type)
                else:
                    pattern.append(child_type)
            result_dict[parent_type].add(tuple(pattern))
        return result_dict


if __name__ == '__main__':
    import os
    all_asts = []
    ast_dir = 'data/parsed_obf'
    for ast_file in tqdm(os.listdir(ast_dir)):
        short_name = str(ast_file.split('/')[-1].split('.')[0])
        ast_file = os.path.join(ast_dir, ast_file)
        gt_ast = GT_AST(ast_file=ast_file)
        name = 'data/pdf_ast_obf/' + short_name + '.pdf'
        gt_ast.show(name)
        ast = gt_ast.graph
        with open('data/asts_obf/{}.p'.format(short_name), 'wb') as f:
            pickle.dump(ast, f)

"""
Module implements a class to handle ASTs based on the Graph Tool library
"""

import graph_tool.all as gt
from graph_tool.draw import graph_draw, radial_tree_layout
from tqdm import tqdm
import pickle


class GT_AST(object):
    def __init__(self, ast_file, ast_string=None):
        """
        :param ast_file: string denotes the rep_contents path to the txt file representing the AST
        which follows the structure:
            ??
            node-type-x id-x
            -----
            node-type-x id-x
            node-type-y id-y
            -----
        That is a list of edges separated by -----
        Due to some artifact the first edge is spurious and can be safely ignored.

        This class has two main structures:
        node_map : mapping AST node-label (node-type-x id-x) to a vertex number
        graph: captures the AST representing each vertex with the corresponding vertex number in `node_map`
                it also maintains the node type through the vertex_property["type"]
        """
        self.node_map = dict()
        self.graph = gt.Graph()
        self.graph.vertex_properties["type"] = self.graph.new_vertex_property("string")
        self.graph.vertex_properties["state"] = self.graph.new_vertex_property("object")
        if ast_string != None:
            ast_string = ast_string.split("\n-----\n")
            ast_string = ast_string[1:-1]
            for double_line in ast_string:
                double_line = double_line.split("\n")
                edge = [0, 0]
                for i in range(2):
                    line = double_line[i]
                    node_label = line[:-1]
                    if node_label not in self.node_map:
                        self.node_map[node_label] = len(self.node_map)
                    edge[i] = self.node_map[node_label]
                self.graph.add_edge(*edge)
        else:
            try:

                with open(ast_file, "r", encoding='utf-8') as f:
                    # skipping the first 3 lines
                    for _ in range(3): next(f)
                    for line in f:
                        edge = [0, 0]
                        for i in range(2):
                            node_label = line[:-1]
                            if node_label not in self.node_map:
                                self.node_map[node_label] = len(self.node_map)
                            edge[i] = self.node_map[node_label]
                            line = next(f)
                        self.graph.add_edge(*edge)
            except UnicodeError:
                with open(ast_file, "r", encoding='utf-16') as f:
                    # skipping the first 3 lines
                    for _ in range(3): next(f)
                    for line in f:
                        edge = [0, 0]
                        for i in range(2):
                            node_label = line[:-1]
                            if node_label not in self.node_map:
                                self.node_map[node_label] = len(self.node_map)
                            edge[i] = self.node_map[node_label]
                            line = next(f)
                        self.graph.add_edge(*edge)
            except EOFError:
                print("EOF Error ", ast_file)

            # except KeyboardInterrupt:
            #     print("Keyboard interrupt!")
            #     exit(1)
            except StopIteration as e:
                print("Stop Iteration exception! " + e + " Exiting...")
                # print(ast_file, "inconsistent file format for reading, please check the file")
                exit(1)

        for k, v in self.node_map.items():
            self.graph.vp.type[v] = k.split()[0]

    def show(self, output_file="ast_viz.pdf"):
        """
        Draw the graph on a pdf
        :param output_file: name of the pdf file (must end with .pdf)
        :return:
        """
        pos = radial_tree_layout(self.graph, self.graph.vertex(0))
        scale = self.graph.num_vertices()

        graph_draw(self.graph, vertex_text=self.graph.vp.type,  # self.graph.vertex_index, #
                   pos=pos, vertex_font_size=scale,
                   output=output_file, output_size=(scale * 200, scale * 200))

    def get_nodes_context(self, node_type_map):
        """
        Parse out feasible node contexts from the AST.
        By context we mean the neighborhood of a node: which a set of nodes that satisfy a certain property with respect
        to the node at hand.
        e.g., a typical context of a node is the immediate children of the node.
        This means for a tree with n vertices (with m < n leaves) we have n - m contexts.
        At the moment, we define the context to be immediate child nodes (similar to Mou 2014) but with li = 1
        :param node_type_map: dict mapping node types to integers, this is useful when learning embeddings of these types.
        :return: a list of lists

        """

        # first let's get the number of nodes of subgraphs rooted at each node, we make use of the special vertex indexing our representation has
        # by traversing in a decreasing order of vertex label, we can compute the number of leaves under each node in one pass
        # this is helpful if we want to compute a special kind of lis

        # num_nodes = {} # table of number of nodes of subgraphs rooted at the given node
        # n = self.graph.num_vertices()
        # for v_idx in reversed(range(n)):
        #     v = self.graph.vertex(v_idx)
        #     if v.out_degree() == 0: # leaf
        #         num_nodes[v_idx] = 1
        #     else:
        #         print(v)
        #         num_nodes[v_idx] = 1
        #         for u in v.out_neighbors():
        #             num_nodes[v_idx] += num_nodes[int(u)]

        # print("ratios:")
        # for v in self.graph.vertices():
        #     if v == 0: continue
        #     print("{}:{}".format(v, num_nodes[v] / (num_nodes[next(v.in_neighbors())] - 1)))

        # dump contexts in the form of:
        # [num_children, parent_type, child_1type, child_2type, ..., l_parent, l_child1, l_child2 ...]
        # where l_s are set to 1 for now
        get_type_idx = lambda _: node_type_map[self.graph.vp['type'][_]]
        contexts = []
        for v in self.graph.vertices():
            n = v.out_degree()  # number of children
            if n != 0:  # not a leaf
                contexts.append([n, get_type_idx(v)])
                for u in v.out_neighbors():
                    contexts[-1].append(get_type_idx(u))
                contexts[-1].extend([1] * (n + 1))
        # print(contexts)
        return contexts

    def get_parents_children_type(self):
        """
        Get the possible child node types for every parent node type
        :return: a dictionary mapping strings to sets
        """
        result_dict = dict()

        for v in self.graph.vertices():
            parent_type = self.graph.vp.type[v]
            if parent_type not in result_dict:
                result_dict[parent_type] = set()
            for u in v.out_neighbors():
                child_type = self.graph.vp.type[u]
                result_dict[parent_type].add(child_type)

        return result_dict

    def get_repeated_children_type(self):
        """
        Get possible node types of repeated children for a parent node type
        :return: a dictionary mapping strings to sets of strings
        """
        result_dict = dict()
        for v in self.graph.vertices():
            parent_type = self.graph.vp.type[v]
            if parent_type not in result_dict:
                result_dict[parent_type] = set()

            children_types = set()
            for u in v.out_neighbors():
                child_type = self.graph.vp.type[u]
                if child_type not in children_types:
                    children_types.add(child_type)
                else:
                    result_dict[parent_type].add(child_type)
        return result_dict

    def get_singular_children_type(self, all_repeated):
        """
        Get the possible configurations with singular children for a parent node type
        The output is a dictionary of node types mapping to a set of tuples which contain
        encountered configurations of singular children. When a contiguous interval of
        repeated children is encountered (one or more repeated children) we insert the
        special symbol "REPEATED".
        Example: {"typeA":{("typeB", "REPEATED", "typeC"), ("REPEATED"),
                  "typeB":{("typeD", "typeB"),("typeC", "typeB")}}}
        :param all_repeated: the dictionary mapping node types to possible
        types of repeated children given the information from ALL ASTs
        :return: a dict mapping strings to sets of tuples of strings
        """
        REP = "REPEATED"
        repeated = all_repeated
        result_dict = dict()
        for v in self.graph.vertices():
            parent_type = self.graph.vp.type[v]
            if parent_type not in result_dict:
                result_dict[parent_type] = set()

            pattern = []
            rep_contents = []
            for u in v.out_neighbors():
                child_type = self.graph.vp.type[u]
                if child_type in repeated[parent_type]:

                    if len(pattern) == 0 or pattern[-1] != REP:
                        pattern.append(REP)
                        rep_contents.append([child_type])
                    elif pattern[-1] == REP:
                        rep_contents[-1].append(child_type)
                        # if child_type != rep_contents[-1][-2]:
                        #     print("WARNING! ", rep_contents[-1])
                else:
                    pattern.append(child_type)
            result_dict[parent_type].add(tuple(pattern))
        return result_dict


if __name__ == "__main__":

    import os

    all_asts = []
    # ast_dir = os.path.join("..", "..", "..", "NFS", "asts", "test_data")
    ast_dir = "data/parsed_obf"
    for ast_file in tqdm(os.listdir(ast_dir)):
        short_name = str(ast_file.split("/")[-1].split(".")[0])
        # if os.path.isfile(os.path.join("/home/sanjas/Desktop/pwsh_ast_sanja-master/src/utils/ps/parsed", name+".txt")):

        ast_file = os.path.join(ast_dir, ast_file)
        gt_ast = GT_AST(ast_file=ast_file)
        # name = "data/pdf_ast_obf/" + name + ".pdf"

        name = "data/pdf_ast_obf/" + short_name + ".pdf"
        gt_ast.show(name)
        ast = gt_ast.graph
        #all_asts.append(ast)
        with open('data/asts_obf/{}.p'.format(short_name), 'wb') as f:
            pickle.dump(ast, f)


    # with open('test_asts.p', 'wb') as f:
    #     pickle.dump(all_asts, f)

    # ast_file = os.path.join("..", "..", "..", "NFS", "asts", "MALICIOUS", "16.txt")
    # gt_ast = GT_AST(ast_file=ast_file)
    # gt_ast.show()
    # print(gt_ast.graph.num_edges())
    # print(gt_ast.graph.num_vertices())
    # for u in gt_ast.graph.vertex(3).out_neighbors():
    #     print(u, gt_ast.graph.vp['type'][u])
    #
    # from helper.py.helper_functions import lines2dict
    # node_type_file= os.path.join("..", "..", "..", "NFS", "mal_dict.txt")
    # node_type_map = lines2dict(node_type_file)