path: root/rigify/utils/node_merger.py

# SPDX-License-Identifier: GPL-2.0-or-later

import bpy
import collections
import heapq
import operator

from mathutils import Vector
from mathutils.kdtree import KDTree

from .errors import MetarigError
from ..base_rig import stage, GenerateCallbackHost
from ..base_generate import GeneratorPlugin


class NodeMerger(GeneratorPlugin):
    """
    Utility that allows rigs to interact based on common points in space.

    Rigs can register node objects representing locations during the
    initialize stage, and at the end the plugin sorts them into buckets
    based on proximity. For each such bucket a group object is created
    and allowed to further process the nodes.

    Nodes chosen by the groups as being 'final' become sub-objects of
    the plugin and receive stage callbacks.

    The domain parameter allows potentially having multiple completely
    separate layers of nodes with different purpose.
    """

    epsilon = 1e-5

    def __init__(self, generator, domain):
        super().__init__(generator)

        assert domain is not None
        assert generator.stage == 'initialize'

        self.domain = domain
        self.nodes = []
        self.final_nodes = []
        self.groups = []
        self.frozen = False

    def register_node(self, node):
        assert not self.frozen
        node.generator_plugin = self
        self.nodes.append(node)

    def initialize(self):
        self.frozen = True

        nodes = self.nodes
        tree = KDTree(len(nodes))

        for i, node in enumerate(nodes):
            tree.insert(node.point, i)

        tree.balance()
        processed = set()
        final_nodes = []
        groups = []

        for i in range(len(nodes)):
            if i in processed:
                continue

            # Find points to merge
            pending = [i]
            merge_set = set(pending)

            while pending:
                added = set()
                for j in pending:
                    point = nodes[j].point
                    eps = max(1, point.length) * self.epsilon
                    for co, idx, dist in tree.find_range(point, eps):
                        added.add(idx)
                pending = added.difference(merge_set)
                merge_set.update(added)

            assert merge_set.isdisjoint(processed)

            processed.update(merge_set)

            # Group the points
            merge_list = [nodes[i] for i in merge_set]
            merge_list.sort(key=lambda x: x.name)

            group_class = merge_list[0].group_class

            for item in merge_list[1:]:
                cls = item.group_class

                if issubclass(cls, group_class):
                    group_class = cls
                elif not issubclass(group_class, cls):
                    raise MetarigError(
                        'Group class conflict: {} and {} from {} of {}'.format(
                            group_class, cls, item.name, item.rig.base_bone,
                        )
                    )

            group = group_class(merge_list)
            group.build(final_nodes)

            groups.append(group)

        self.final_nodes = self.rigify_sub_objects = final_nodes
        self.groups = groups


class MergeGroup(object):
    """
    Standard node group, merges nodes based on certain rules and priorities.

    1. Nodes are classified into main and query nodes; query nodes are not merged.
    2. Nodes owned by the same rig cannot merge with each other.
    3. Node can only merge into target if node.can_merge_into(target) is true.
    4. Among multiple candidates in one rig, node.get_merge_priority(target) is used.
    5. The largest clusters of nodes that can merge are picked until none are left.

    The master nodes of the chosen clusters, plus query nodes, become 'final'.
    """

    def __init__(self, nodes):
        self.nodes = nodes

        for node in nodes:
            node.group = self

        def is_main(node):
            return isinstance(node, MainMergeNode)

        self.main_nodes = [n for n in nodes if is_main(n)]
        self.query_nodes = [n for n in nodes if not is_main(n)]

    def build(self, final_nodes):
        main_nodes = self.main_nodes

        # Sort nodes into rig buckets - can't merge within the same rig
        rig_table = collections.defaultdict(list)

        for node in main_nodes:
            rig_table[node.rig].append(node)

        # Build a 'can merge' table
        merge_table = {n: set() for n in main_nodes}

        for node in main_nodes:
            for rig, tgt_nodes in rig_table.items():
                if rig is not node.rig:
                    nodes = [n for n in tgt_nodes if node.can_merge_into(n)]
                    if nodes:
                        best_node = max(nodes, key=node.get_merge_priority)
                        merge_table[best_node].add(node)

        # Output groups starting with largest
        self.final_nodes = []

        pending = set(main_nodes)

        while pending:
            # Find largest group
            nodes = [n for n in main_nodes if n in pending]
            max_len = max(len(merge_table[n]) for n in nodes)

            nodes = [n for n in nodes if len(merge_table[n]) == max_len]

            # If a tie, try to resolve using comparison
            if len(nodes) > 1:
                weighted_nodes = [
                    (n, sum(
                        1 if (n.is_better_cluster(n2)
                              and not n2.is_better_cluster(n)) else 0
                        for n2 in nodes
                    ))
                    for n in nodes
                ]
                max_weight = max(wn[1] for wn in weighted_nodes)
                nodes = [wn[0] for wn in weighted_nodes if wn[1] == max_weight]

            # Final tie breaker is the name
            best = min(nodes, key=lambda n: n.name)
            child_set = merge_table[best]

            # Link children
            best.point = sum((c.point for c in child_set),
                             best.point) / (len(child_set) + 1)

            for child in [n for n in main_nodes if n in child_set]:
                child.point = best.point
                best.merge_from(child)
                child.merge_into(best)

            final_nodes.append(best)
            self.final_nodes.append(best)

            best.merge_done()

            # Remove merged nodes from the table
            pending.remove(best)
            pending -= child_set

            for children in merge_table.values():
                children &= pending

        for node in self.query_nodes:
            node.merge_done()

        final_nodes += self.query_nodes


class BaseMergeNode(GenerateCallbackHost):
    """Base class of mergeable nodes."""

    merge_domain = None
    merger = NodeMerger
    group_class = MergeGroup

    def __init__(self, rig, name, point, *, domain=None):
        self.rig = rig
        self.obj = rig.obj
        self.name = name
        self.point = Vector(point)

        merger = self.merger(rig.generator, domain or self.merge_domain)
        merger.register_node(self)

    def register_new_bone(self, new_name, old_name=None):
        self.generator_plugin.register_new_bone(new_name, old_name)

    def can_merge_into(self, other):
        raise NotImplementedError()

    def get_merge_priority(self, other):
        "Rank candidates to merge into."
        return 0


class MainMergeNode(BaseMergeNode):
    """
    Base class of standard mergeable nodes. Each node can either be
    a master of its cluster or a merged child node. Children become
    sub-objects of their master to receive callbacks in defined order.
    """

    def __init__(self, rig, name, point, *, domain=None):
        super().__init__(rig, name, point, domain=domain)

        self.merged_into = None
        self.merged = []

    def get_merged_siblings(self):
        master = self.merged_master
        return [master, *master.merged]

    def is_better_cluster(self, other):
        "Compare with the other node to choose between cluster masters."
        return False

    def can_merge_from(self, other):
        return True

    def can_merge_into(self, other):
        return other.can_merge_from(self)

    def merge_into(self, other):
        self.merged_into = other

    def merge_from(self, other):
        self.merged.append(other)

    @property
    def is_master_node(self):
        return not self.merged_into

    def merge_done(self):
        self.merged_master = self.merged_into or self
        self.rigify_sub_objects = list(self.merged)

        for child in self.merged:
            child.merge_done()


class QueryMergeNode(BaseMergeNode):
    """Base class for special nodes used only to query which nodes are at a certain location."""

    is_master_node = False
    require_match = True

    def merge_done(self):
        self.matched_nodes = [
            n for n in self.group.final_nodes if self.can_merge_into(n)
        ]
        self.matched_nodes.sort(key=self.get_merge_priority, reverse=True)

        if self.require_match and not self.matched_nodes:
            self.rig.raise_error(
                'Could not match control node for {}', self.name)