from graphviz import Digraph

# 獲取所有節點中最多子節點的葉節點
def getMaxLeafs(myTree):
    numLeaf = len(myTree.keys())
    for key, value in myTree.items():
        if isinstance(value, dict):
            sum_numLeaf = getMaxLeafs(value)
            if sum_numLeaf > numLeaf:
                numLeaf = sum_numLeaf
    return numLeaf

def plot_model(tree, name):
    g = Digraph("G", filename=name, format='png', strict=False)
    first_label = list(tree.keys())[0]
    g.node("0", first_label)
    _sub_plot(g, tree, "0")
    leafs = str(getMaxLeafs(tree) // 10)
    g.attr(rankdir='LR', ranksep=leafs)
    g.view()

root = "0"


def _sub_plot(g, tree, inc):
    global root

    first_label = list(tree.keys())[0]
    ts = tree[first_label]
    for i in ts.keys():
        if isinstance(tree[first_label][i], dict):
            root = str(int(root) + 1)
            g.node(root, list(tree[first_label][i].keys())[0])
            g.edge(inc, root, str(i))
            _sub_plot(g, tree[first_label][i], root)
        else:
            root = str(int(root) + 1)
            g.node(root, tree[first_label][i])
            g.edge(inc, root, str(i))

tree = {
        "tearRate": {
            "reduced": "no lenses",
            "normal": {
                "astigmatic": {
                    "yes": {
                        "prescript": {
                            "myope": "hard",
                            "hyper": {
                                "age": {
                                    "young": "hard",
                                    "presbyopic": "no lenses",
                                    "pre": "no lenses"
                                }
                            }
                        }
                    },
                    "no": {
                        "age": {
                            "young": "soft",
                            "presbyopic": {
                                "prescript": {
                                    "myope": "no lenses",
                                    "hyper": "soft"
                                }
                            },
                            "pre": "soft"
                        }
                    }
                }
            }
        }
    }
plot_model(tree, "tree.gv")

# 生成樹結構
def get_trees(data,
              key_column='elementId',
              parent_column='parentId',
              child_column='children'):
    """
    :param data: 數據列表
    :param key_column: 主鍵字段，默認id
    :param parent_column: 父ID字段名，父ID默認從0開始
    :param child_column: 子列表字典名稱
    :return: 樹結構
    """
    data_dic = {}
    for d in data:
        data_dic[d.get(key_column)] = d  # 以自己的權限主鍵為鍵,以新構建的字典為值,構造新的字典
 
    data_tree_list = []  # 整個數據大列表
    for d_id, d_dic in data_dic.items():
        pid = d_dic.get(parent_column)  # 取每一個字典中的父id
        if not pid:  # 父id=0，就直接加入數據大列表
            data_tree_list.append(d_dic)
        else:  # 父id>0 就加入父id隊對應的那個的節點列表
            try:  # 判斷異常代表有子節點，增加子節點列表=[]
                data_dic[pid][child_column].append(d_dic)
            except KeyError:
                data_dic[pid][child_column] = []
                data_dic[pid][child_column].append(d_dic)
    return data_tree_list
 
def recursion(data, l=None):
    if l is None:
        l = []
    for i in data:
        if 'children' in i:
            children=i.pop('children')
            l.append(i)
            recursion(children,l)
        else:
            l.append(i)
    return l