gp_render/standalone_scripts/rn_gen_tree.py

info = {
    'icon': 'SHADERFX',
    'description': 'create GP render nodes',
}

import fnmatch
import glob
import os
import re
from math import degrees, radians
from os import listdir
from os.path import basename, dirname, exists, isdir, isfile, join, splitext
from pathlib import Path
from mathutils import Vector
import bpy

from collections import defaultdict
C = bpy.context
D = bpy.data

def real_loc(n):
    if not n.parent:
        return n.location
    return n.location + real_loc(n.parent)

def get_frame_transform(f, node_tree):
    '''Return real transform location of a frame node
    only works with one level of nesting (not recursive)
    '''

    if f.type != 'FRAME':
        return
        # return real_loc(f), f.dimensions

    childs = [n for n in node_tree.nodes if n.parent == f]
    # real_locs = [f.location + n.location for n in childs]
    
    xs = [n.location.x for n in childs] + [n.location.x + n.dimensions.x for n in childs]
    ys = [n.location.y for n in childs] + [n.location.y - n.dimensions.y for n in childs]
    xs.sort(key=lambda loc: loc) # x val : ascending 
    ys.sort(key=lambda loc: loc) # ascending # , reversed=True) # y val : descending
    
    loc = Vector((min(xs), max(ys)))
    dim = Vector((max(xs) - min(xs) + 60, max(ys) - min(ys) + 60))
    
    return loc, dim


def bbox(f, frames):
    xs=[]
    ys=[]
    for n in frames[f]: # nodes of passed frame
        # Better as Vectors ?
        if n.type == 'FRAME':
            if n not in frames.keys():
                # print(f'frame {n.name} not in frame list')
                continue
            all_xs, all_ys = bbox(n, frames) # frames[n]
            xs += all_xs
            ys += all_ys

        else:
            loc = real_loc(n)
            xs += [loc.x, loc.x + n.dimensions.x] # + (n.dimensions.x/get_dpi_factor())
            ys += [loc.y, loc.y - n.dimensions.y] # - (n.dimensions.y/get_dpi_factor())


    # margin ~= 30 
    # return xs and ys
    return [min(xs)-30, max(xs)+30], [min(ys)-30, max(ys)+30]

def get_frames_bbox(node_tree):
    '''Return a dic with all frames
    ex: {frame_node: (location, dimension), ...}
    '''

    # create dic of frame object with his direct child nodes nodes
    frames = defaultdict(list)
    frames_bbox = {}
    for n in node_tree.nodes:
        if not n.parent:
            continue
        # also contains frames
        frames[n.parent].append(n)
    
    # Dic for bbox coord
    for f, nodes in frames.items():
        if f.parent:
            continue
        
        xs, ys = bbox(f, frames)
        # xs, ys = bbox(nodes, frames)
        
        ## returning: list of corner coords
        # coords = [
        #     Vector((xs[0], ys[1])),
        #     Vector((xs[1], ys[1])),
        #     Vector((xs[1], ys[0])),
        #     Vector((xs[0], ys[0])),
        # ]
        # frames_bbox[f] = coords

        ## returning: (loc vector, dimensions vector)
        frames_bbox[f] = Vector((xs[0], ys[1])), Vector((xs[1] - xs[0], ys[1] - ys[0]))

    return frames_bbox

def create_node(type, tree=None, **kargs):
    tree = tree or bpy.context.scene.node_tree

    node = tree.nodes.new(type)
    for k,v in kargs.items():
        setattr(node, k, v)
    
    return node

def new_aa_node(tree):
    '''create AA node'''
    aa = create_node('CompositorNodeAntiAliasing', tree) # type = ANTIALIASING
    aa.threshold = 0.5
    aa.contrast_limit = 0.5
    aa.corner_rounding = 0.25
    aa.hide = True
    return aa

def get_render_scene():
    render = bpy.data.scenes.get('Render')
    if not render:
        render = bpy.data.scenes.new('Render')
    render.use_nodes = True
    return render

def set_settings(scene=None):
    if not scene:
        scene = bpy.context.scene
    # specify scene settings for these kind of render
    scene = bpy.context.scene
    scene.eevee.taa_render_samples = 1
    scene.grease_pencil_settings.antialias_threshold = 0

def get_view_layer(name, scene=None):
    '''get viewlayer name
    return existing/created viewlayer
    '''
    if not scene:
        # scene = bpy.context.scene
        scene = get_render_scene()
    ### pass double letter prefix as suffix
    ## pass_name = re.sub(r'^([A-Z]{2})(_)(.*)', r'\3\2\1', 'name')
    ## pass_name = f'{name}_{passe}'
    pass_vl = scene.view_layers.get(name)
    if not pass_vl:
        pass_vl = scene.view_layers.new(name)
    return pass_vl

def add_rlayer(layer_name, scene=None, location=None, color=None, node_name=None, width=400):
    '''create a render layer node
    if node_name is not specified, use passed layer name
    '''

    if not node_name:
        node_name = layer_name # 'RL_' +

    if not scene:
        scene=bpy.context.scene 
    
    nodes = scene.node_tree.nodes
    
    comp = nodes.get(node_name)
    if comp:
        if comp.layer == node_name:
            return comp
        else:
            # TODO : delete rlayer with bad VL name !
            pass

    comp = nodes.new('CompositorNodeRLayers')
    comp.name = node_name
    comp.scene = scene
    comp.layer = layer_name
    comp.label = layer_name
    if location:
        comp.location = location
    if color:
        comp.color = color
    
    if width:
        comp.width = width
    comp.show_preview = False
    return comp

def clear_nodegroup(name, full_clear=False):
    '''remove duplication of a nodegroup (.???)
    also remove the base one if full_clear True
    '''
    for ng in reversed(bpy.data.node_groups):
        pattern = name + r'\.\d{3}'
        if re.search(pattern, ng.name):
            bpy.data.node_groups.remove(ng)

        if full_clear and ng.name == name:
            # if full clear
            bpy.data.node_groups.remove(ng)

def rearrange_frames(node_tree):
    print('> re-arrange node_tree')
    frame_d = get_frames_bbox(node_tree) # dic : {frame_node:(loc vector, dimensions vector), ...}
    if not frame_d:
        return
    ## order the dict by frame.y location
    frame_d = {key: value for key, value in sorted(frame_d.items(), key=lambda pair: pair[1][0].y - pair[1][1].y, reversed=True)}
    frame_d = {key: value for key, value in sorted(frame_d.items(), key=lambda pair: pair[1][0].y - pair[1][1].y, reversed=True)}

    frames = [[f, v[0], v[1].y] for f, v in frame_d.items()] # [frame_node, real_loc, real dimensions]
    # frames.sort(key=lambda n: n.location.y - n.dimensions.y, reverse=True)
    # top = frames[0].location.y
    top = frames[0][1].y # upper node location.y
    offset = 0
    for f in frames:
        # n.location.y = top - offset
        f[0].location.y = (f[1].y - f[0].location.y) + top - offset
        offset += f[2] + 50 # gap # f[0].dimensions.y


def reorder_inputs(ng):
    rl_nodes = [s.links[0].from_node for s in ng.inputs if s.is_linked and s.links and s.links[0].from_node.type == 'R_LAYERS']
    rl_nodes.sort(key=lambda x: x.location.y, reverse=True)
    names = [n.layer for n in rl_nodes]
    inputs_names = [s.name for s in ng.inputs]
    filtered_names = [n for n in names if n in inputs_names]

    for dest, name in enumerate(filtered_names):
        ## rebuild list at each iteration so index are good
        inputs_names = [s.name for s in ng.inputs]
        src = inputs_names.index(name)
        # reorder on node_tree not directly on node!
        ng.node_tree.inputs.move(src, dest)

def reorder_outputs(ng):
    ordered_out_name = [nis.name for nis in ng.inputs if nis.name in [o.name for o in ng.outputs]]
    for s_name in ordered_out_name:
        all_outnames = [o.name for o in ng.outputs]
        # reorder on nodetree, not on node !
        ng.node_tree.outputs.move(all_outnames.index(s_name), ordered_out_name.index(s_name))

def clear_disconnected(fo):
    for inp in reversed(fo.inputs):
        if not inp.is_linked:
            print(f'Deleting unlinked fileout slot: {inp.name}')
            fo.inputs.remove(inp)

def reorder_fileout(fo, ng=None):
    if not ng: # get connected nodegroup
        for s in fo.inputs:
            if s.is_linked and s.links and s.links[0].from_node.type == 'GROUP':
                ng = s.links[0].from_node
                break
    if not ng:
        print(f'No nodegroup to refer to filter {fo.name}')
        return
    ordered = [o.links[0].to_socket.name for o in ng.outputs if o.is_linked and o.is_linked and o.links[0].to_node == fo]
    for s_name in ordered:
        all_outnames = [s.name for s in fo.inputs] # same as [fs.path for fs in fo.file_slots]
        fo.inputs.move(all_outnames.index(s_name), ordered.index(s_name))

def connect_to_group_output(n):
    for o in n.outputs:
        if o.is_linked:
            if o.links[0].to_node.type == 'GROUP_OUTPUT':
                return o.links[0].to_socket
            val = connect_to_group_output(o.links[0].to_node)
            if val:
                return val
    return False

def connect_to_group_input(n):
    for i in n.inputs:
        if i.is_linked:
            if i.links[0].from_node.type == 'GROUP_INPUT':
                return i.links[0].from_socket
            val = connect_to_group_input(i.links[0].from_node)
            if val:
                return val
    return False

def connect_render_layer(rlayer, ng=None, out=None, frame=None):
    scene = get_render_scene()
    nodes = scene.node_tree.nodes
    links = scene.node_tree.links

    vl_name = rlayer.layer
    if not vl_name or vl_name == 'View Layer':
        print(f'Bad layer for node {rlayer.name}')
    
    if not ' / ' in vl_name:
        print(f'no slash (" / ") separator in vl_name {vl_name}, should be "obj.name / layer_name"')
        return

    obname, lname = vl_name.split(' / ')
    lname = bpy.path.clean_name(lname)
    
    if not frame:
        if rlayer.parent:
            frame=rlayer.parent
        else:
            print(f'render_layer has not parent frame: {rlayer.name}')
            frame=None

    ng_name = f'NG_{obname}' # only object name

    # get set nodegroup from vlayer name

    ## clear nodes groups duplication (.00?)
    clear_nodegroup(ng_name, full_clear=False)

    if not ng:
        ng = nodes.get(ng_name)

    if not ng:
        ngroup = bpy.data.node_groups.get(ng_name)
        if not ngroup:
            # delete and recreate ?
            print(f'create nodegroup {ng_name}')
            ngroup = bpy.data.node_groups.new(ng_name, 'CompositorNodeTree')
            
        ng = create_node('CompositorNodeGroup', tree=scene.node_tree, location=(rlayer.location[0] + 600, rlayer.location[1]), width=400)
        if frame:
            ng.parent= frame
        ng.node_tree = ngroup
        ng.name = ngroup.name

        ng_in = create_node('NodeGroupInput', tree=ngroup, location=(-600,0))
        ng_out = create_node('NodeGroupOutput', tree=ngroup, location=(600,0))
    
    else:
        print(f'found group node {ng.name}')
        ngroup = ng.node_tree
        ng_in = ngroup.nodes.get('Group Input')
        ng_out = ngroup.nodes.get('Group Output')

    # Connect rlayer to nodegroup
    if not rlayer.outputs['Image'].is_linked:
        sockin = ng.inputs.get(vl_name)
        if not sockin:
            print('creating socket', vl_name)
            sockin = ng.inputs.new('NodeSocketColor', vl_name)
            sockin = ng.inputs[-1]
        
        links.new(rlayer.outputs['Image'], sockin)

    ## get nodes from frame
    rl_nodes = [n for n in nodes if n.type == 'R_LAYERS' and n.layer != 'View Layer' and n.parent == frame]
    
    # auto clean : if an input exists but is not linked and name not exists in rlayers of current frame
    for s in reversed(ng.inputs):
        if not s.is_linked: # and not any(x.layer == s.name for x in rl_nodes)
            print(f'removing grp unlinked input {s.name}')
            ng.inputs.remove(s)
    
    ## get nodes from linked NG inputs ??? maybe more clear...
    # rl_nodes = [s.links[0].from_node for s in ng.inputs if s.links and s.links[0].from_node and s.links[0].from_node.type == 'R_LAYERS']
    
    ## reorder
    reorder_inputs(ng)
    
    # CREATE NG outsocket (individual, without taking merge)

    connected = False

    if ng_in.outputs[vl_name].is_linked:
        # check if connect to the other side
        socket = connect_to_group_output(ng_in.outputs[vl_name].links[0].to_node) #if ng_in.outputs[vl_name].links[0].to_node.type == 'ALPHAOVER':
        if socket:
            connected = True
            groupout = ng.outputs.get(socket.name)

    if not connected:
        print('need to connect')
        # add AA and connect
        aa = new_aa_node(ngroup)
        groupout = ng.outputs.get(vl_name)
        if not groupout:
            print('create group out-socket')
            ng.outputs.new('NodeSocketColor', vl_name) # assigning direcly doesn't link well
            groupout = ng.outputs[-1]

        print('ng_out.inputs.get(vl_name): ', ng_out.inputs.get(vl_name))
        # ng_in.outputs[vl_name]
        ngroup.links.new(ng_in.outputs[vl_name], aa.inputs[0]) # node_tree
        ngroup.links.new(aa.outputs[0], ng_out.inputs[vl_name]) # node_tree

    # clean outputs
    for o in reversed(ngroup.outputs):
        if not o.name in [o.name for o in ngroup.inputs]:
            print(f'removing group output {o.name} (name not exists in group inputs)')
            ngroup.outputs.remove(o)

    # reorder output to match inputs
    reorder_outputs(ng)

    # Clear : delete orphan nodes that are not connected from ng_in
    for n in reversed(ngroup.nodes):
        if n.type in ('GROUP_INPUT', 'GROUP_OUTPUT'):
            continue
        if not connect_to_group_input(n) and not connect_to_group_output(n): # is disconnected from both side
            ngroup.nodes.remove(n)
        # TODO clear nodes that are disconnected from input side ?

    if groupout.links and groupout.links[0].to_node.type == 'OUTPUT_FILE':
        # if already connected to outfile just skip cause user might have customised the name
        return
    
    slot_name = f'{lname}/{lname}_'
    out_name = f'OUT_{obname}' # or get output from frame
    if not out:
        out = nodes.get(out_name)
    if not out:
        out = create_node('CompositorNodeOutputFile', tree=scene.node_tree, location=(ng.location[0]+600, ng.location[1]+50), width=600) # color = (0.2,0.3,0.5)
        out.name = out_name
        out.parent = frame
        out.base_path = f'//render/{bpy.path.clean_name(obname)}'

    out_input = out.inputs.get(slot_name)
    if not out_input:
        out.file_slots.new(slot_name)
        out_input = out.inputs[-1] # assigning directly above doesn't link afterwards
        print(f'new filouput entry: {out_input}')

    # link to FileOut
    links.new(groupout, out_input)

    # clean fileout
    clear_disconnected(out)
    reorder_fileout(out, ng=ng)

    return ng, out


def get_set_viewlayer_from_gp(ob, l, scene=None):
    if not scene:
        # scene = bpy.context.scene
        scene = get_render_scene() # create if necessary
    node_tree = scene.node_tree
    nodes = scene.node_tree.nodes

    in_rds = scene.collection.all_objects.get(ob.name)
    if not in_rds:
        scene.collection.objects.link(ob)
    
    # create viewlayer
    vl_name = f'{ob.name} / {l.info}'
    vl = get_view_layer(vl_name, scene=scene)
    vl_name = vl.name
    # affect layer to this vl
    l.viewlayer_render = vl_name

    # check if already exists
    rlayer_list = [n for n in nodes if n.type == 'R_LAYERS' and n.layer == vl_name]

    # get frame object and their contents
    # dict like : {objname : [layer_nodeA, layer_nodeB,...]}
    frame_dic = {f.label: [n for n in nodes if n.type == 'R_LAYERS' and n.parent and n.parent.name == f.name and n.layer != 'View Layer']
        for f in nodes if f.type == 'FRAME'}

    # debug print
    for k,v in frame_dic.items():
        print('-', k)
        for n in v:
            print('---', n.layer)

    if rlayer_list: # rlayer exists
        print(f'{len(rlayer_list)} nodes using {vl_name}')
        
        # affect only the one within an object frame
        framed_rl = [n for n in rlayer_list if n.parent and n.parent.label == ob.name]
        if framed_rl:
            if len(framed_rl) > 1:
                print(f'! More than one nodes using {vl_name} in a frame ({len(framed_rl)}) !')

            # sort top to bottom and take upper node
            framed_rl.sort(key=lambda x:x.location.y, reverse=True)
            cp = framed_rl[0]
            cp.select = True # select so the user see that it existed
            return vl, cp

    # Returned if existed and OK    

    if not ob.name in frame_dic.keys(): # and len(frame_dic[ob.name])
        print(f'\n{ob.name} -> {l.info} (first generation)')
        # frame not exists, add the RL and frame at the very bottom of all render_layers
        # check position of frame type ? all type ?
        all_frames = [n for n in nodes if n.type == 'FRAME']
        # all_rl_x = [n.location.x for n in nodes if n.type == 'R_LAYERS' and n.layer != 'View Layer']
        if all_frames:
            # all_frames.sort(key=lambda x: x.location.y, reverse=True)
            # loc.y -  dim.y
            y_loc = min(get_frame_transform(f)[0].y - get_frame_transform(f)[1].y for f in all_frames)
            loc = (0, y_loc)
        else:
            loc = (0,0)

        print('loc: ', loc)
        # create frame at new rl position
        frame = nodes.new('NodeFrame')
        frame.label = ob.name
        frame.label_size = 50
        frame.location = (loc[0], loc[1] + 20)
    
        cp = add_rlayer(vl_name, scene=scene, location=loc)
        cp.parent = frame

        connect_render_layer(cp, frame=frame)

        """ # Create omega-node group
        ngroup = bpy.data.node_groups.new('NG_' + vl_name, 'CompositorNodeTree')
        ng = create_node('CompositorNodeGroup', tree=scene.node_tree, location=(x_loc + 600, y_loc), width=400)
        ng.parent=frame
        ng.node_tree = ngroup
        ng.name = ngroup.name
        print('ng.node_tree: ', ng.node_tree)

        # add GROUP_INPUT(NodeGroupInput) && GROUP_OUTPUT(NodeGroupOutput)
        ng_in = create_node('NodeGroupInput', tree=ngroup, location=(-600,0))
        ng_out = create_node('NodeGroupOutput', tree=ngroup, location=(600,0))

        # add AA and connect
        aa = new_aa_node(ngroup)
        ngroup.inputs.new('NodeSocketColor', vl_name)
        ngroup.outputs.new('NodeSocketColor', vl_name)
        
        scene.node_tree.links.new(cp.outputs[0], ng.inputs[0])

        ngroup.links.new(ng_in.outputs[0], aa.inputs[0]) # node_tree.
        ngroup.links.new(aa.outputs[0], ng_out.inputs[0]) # node_tree.

        # --- add fileout node
        # CompositorNodeOutputFile OUTPUT_FILE

        out = create_node('CompositorNodeOutputFile', tree=scene.node_tree, location=(x_loc+1200, y_loc+50), width=600) # color = (0.2,0.3,0.5)
        out.name = f'OUT_{vl_name}'
        out.parent=frame
        out.base_path = f'//render/{bpy.path.clean_name(ob.name)}' # TODO hardcoded base path
        out.file_slots[0].path = f'{bpy.path.clean_name(l.info)}/{bpy.path.clean_name(l.info)}_'
        
        scene.node_tree.links.new(ng.outputs[0], out.inputs[0])
        """
        return vl, cp


    print(f'\n {ob.name} -> {l.info} (connect to existing)')
    # ng = nodes.get(f'NG_{vl_name}')
    # if not ng:
    #     print('nodegroup not found') # TODO generate if necessary
    #     return

    # out = nodes.get(f'OUT_{vl_name}')
    # if not out:
    #     print('output not found') # TODO generate if necessary
    #     return

    ## object frame exists: get framing and insert
    cp = add_rlayer(vl_name, scene=scene, location=(0,0))
    if cp.layer != vl_name:
        print(f'problem with {cp}: {cp.layer} != {vl_name}')
        return

    frame = [f for f in nodes if f.type == 'FRAME' and f.label == ob.name][0]
    rl_nodes = frame_dic[frame.label]
    if rl_nodes:
        # get nodes order to insert
        rl_nodes.sort(key=lambda n: real_loc(n).y, reverse=True)
        top_loc = real_loc(rl_nodes[0])
    else:
        top_loc = get_frame_transform(frame[1]) -60
        
    # cp.location = (top_loc[0], top_loc[1] + 100) # temp location to adjust x loc

    # list of layer names in nodes order
    rl_names = [n.layer.split(' / ')[1] for n in rl_nodes] # get True layer name from rl
    # names with the right order WITH the new layer included
    names = [lay.info for lay in ob.data.layers if lay.info in rl_names or lay == l]
    
    rl_nodes.append(cp)
    
    # filter by getting index(layer_name)
    cp.parent = frame
    rl_nodes.sort(key=lambda x : names.index(x.layer.split(' / ')[1])) # Sort True layer name from rl

    offset = 0
    print(f'number of nodes in frame: {len(rl_nodes)}')
    ref_node = rl_nodes[0]
    
    print('ref_node: ', ref_node.name, ref_node.location)
    for n in rl_nodes:
        # set x loc from first node in list (maybe use leftmost ?)
        n.location = (ref_node.location[0], top_loc[1] - offset)
        offset += 180
    
    # reorder render layers nodes within frame

    connect_render_layer(cp, frame=frame)

    # re-arrange all frames (since the offset probably overlapped)
    rearrange_frames(node_tree)

    return vl, cp

# def generate_all_layer(ob):
#     '''Basic layer generation'''
#     bpy.context.scene.use_nodes = True
#     for l in ob.data.layers:
#         if l.hide:
#             continue
#         get_set_viewlayer_from_gp(ob, l)

def generate_full_render_output(ob):
    bpy.context.scene.use_nodes = True
    
    # Create another scene. link the GP colleciton (or selected GP object when) in it
    
    # create the render scene
    rds = get_render_scene()
    scn = bpy.data.scenes.get('Scene')
    if not scn:
        if bpy.context.scene != rds:
            scn = rds
        else:
            # return
            all_scenes = [s for s in bpy.data.scenes if s != rds]
            if not all_scenes:
                print('! there is no default scene !')
                return
            scn = all_scenes[0]
    
    # bpy.context.window.scene = rds # switch to render scene ?

    # Link GP or 2D collections ? or GP by GP (in a dedicated collections ? not necessary)
    # gp_col = bpy.data.collections.get('GP')
    # if gp_col:
    #     rds.collection.children.link(gp_col)
    # two_d = bpy.data.collections.get('2D')
    # if two_d:
    #     rds.collection.children.link(gp_col)

    ## better to link selected objects (or all GP objects) on the fly...
    set_settings(scene=rds)
    
    ## setup world, link a specific world or use the current one
    if not rds.world:
        rds.world = scn.world

    ## put in an "output" collection ?
    # out_col = rds.collections.children.get('output')
    # if not out_col:
    #     out_col = bpy.data.collections.new('output')
    #     rds.collection.children.link(out_col)  

    # Clear this/all object(s) and start from scratch 

    # nodes = rds.node_tree.nodes

    exclude_list = ['MA', 'IN']
    for l in ob.data.layers:
        if any(x + '_' in l.info for x in exclude_list):
            continue
        if l.hide:
            continue
        ## Create associated nodegroup later if needed (same operation as selecting multiple and run)
        vl, cp = get_set_viewlayer_from_gp(ob, l, scene=rds)


    ## if an objects nodes are already there, should create in the same area (and offset the all the unrelated bottom nodes)
    

# def generate_all_objects():

#     ## filter the objects by depth ? -> not reliable since animators use the X-ray sometimes...
#     for o in bpy.context.selected_objects:
#         if o.type != 'GPENCIL':
#             continue
#         if not o.select_get():
#             continue
#         generate_all_layer(o)


# generate_full_render_output(C.object)
get_set_viewlayer_from_gp(C.object, C.object.data.layers.active)