io_scene_psk_psa/psa/importer.py

import typing
from typing import List, Optional

import bpy
import numpy as np
from bpy.types import FCurve, Object, Context
from mathutils import Vector, Quaternion

from .config import PsaConfig, REMOVE_TRACK_LOCATION, REMOVE_TRACK_ROTATION
from .data import Psa
from .reader import PsaReader


class PsaImportOptions(object):
    def __init__(self):
        self.should_use_fake_user = False
        self.should_stash = False
        self.sequence_names = []
        self.should_overwrite = False
        self.should_write_keyframes = True
        self.should_write_metadata = True
        self.action_name_prefix = ''
        self.should_convert_to_samples = False
        self.bone_mapping_mode = 'CASE_INSENSITIVE'
        self.fps_source = 'SEQUENCE'
        self.fps_custom: float = 30.0
        self.should_use_config_file = True
        self.psa_config: PsaConfig = PsaConfig()


class ImportBone(object):
    def __init__(self, psa_bone: Psa.Bone):
        self.psa_bone: Psa.Bone = psa_bone
        self.parent: Optional[ImportBone] = None
        self.armature_bone = None
        self.pose_bone = None
        self.original_location: Vector = Vector()
        self.original_rotation: Quaternion = Quaternion()
        self.post_rotation: Quaternion = Quaternion()
        self.fcurves: List[FCurve] = []


def _calculate_fcurve_data(import_bone: ImportBone, key_data: typing.Iterable[float]):
    # Convert world-space transforms to local-space transforms.
    key_rotation = Quaternion(key_data[0:4])
    key_location = Vector(key_data[4:])
    q = import_bone.post_rotation.copy()
    q.rotate(import_bone.original_rotation)
    rotation = q
    q = import_bone.post_rotation.copy()
    if import_bone.parent is None:
        q.rotate(key_rotation.conjugated())
    else:
        q.rotate(key_rotation)
    rotation.rotate(q.conjugated())
    location = key_location - import_bone.original_location
    location.rotate(import_bone.post_rotation.conjugated())
    return rotation.w, rotation.x, rotation.y, rotation.z, location.x, location.y, location.z


class PsaImportResult:
    def __init__(self):
        self.warnings: List[str] = []


def _get_armature_bone_index_for_psa_bone(psa_bone_name: str, armature_bone_names: List[str], bone_mapping_mode: str = 'EXACT') -> Optional[int]:
    """
    @param psa_bone_name: The name of the PSA bone.
    @param armature_bone_names: The names of the bones in the armature.
    @param bone_mapping_mode: One of 'EXACT' or 'CASE_INSENSITIVE'.
    @return: The index of the armature bone that corresponds to the given PSA bone, or None if no such bone exists.
    """
    for armature_bone_index, armature_bone_name in enumerate(armature_bone_names):
        if bone_mapping_mode == 'CASE_INSENSITIVE':
            if armature_bone_name.lower() == psa_bone_name.lower():
                return armature_bone_index
        else:
            if armature_bone_name == psa_bone_name:
                return armature_bone_index
    return None

def _get_sample_frame_times(source_frame_count: int, frame_step: float) -> typing.Iterable[float]:
    # TODO: for correctness, we should also emit the target frame time as well (because the last frame can be a
    #  fractional frame).
    time = 0.0
    while time < source_frame_count - 1:
        yield time
        time += frame_step
    yield source_frame_count - 1

def _resample_sequence_data_matrix(sequence_data_matrix: np.ndarray, frame_step: float = 1.0) -> np.ndarray:
    """
    Resamples the sequence data matrix to the target frame count.
    @param sequence_data_matrix: FxBx7 matrix where F is the number of frames, B is the number of bones, and X is the
    number of data elements per bone.
    @param frame_step: The step between frames in the resampled sequence.
    @return: The resampled sequence data matrix, or sequence_data_matrix if no resampling is necessary.
    """
    if frame_step == 1.0:
        # No resampling is necessary.
        return sequence_data_matrix

    source_frame_count, bone_count = sequence_data_matrix.shape[:2]
    sample_frame_times = list(_get_sample_frame_times(source_frame_count, frame_step))
    target_frame_count = len(sample_frame_times)
    resampled_sequence_data_matrix = np.zeros((target_frame_count, bone_count, 7), dtype=float)

    for sample_frame_index, sample_frame_time in enumerate(sample_frame_times):
        frame_index = int(sample_frame_time)
        if sample_frame_time % 1.0 == 0.0:
            # Sample time has no fractional part, so just copy the frame.
            resampled_sequence_data_matrix[sample_frame_index, :, :] = sequence_data_matrix[frame_index, :, :]
        else:
            # Sample time has a fractional part, so interpolate between two frames.
            next_frame_index = frame_index + 1
            for bone_index in range(bone_count):
                source_frame_1_data = sequence_data_matrix[frame_index, bone_index, :]
                source_frame_2_data = sequence_data_matrix[next_frame_index, bone_index, :]
                factor = sample_frame_time - frame_index
                q = Quaternion((source_frame_1_data[:4])).slerp(Quaternion((source_frame_2_data[:4])), factor)
                q.normalize()
                l = Vector(source_frame_1_data[4:]).lerp(Vector(source_frame_2_data[4:]), factor)
                resampled_sequence_data_matrix[sample_frame_index, bone_index, :] = q.w, q.x, q.y, q.z, l.x, l.y, l.z

    return resampled_sequence_data_matrix


def import_psa(context: Context, psa_reader: PsaReader, armature_object: Object, options: PsaImportOptions) -> PsaImportResult:
    result = PsaImportResult()
    sequences = [psa_reader.sequences[x] for x in options.sequence_names]
    armature_data = typing.cast(bpy.types.Armature, armature_object.data)

    # Create an index mapping from bones in the PSA to bones in the target armature.
    psa_to_armature_bone_indices = {}
    armature_to_psa_bone_indices = {}
    armature_bone_names = [x.name for x in armature_data.bones]
    psa_bone_names = []
    duplicate_mappings = []

    for psa_bone_index, psa_bone in enumerate(psa_reader.bones):
        psa_bone_name: str = psa_bone.name.decode('windows-1252')
        armature_bone_index = _get_armature_bone_index_for_psa_bone(psa_bone_name, armature_bone_names, options.bone_mapping_mode)
        if armature_bone_index is not None:
            # Ensure that no other PSA bone has been mapped to this armature bone yet.
            if armature_bone_index not in armature_to_psa_bone_indices:
                psa_to_armature_bone_indices[psa_bone_index] = armature_bone_index
                armature_to_psa_bone_indices[armature_bone_index] = psa_bone_index
            else:
                # This armature bone has already been mapped to a PSA bone.
                duplicate_mappings.append((psa_bone_index, armature_bone_index, armature_to_psa_bone_indices[armature_bone_index]))
            psa_bone_names.append(armature_bone_names[armature_bone_index])
        else:
            psa_bone_names.append(psa_bone_name)

    # Warn about duplicate bone mappings.
    if len(duplicate_mappings) > 0:
        for (psa_bone_index, armature_bone_index, mapped_psa_bone_index) in duplicate_mappings:
            psa_bone_name = psa_bone_names[psa_bone_index]
            armature_bone_name = armature_bone_names[armature_bone_index]
            mapped_psa_bone_name = psa_bone_names[mapped_psa_bone_index]
            result.warnings.append(f'PSA bone {psa_bone_index} ({psa_bone_name}) could not be mapped to armature bone {armature_bone_index} ({armature_bone_name}) because the armature bone is already mapped to PSA bone {mapped_psa_bone_index} ({mapped_psa_bone_name})')

    # Report if there are missing bones in the target armature.
    missing_bone_names = set(psa_bone_names).difference(set(armature_bone_names))
    if len(missing_bone_names) > 0:
        result.warnings.append(
            f'The armature \'{armature_object.name}\' is missing {len(missing_bone_names)} bones that exist in '
            'the PSA:\n' +
            str(list(sorted(missing_bone_names)))
        )
    del armature_bone_names

    # Create intermediate bone data for import operations.
    import_bones = []
    psa_bone_names_to_import_bones = dict()

    for (psa_bone_index, psa_bone), psa_bone_name in zip(enumerate(psa_reader.bones), psa_bone_names):
        if psa_bone_index not in psa_to_armature_bone_indices:
            # PSA bone does not map to armature bone, skip it and leave an empty bone in its place.
            import_bones.append(None)
            continue
        import_bone = ImportBone(psa_bone)
        import_bone.armature_bone = armature_data.bones[psa_bone_name]
        import_bone.pose_bone = armature_object.pose.bones[psa_bone_name]
        psa_bone_names_to_import_bones[psa_bone_name] = import_bone
        import_bones.append(import_bone)

    bones_with_missing_parents = []

    for import_bone in filter(lambda x: x is not None, import_bones):
        armature_bone = import_bone.armature_bone
        has_parent = armature_bone.parent is not None
        if has_parent:
            if armature_bone.parent.name in psa_bone_names:
                import_bone.parent = psa_bone_names_to_import_bones[armature_bone.parent.name]
            else:
                # Add a warning if the parent bone is not in the PSA.
                bones_with_missing_parents.append(armature_bone)
        # Calculate the original location & rotation of each bone (in world-space maybe?)
        if has_parent:
            import_bone.original_location = armature_bone.matrix_local.translation - armature_bone.parent.matrix_local.translation
            import_bone.original_location.rotate(armature_bone.parent.matrix_local.to_quaternion().conjugated())
            import_bone.original_rotation = armature_bone.matrix_local.to_quaternion()
            import_bone.original_rotation.rotate(armature_bone.parent.matrix_local.to_quaternion().conjugated())
            import_bone.original_rotation.conjugate()
        else:
            import_bone.original_location = armature_bone.matrix_local.translation.copy()
            import_bone.original_rotation = armature_bone.matrix_local.to_quaternion().conjugated()

        import_bone.post_rotation = import_bone.original_rotation.conjugated()

    # Warn about bones with missing parents.
    if len(bones_with_missing_parents) > 0:
        count = len(bones_with_missing_parents)
        message = f'{count} bone(s) have parents that are not present in the PSA:\n' + str([x.name for x in bones_with_missing_parents])
        result.warnings.append(message)

    context.window_manager.progress_begin(0, len(sequences))

    # Create and populate the data for new sequences.
    actions = []
    for sequence_index, sequence in enumerate(sequences):
        # Add the action.
        sequence_name = sequence.name.decode('windows-1252')
        action_name = options.action_name_prefix + sequence_name

        # Get the bone track flags for this sequence, or an empty dictionary if none exist.
        sequence_bone_track_flags = dict()
        if sequence_name in options.psa_config.sequence_bone_flags.keys():
            sequence_bone_track_flags = options.psa_config.sequence_bone_flags[sequence_name]

        if options.should_overwrite and action_name in bpy.data.actions:
            action = bpy.data.actions[action_name]
        else:
            action = bpy.data.actions.new(name=action_name)

        # Calculate the target FPS.
        match options.fps_source:
            case 'CUSTOM':
                target_fps = options.fps_custom
            case 'SCENE':
                target_fps = context.scene.render.fps
            case 'SEQUENCE':
                target_fps = sequence.fps
            case _:
                raise ValueError(f'Unknown FPS source: {options.fps_source}')

        if options.should_write_keyframes:
            # Remove existing f-curves.
            action.fcurves.clear()

            # Create f-curves for the rotation and location of each bone.
            for psa_bone_index, armature_bone_index in psa_to_armature_bone_indices.items():
                bone_track_flags = sequence_bone_track_flags.get(psa_bone_index, 0)
                import_bone = import_bones[psa_bone_index]
                pose_bone = import_bone.pose_bone
                rotation_data_path = pose_bone.path_from_id('rotation_quaternion')
                location_data_path = pose_bone.path_from_id('location')
                add_rotation_fcurves = (bone_track_flags & REMOVE_TRACK_ROTATION) == 0
                add_location_fcurves = (bone_track_flags & REMOVE_TRACK_LOCATION) == 0
                import_bone.fcurves = [
                    action.fcurves.new(rotation_data_path, index=0, action_group=pose_bone.name) if add_rotation_fcurves else None,  # Qw
                    action.fcurves.new(rotation_data_path, index=1, action_group=pose_bone.name) if add_rotation_fcurves else None,  # Qx
                    action.fcurves.new(rotation_data_path, index=2, action_group=pose_bone.name) if add_rotation_fcurves else None,  # Qy
                    action.fcurves.new(rotation_data_path, index=3, action_group=pose_bone.name) if add_rotation_fcurves else None,  # Qz
                    action.fcurves.new(location_data_path, index=0, action_group=pose_bone.name) if add_location_fcurves else None,  # Lx
                    action.fcurves.new(location_data_path, index=1, action_group=pose_bone.name) if add_location_fcurves else None,  # Ly
                    action.fcurves.new(location_data_path, index=2, action_group=pose_bone.name) if add_location_fcurves else None,  # Lz
                ]

            # Read the sequence data matrix from the PSA.
            sequence_data_matrix = psa_reader.read_sequence_data_matrix(sequence_name)

            # Convert the sequence's data from world-space to local-space.
            for bone_index, import_bone in enumerate(import_bones):
                if import_bone is None:
                    continue
                for frame_index in range(sequence.frame_count):
                    # This bone has writeable keyframes for this frame.
                    key_data = sequence_data_matrix[frame_index, bone_index]
                    # Calculate the local-space key data for the bone.
                    sequence_data_matrix[frame_index, bone_index] = _calculate_fcurve_data(import_bone, key_data)

            # Resample the sequence data to the target FPS.
            # If the target frame count is the same as the source frame count, this will be a no-op.
            resampled_sequence_data_matrix = _resample_sequence_data_matrix(sequence_data_matrix,
                                                                            frame_step=sequence.fps / target_fps)

            # Write the keyframes out.
            # Note that the f-curve data consists of alternating time and value data.
            target_frame_count = resampled_sequence_data_matrix.shape[0]
            fcurve_data = np.zeros(2 * target_frame_count, dtype=float)
            fcurve_data[0::2] = range(0, target_frame_count)

            for bone_index, import_bone in enumerate(import_bones):
                if import_bone is None:
                    continue
                for fcurve_index, fcurve in enumerate(import_bone.fcurves):
                    if fcurve is None:
                        continue
                    fcurve_data[1::2] = resampled_sequence_data_matrix[:, bone_index, fcurve_index]
                    fcurve.keyframe_points.add(target_frame_count)
                    fcurve.keyframe_points.foreach_set('co', fcurve_data)
                    for fcurve_keyframe in fcurve.keyframe_points:
                        fcurve_keyframe.interpolation = 'LINEAR'

            if options.should_convert_to_samples:
                # Bake the curve to samples.
                for fcurve in action.fcurves:
                    fcurve.convert_to_samples(start=0, end=sequence.frame_count)

        # Write meta-data.
        if options.should_write_metadata:
            action.psa_export.fps = target_fps

        action.use_fake_user = options.should_use_fake_user

        actions.append(action)

        context.window_manager.progress_update(sequence_index)

    # If the user specifies, store the new animations as strips on a non-contributing NLA track.
    if options.should_stash:
        if armature_object.animation_data is None:
            armature_object.animation_data_create()
        for action in actions:
            nla_track = armature_object.animation_data.nla_tracks.new()
            nla_track.name = action.name
            nla_track.mute = True
            nla_track.strips.new(name=action.name, start=0, action=action)

    context.window_manager.progress_end()

    return result