framework/cg/blender/processing/restruct_hierarchy_by_lcs.py

# -*- coding: utf-8 -*-
# Copyright (C) 2023 Ilia Kurochkin <brothermechanic@gmail.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.1'
import logging
import math

import bpy
from mathutils import Matrix

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

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


# COLLECTIONS NAMIG CONVENTION
parts_col_name = 'Parts'
lcs_col_name = 'LCS'
hierarchy_col_name = 'Hierarchy'
lowpoly_col_name = 'Lowpoly'
# LCS POINT'S SUFFIXES CONVENTION
inlet = '_in'
outlet = '_out'
root = '_root'
# CG ASSETS SUFFIXES CONVENTION
hightpoly = '_hp'
lowpoly = '_lp'
render = '_render'


def retree_by_lcs(lcs_objects, root_lcs):
    ''' Organizing project structure based on LCS. '''
    for lcs in lcs_objects:
        locator = lcs.parent
        if lcs.name.endswith(inlet):
            unparenting(lcs)
            round_transforms(lcs)
            if locator.parent:
                unparenting(locator)
            parenting(lcs, locator)
            parenting(root_lcs, lcs)
    for lcs in lcs_objects:
        if lcs.name.endswith(outlet):
            unparenting(lcs)
            round_transforms(lcs)
            parenting(
                lcs_objects[lcs_objects.index(
                    bpy.data.objects[
                        '{}{}'.format(lcs.name.split(outlet)[0], inlet)])],
                lcs)

    root_lcs.matrix_world = Matrix()
    return lcs_objects


def closest_lcs(lcs_objects):
    ''' Finding closest outlet to inlet LCS. '''
    target_dists = {}
    for target in lcs_objects:
        if target.name.endswith(inlet):
            dists = {}
            for lcs in lcs_objects:
                if lcs.name.endswith(outlet):
                    dist = math.dist(
                        target.matrix_world.translation,
                        lcs.matrix_world.translation)
                    dists[lcs.name] = dist
            min_dist = min(dists.values())
            if min_dist < 0.01:
                min_lcs = [k for k, v in dists.items() if v == min_dist][0]
                target_dists[target.name] = min_lcs
    return target_dists


def lcs_constrainting(lcs_objects, root_lcs):
    ''' Placing inlet right on outlet LCS. '''
    closests = closest_lcs(lcs_objects)
    for lcs in lcs_objects:
        if lcs.name in closests:
            constraint = lcs.constraints.new(type='COPY_TRANSFORMS')
            constraint.target = bpy.data.objects[closests[lcs.name]]
        if lcs.name.endswith(outlet):
            constraint = lcs.constraints.new(type='COPY_TRANSFORMS')
            constraint.target = root_lcs
            constraint.enabled = False
    for lcs in lcs_objects:
        if len(lcs.constraints) == 0:
            constraint = lcs.constraints.new(type='COPY_TRANSFORMS')
            constraint.target = root_lcs
            constraint.enabled = False
    return lcs_objects


def unlink_from_col(obj):
    ''' Unlinking object from all collections. '''
    for col in bpy.data.collections:
        if obj.name in col.objects:
            col.objects.unlink(obj)
    return obj


def lcs_collections(root_lcs, lcs_objects):
    ''' Create LCS based hierarchy. '''
    for lcs in root_lcs.children:
        lcs_col = bpy.data.collections.new(
            '{}{}'.format(lcs.name.split(inlet)[0], hightpoly))
        bpy.data.collections[parts_col_name].children.link(lcs_col)
        for obj in lcs.children_recursive:
            unlink_from_col(obj)
            lcs_col.objects.link(obj)
        if lcs not in lcs_objects:
            unlink_from_col(lcs)
            lcs_col.objects.link(lcs)
    return root_lcs.children


def restruct_hierarchy(lcs_names):
    ''' Execute restructurisation. '''

    #lcs_objects = bpy.data.collections[lcs_col_name].objects
    lcs_objects = []
    root_lcs = None
    if lcs_names:
        for obj_name in lcs_names:
            if obj_name.endswith(root):
                root_lcs = bpy.data.objects[obj_name]
            lcs_objects.append(bpy.data.objects[obj_name])

        main_locators = [obj for obj in bpy.data.objects if not obj.parent]
        if len(main_locators) > 1:
            logger.info('Scene has several main (root) locators! '
                        'This may cause an error!')

        if root_lcs:
            lcs_objects = [lcs for lcs in lcs_objects if lcs != root_lcs]
            root_locator = root_lcs.parent
            unparenting(root_lcs)
            round_transforms(root_lcs)
            unparenting(root_locator)
            parenting(root_lcs, root_locator)
            parenting(root_lcs, main_locators[0])

            retree_by_lcs(lcs_objects, root_lcs)
            lcs_constrainting(lcs_objects, root_lcs)

            lcs_collections(root_lcs, lcs_objects)

            # remove unused for now collection
            bpy.data.collections.remove(bpy.data.collections[hierarchy_col_name])

            return logger.info('Restructuring pipeline finished!')
        else:
            return logger.info('Lost root LCS object!')
    else:
        return logger.info('Restructuring pipeline canceled!')