import json
import argparse
import pickle

parser = argparse.ArgumentParser()
parser.add_argument('--id', type=str, default=None)
parser.add_argument('--name', type=str, default=None)
args = parser.parse_args()

class Node:
    def __init__(self, id_: str, name: str, *, level: int = None):
        self.id_ = id_
        self.name = name
        self.children = []
        self.level = -1
        if level:
            self.level = level

    def add_child(self, elem):
        self.children.append(elem)

    def update_level(self, value: int):
        self.level = value

    def __str__(self):
        return f"[{self.id_} {self.name}]".format(self.id_, self.name)

    __repr__ = __str__

    @staticmethod
    def _all_children_ids(node):
        print(node)
        for item in node.children:
            Node._all_children_ids(item)

    @staticmethod
    def all_children_ids(node):
        Node._all_children_ids(node)

    @staticmethod
    def _all_children_ids_as_list(node, res):
        res.append(node.id_)
        for item in node.children:
            Node._all_children_ids_as_list(item, res)

    @staticmethod
    def all_children_ids_as_list(node):
        res = []
        Node._all_children_ids_as_list(node, res)
        return res



root = Node("null", "null", level=0)

with open("mobilenet.json", "r") as f:
    json_h = json.load(f)

mpp = {}

def build_tree_recursive(node, json_node, level):
    global mpp
    if "children" in json_node:
        for item_ in json_node["children"]:
            child_node = Node(item_["id"], item_["name"], level=level + 1)
            mpp[item_["id"]] = item_["name"]
            build_tree_recursive(child_node, item_, level + 1)
            node.add_child(child_node)

build_tree_recursive(root, json_h, 0)

def find_with_id(node, id):
    if node.id_ == id:
        return node

    for item in node.children:
        elem = find_with_id(item, id)
        if elem:
            return elem
    return None

def find_with_name(node, name):
    if node.name == name:
        return node

    for item in node.children:
        elem = find_with_name(item, name)
        if elem:
            return elem
    return None


if args.id:
    elem = find_with_id(root, args.id)
    print(elem.all_children_ids(elem))
    exit()

if args.name:
    elem = find_with_name(root, args.name)
    list_ = elem.all_children_ids_as_list(elem)
    with open("../ids.pkl", "wb") as f:
        pickle.dump(list_, f)
    exit()

print(root.all_children_ids(root))

with open("mmpap.pkl", "wb") as f:
    pickle.dump(mpp, f)
