framework/cg/blender/processing/restruct_hierarchy_by_lcs.py

# coding: utf-8
# Copyright (C) 2023 Ilia Kurochkin <brothermechanic@yandex.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
'''
DESCRIPTION.
Reorganization and restructuring of assembly structure
based on LCS point objects.
'''
__version__ = '0.2'
import logging

import bpy
import mathutils

from blender.utils.object_relations import (parenting,
                                            unparenting)
from blender.utils.object_transforms import round_transforms
from blender.utils.collection_tools import unlink_from_collections

logger = logging.getLogger(__name__)


def restruct_hierarchy(lcs_names, parts_sequence=None, **cg_config):
    ''' Execute restructurisation. '''

    main_locators = [obj for obj in bpy.data.objects if not obj.parent]
    assert len(main_locators) == 1, (
        'Scene should have only one root parent locator!')

    lcs_inlet_objects = []
    lcs_outlet_objects = []
    for name in lcs_names:
        if name.endswith(cg_config['lcs_inlet']):
            lcs_inlet_objects.append(bpy.data.objects[name])
        if name.endswith(cg_config['lcs_outlet']):
            lcs_outlet_objects.append(bpy.data.objects[name])

    if len(lcs_inlet_objects) > 1:
        assert parts_sequence, (
            'Parts sequence do not assign! Process stopped!')
        for locator in main_locators[0].children:
            assert locator.name in parts_sequence, (
                    'Can not find {} in "parts_sequence" config!'
                    .format(locator.name))
        first_part_obj = bpy.data.objects[parts_sequence[0]]
    elif len(lcs_inlet_objects) == 1:
        first_part_obj = lcs_inlet_objects[0].parent
    else:
        # TODO
        first_part_obj = None
        assert lcs_inlet_objects, (
            'Scene do not contain any inlet lcs! Process stopped!')

    # create root lcs by parts sequence
    root_lcs = None
    for lcs in first_part_obj.children:
        if lcs.name not in lcs_names:
            continue
        if lcs.name.endswith(cg_config['lcs_outlet']):
            continue
        root_lcs_name = cg_config['lcs_root']
        root_lcs = bpy.data.objects.new(root_lcs_name, None)
        root_lcs.empty_display_type = 'ARROWS'
        root_lcs.empty_display_size = 0.15
        root_lcs.show_in_front = True
        root_lcs.location = lcs.location
        root_lcs.rotation_euler = lcs.rotation_euler
        root_lcs.parent = lcs.parent
        bpy.data.collections[cg_config['lcs_col_name']].objects.link(root_lcs)
        logger.info('Root Inlet LCS object created!')

    unparenting(root_lcs)
    round_transforms(root_lcs)
    parenting(root_lcs, main_locators[0])

    # retree_by lcs
    for lcs in lcs_inlet_objects:
        locator = lcs.parent
        unparenting(lcs)
        round_transforms(lcs)
        if locator:
            if locator.parent:
                unparenting(locator)
            parenting(lcs, locator)
        parenting(root_lcs, lcs)
    for lcs in lcs_outlet_objects:
        unparenting(lcs)
        round_transforms(lcs)
        parenting(
            lcs_inlet_objects[lcs_inlet_objects.index(bpy.data.objects[
                '{}_{}'.format(lcs.name.rpartition('_')[0], cg_config['lcs_inlet'])
            ])],
            lcs)
    # reset transforms for root_lcs
    root_lcs.matrix_world = mathutils.Matrix()

    # lcs collections
    part_names = []
    for lcs in root_lcs.children:
        # remove unmarked parts
        if lcs not in lcs_inlet_objects:
            for obj in lcs.children_recursive:
                bpy.data.objects.remove(obj, do_unlink=True)
            bpy.data.objects.remove(lcs, do_unlink=True)
            continue
        # collect part names
        part_name = None
        for locator in lcs.children:
            if locator not in lcs_outlet_objects:
                part_name = locator.name
                part_names.append(part_name)
        # pack parts to collections
        part_col = bpy.data.collections.new('{}_{}'.format(
            part_name, cg_config['hightpoly']))
        bpy.data.collections[cg_config['parts_col_name']].children.link(part_col)
        for obj in lcs.children_recursive:
            # outlet lcs objects are already in place, don't move it
            if obj in lcs_outlet_objects:
                continue
            unlink_from_collections(obj)
            part_col.objects.link(obj)

    # parts assembling TODO clones
    if len(lcs_inlet_objects) > 1:
        for idx, part_name in enumerate(parts_sequence):
            # for clones
            if part_name not in part_names:
                continue

            lcs_in = bpy.data.objects[part_name].parent
            constraint = lcs_in.constraints.new(type='COPY_TRANSFORMS')
            if idx == 0:
                constraint.target = root_lcs
                continue
            # if asm pair exists
            if bpy.data.objects.get(parts_sequence[idx - 1]):
                lcs_target = [
                    lcs_out
                    for lcs_out in bpy.data.objects[parts_sequence[idx - 1]].parent.children
                    if lcs_out in lcs_outlet_objects][0]
                constraint.target = lcs_target
            else:
                constraint.target = root_lcs
                constraint.enabled = False
        # for reseet transforms when exporting
        for lcs in lcs_outlet_objects:
            constraint = lcs.constraints.new(type='COPY_TRANSFORMS')
            constraint.target = root_lcs
            constraint.enabled = False

    logger.info('Restructuring pipeline by LCS finished!')

    return part_names