import ast
import os

current_path = os.path.dirname(__file__)


def init_node_map():
    with open(current_path + "/astfile/_ast.py", 'r') as file:
        content = file.read()
    ast.parse(content)
    root = ast.parse(content)
    mp = dict()
    index = 1
    edge_mp = dict()
    edge_index = 1
    for node in root.body:
        if isinstance(node, ast.ClassDef):
            mp[node.name] = index
            index += 1
            for son in node.body:
                if isinstance(son, ast.Assign) and son.targets[0].id == '_fields':
                    for field in son.value.dims:
                        if not field.n in edge_mp:
                            edge_mp[field.n] = edge_index
                            edge_index += 1
    return mp, edge_mp


class NodeType:
    node_map, edge_map = init_node_map()


def write_to_txt():
    with open("/Users/liufan/program/PYTHON/sap2nd/GnnForPrivacyScan/data/node.txt", 'w') as file:
        for (node, t) in NodeType.node_map.items():
            file.write(node + " " + str(t) + "\n")
    with open("/Users/liufan/program/PYTHON/sap2nd/GnnForPrivacyScan/data/edge.txt", 'w') as file2:
        for (node, t) in NodeType.edge_map.items():
            file2.write(node + " " + str(t) + "\n")

write_to_txt()