"""Silhouette-preserving helpers shared by the animation retargeters.

The source animation owns the direction of each major body segment. Copying
only a bone's rest-relative rotation preserves a target rig's modelling pose;
characters authored with crossed or rearward arms then keep those directions
in every idle, dance, and locomotion clip.

These helpers transfer animated joint-to-joint directions for the torso and
limbs while retaining the target bone's own roll. This is an offline bake of
the source performance, not a runtime procedural animation layer.
"""

from mathutils import Quaternion, Vector


def build_segment_pairs(mapping, source_roles, target_roles):
    """Return target bone -> (target child, source bone, source child)."""
    chains = [
        [target_roles.hips] + list(target_roles.spine),
    ]
    for side in ("L", "R"):
        chains.append([
            target_roles.limb.get((role, side))
            for role in ("shoulder", "upper_arm", "forearm", "hand")
        ])
        chains.append([
            target_roles.limb.get((role, side))
            for role in ("thigh", "shin", "foot", "toe")
        ])

    pairs = {}
    for chain in chains:
        chain = [name for name in chain if name]
        for index, target_name in enumerate(chain[:-1]):
            source_name = mapping.get(target_name)
            if not source_name:
                continue
            # Chain resampling can map neighbouring target spine bones to the
            # same source joint. Look ahead to the first genuinely new joint.
            for target_child in chain[index + 1:]:
                source_child = mapping.get(target_child)
                if source_child and source_child != source_name:
                    pairs[target_name] = (
                        target_child, source_name, source_child
                    )
                    break
    return pairs


def _twist_about(rotation, axis):
    """Extract the twist component of a world-space quaternion around axis."""
    axis = axis.normalized()
    vector = Vector((rotation.x, rotation.y, rotation.z))
    projected = axis * vector.dot(axis)
    twist = Quaternion((rotation.w, projected.x, projected.y, projected.z))
    if twist.magnitude < 1e-8:
        return Quaternion()
    twist.normalize()
    return twist


def heading_inverse(pose_rotation, rest_rotation, up):
    """World-up heading that removes a capture's initial stage orientation."""
    return _twist_about(
        pose_rotation @ rest_rotation.inverted(), up
    ).inverted()


def authored_world_rotation(
    target_name,
    source_name,
    source_eval,
    source_world,
    source_rest_rot,
    target_rest,
    target_rest_rot,
    yaw,
    yaw_inverse,
    segment_pairs,
    source_heading=Quaternion(),
):
    """Retarget one bone, using authored segment direction where available."""
    pose_rotation = source_heading @ (
        source_world @ source_eval.pose.bones[source_name].matrix
    ).to_quaternion()
    delta = pose_rotation @ source_rest_rot[source_name].inverted()
    full_rotation = (
        yaw @ delta @ yaw_inverse
    ) @ target_rest_rot[target_name]

    pair = segment_pairs.get(target_name)
    if pair is None:
        return full_rotation
    target_child, pair_source, source_child = pair
    if (
        pair_source not in source_eval.pose.bones
        or source_child not in source_eval.pose.bones
        or target_child not in target_rest
    ):
        return full_rotation

    source_head = (
        source_world @ source_eval.pose.bones[pair_source].matrix
    ).translation
    source_child_head = (
        source_world @ source_eval.pose.bones[source_child].matrix
    ).translation
    desired_direction = yaw @ (
        source_heading @ (source_child_head - source_head)
    )
    target_direction = (
        target_rest[target_child].translation
        - target_rest[target_name].translation
    )
    if desired_direction.length < 1e-7 or target_direction.length < 1e-7:
        return full_rotation
    desired_direction.normalize()
    target_direction.normalize()

    swing = target_direction.rotation_difference(desired_direction)
    swung_rotation = swing @ target_rest_rot[target_name]
    residual = full_rotation @ swung_rotation.inverted()
    return (
        _twist_about(residual, desired_direction) @ swung_rotation
    ).normalized()