mantis/visualize.py

""" Optional file providing a tool to visualize Mantis Graphs, for debugging and development"""

from bpy.types import Node, NodeTree, Operator

from .utilities import (prRed, prGreen, prPurple, prWhite,
                              prOrange,
                              wrapRed, wrapGreen, wrapPurple, wrapWhite,
                              wrapOrange,)

class MantisVisualizeTree(NodeTree):
    '''A custom node tree type that will show up in the editor type list'''
    bl_idname = 'MantisVisualizeTree'
    bl_label = "mantis output"
    bl_icon = 'HIDE_OFF'

class MantisVisualizeNode(Node):
    bl_idname = "MantisVisualizeNode"
    bl_label = "Node"
    @classmethod
    def poll(cls, ntree):
        return (ntree.bl_idname in ['MantisVisualizeTree'])
    
    def init(self, context):
        pass
    
    def gen_data(self, mantis_node, mode='DEBUG_CONNECTIONS'):
        from .utilities import get_node_prototype
        if mantis_node.node_type in ['SCHEMA', 'DUMMY']:
            np=None
        else:
            np=get_node_prototype(mantis_node.ui_signature, mantis_node.base_tree)
        self.use_custom_color = True
        match mantis_node.node_type:
            case 'XFORM':        self.color = (1.0 ,0.5, 0.0)
            case 'LINK':         self.color = (0.4 ,0.2, 1.0)
            case 'UTILITY':      self.color = (0.2 ,0.2, 0.2)
            case 'DRIVER':       self.color = (0.7, 0.05, 0.8)
            case 'DUMMY_SCHEMA': self.color = (0.85 ,0.95, 0.9)
            case 'DUMMY':        self.color = (0.05 ,0.05, 0.15)
        self.name = '.'.join(mantis_node.signature[1:])

        if np:
            if np.label:
                self.label=np.label
            else:
                self.label=np.name
            for inp in mantis_node.inputs:
                match mode:
                    case "DEBUG_CONNECTIONS":
                        if not inp.is_connected:
                            continue
                s = self.inputs.new('WildcardSocket', inp.name)
                try:
                    if sock := np.inputs.get(inp.name):
                        s.color = sock.color_simple
                except AttributeError: #default bl_idname types like Float and Vector, no biggie
                    pass
                except KeyError:
                    pass
            for out in mantis_node.outputs:
                match mode:
                    case "DEBUG_CONNECTIONS":
                        if not out.is_connected:
                            continue

                s = self.outputs.new('WildcardSocket', out.name)
                try:
                    if sock := np.outputs.get(out.name):
                        s.color = sock.color_simple
                except AttributeError: #default bl_idname types like Float and Vector, no biggie
                    pass
                except KeyError:
                    pass
            self.location = np.location # will get overwritten by Grandalf later.
        else:
            self.label = mantis_node.signature[-1] # which is be the unique name.
            for inp in mantis_node.inputs:
                match mode:
                    case "DEBUG_CONNECTIONS":
                        if not inp.is_connected:
                            continue
                self.inputs.new('WildcardSocket', inp.name)
            for out in mantis_node.outputs:
                match mode:
                    case "DEBUG_CONNECTIONS":
                        if not out.is_connected:
                            continue
                self.outputs.new('WildcardSocket', out.name)

def gen_vis_node( mantis_node, 
                  vis_tree, 
                  links,
                  omit_simple=True,
                 ):
    if mantis_node.node_type == 'UTILITY':
        return
    base_tree= mantis_node.base_tree
    vis = vis_tree.nodes.new('MantisVisualizeNode')
    vis.gen_data(mantis_node)
    for inp in mantis_node.inputs.values():
        for l in inp.links:
            if l.from_node in mantis_node.hierarchy_dependencies:
                links.add(l)
    for out in mantis_node.outputs.values():
        for l in out.links:
            if l.to_node in mantis_node.hierarchy_connections:
                links.add(l)
    return vis
                
def visualize_tree(m_nodes, base_tree, context):
    # first create a MantisVisualizeTree
    from .readtree import check_and_add_root
    from .utilities import trace_all_nodes_from_root
    import bpy

    base_tree.is_executing=True
    import cProfile
    import pstats, io
    from pstats import SortKey
    with cProfile.Profile() as pr:
        try:
            trace_all_nodes = True
            all_links = set()
            mantis_nodes=set()
            nodes={}
            if trace_all_nodes:
                roots=[]
                for n in m_nodes.values():
                    check_and_add_root(n, roots)
                for r in roots:
                    trace_all_nodes_from_root(r, mantis_nodes)
                if len(mantis_nodes) ==  0:
                    print ("No nodes to visualize")
                    return
            else:
                mantis_nodes = list(base_tree.parsed_tree.values())

            vis_tree = bpy.data.node_groups.new(base_tree.name+'_visualized', type='MantisVisualizeTree')

            for m in mantis_nodes:
                nodes[m.signature]=gen_vis_node(m, vis_tree,all_links)

            for l in all_links:
                if l.to_node.node_type in ['DUMMY_SCHEMA', 'DUMMY'] or \
                l.from_node.node_type in ['DUMMY_SCHEMA', 'DUMMY']:
                    pass
                from_node=nodes.get(l.from_node.signature)
                to_node=nodes.get(l.to_node.signature)
                from_socket, to_socket = None, None
                if from_node and to_node:
                    from_socket = from_node.outputs.get(l.from_socket)
                    to_socket = to_node.inputs.get(l.to_socket)
                if from_socket and to_socket:
                    try:
                        vis_tree.links.new(
                            input=from_socket,
                            output=to_socket,
                            )
                    except Exception as e:
                        print (type(e)); print(e)
                        raise e
            # at this point not all links are in the tree yet!

            def has_links (n):
                for input in n.inputs:
                    if input.is_linked:
                        return True
                for output in n.outputs:
                    if output.is_linked:
                        return True
                return False
            
            no_links=[]

            for n in vis_tree.nodes:
                if not has_links(n):
                    no_links.append(n)
            
            while (no_links):
                n = no_links.pop()
                vis_tree.nodes.remove(n)
            
            # def side_len(n):
            #     from math import floor
            #     side = floor(n**(1/2)) + 1
            #     return side
            # side=side_len(len(no_links))
            # break_me = True
            # for i in range(side):
            #     for j in range(side):
            #         index = side*i+j
            #         try:
            #             n = no_links[index]
            #             n.location.x = i*200
            #             n.location.y = j*200
            #         except IndexError:
            #             break_me = True # it's too big, that's OK the square is definitely bigger
            #             break
            #     if break_me:
            #         break
            # from .utilities import SugiyamaGraph
            # SugiyamaGraph(vis_tree, 1) # this can take a really long time
        finally:
            s = io.StringIO()
            sortby = SortKey.TIME
            ps = pstats.Stats(pr, stream=s).strip_dirs().sort_stats(sortby)
            ps.print_stats(40) # print the top 20
            print(s.getvalue())
            base_tree.prevent_next_exec=True
            base_tree.is_executing=False


from .ops_nodegroup import mantis_tree_poll_op

class MantisVisualizeOutput(Operator):
    """Mantis Visualize Output Operator"""
    bl_idname = "mantis.visualize_output"
    bl_label = "Visualize Output"

    @classmethod
    def poll(cls, context):
        return (mantis_tree_poll_op(context))

    def execute(self, context):
        from time import time
        from .utilities import wrapGreen, prGreen
        
        tree=context.space_data.path[0].node_tree
        tree.update_tree(context)
        # tree.execute_tree(context)
        prGreen(f"Visualize Tree: {tree.name}")
        nodes = tree.parsed_tree
        visualize_tree(nodes, tree, context)
        return {"FINISHED"}