anim.py

import bpy
from mathutils import Quaternion, Vector, Euler
import math
import os
import re
import xml.etree.ElementTree as ET

def _unwrap_deg(prev, cur):
    """Return cur adjusted by +/-360 to minimise jump from prev."""
    if prev is None:
        return cur
    # choose among cur-360, cur, cur+360
    options = (cur - 360.0, cur, cur + 360.0)
    best = min(options, key=lambda v: abs(v - prev))
    return best

template = """<?xml version="1.0" encoding="utf-8"?>
<WadAnimation xmlns:xsd="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance">
    <FrameRate>1</FrameRate>
    <StateId>2</StateId>
    <EndFrame>43</EndFrame>
    <NextAnimation>103</NextAnimation>
    <NextFrame>0</NextFrame>
    <Name>STAND_IDLE</Name>
    <StartVelocity>0</StartVelocity>
    <EndVelocity>0</EndVelocity>
    <StartLateralVelocity>0</StartLateralVelocity>
    <EndLateralVelocity>0</EndLateralVelocity>
    <KeyFrames />
    <StateChanges />
    <AnimCommands />
</WadAnimation>
"""

def _best_posebone_name(obj, token):
    """Pick the most likely pose bone name for a given TR token (e.g. 'LEFT_THIGH')."""
    matches = [pb.name for pb in obj.pose.bones if token in pb.name]
    if not matches:
        return None

    # Prefer exact-ish endings first.
    preferred = [
        f"{token}_BONE",
        f"_{token}_BONE",
        token,
        f"_{token}",
    ]
    for p in preferred:
        for m in matches:
            if m.endswith(p):
                return m

    # Otherwise pick the shortest (usually the cleanest name).
    return sorted(matches, key=len)[0]

def export_anim(template_anim, zoffset, rig_name):
    # Load template (or use default)
    if os.path.exists(template_anim):
        with open(template_anim, encoding='utf-8', errors='ignore') as f:
            xml_string = '\n'.join(f.readlines())
    else:
        xml_string = template

    xml_string = xml_string.replace('utf-16', 'utf8')
    tree = ET.ElementTree(ET.fromstring(xml_string))
    root = tree.getroot()

    # Find keyframes node and wipe template frames
    keyframes_node = tree.find("KeyFrames")
    for keyframe in list(keyframes_node.findall("WadKeyFrame")):
        keyframes_node.remove(keyframe)

    # Resolve rig object
    obj = bpy.data.objects.get(rig_name)
    if obj is None:
        obj = bpy.context.view_layer.objects.active
    if obj is None or obj.type != 'ARMATURE':
        raise RuntimeError(f"WADBlender: couldn't find an armature named '{rig_name}', and no active armature selected.")

    if obj.animation_data is None or obj.animation_data.action is None:
        raise RuntimeError("WADBlender: the selected armature has no active Action to export.")

    action = obj.animation_data.action
    scene = bpy.context.scene

    # Lara skin order as expected by WadTool
    lara_skin_names = [
        'HIPS', 'LEFT_THIGH', 'LEFT_SHIN', 'LEFT_FOOT',
        'RIGHT_THIGH', 'RIGHT_SHIN', 'RIGHT_FOOT',
        'TORSO', 'RIGHT_UPPER_ARM', 'RIGHT_FOREARM', 'RIGHT_HAND',
        'LEFT_UPPER_ARM', 'LEFT_FOREARM', 'LEFT_HAND', 'HEAD'
    ]

    # Map TR tokens -> actual pose bone names in this rig (supports prefixes like 'LARA_HIPS_BONE')
    bone_map = {}
    for token in lara_skin_names:
        pb_name = _best_posebone_name(obj, token)
        if pb_name is None:
            raise RuntimeError(f"WADBlender: couldn't find a bone containing '{token}' in the armature '{obj.name}'.")
        bone_map[token] = pb_name

    # Root motion bone:
    # - Prefer TR4-style dummy root 'mesh_00_anim' if present (common in TR4 pipelines)
    # - Otherwise fall back to the hips bone (works for classic WADBlender rigs)
    root_pb = obj.pose.bones.get('mesh_00_anim')
    if root_pb is None:
        root_pb = obj.pose.bones[bone_map['HIPS']]

    # Use action frame range (not keyframe count in the first fcurve)
    frame_start = int(round(action.frame_range[0]))
    frame_end   = int(round(action.frame_range[1]))
    keyframes_count = max(1, frame_end - frame_start + 1)

    # WadTool uses 0-based EndFrame index
    endframe_node = tree.find("EndFrame")
    endframe_node.text = str(keyframes_count - 1)

    # Store current scene frame so we can restore it
    cur_frame = scene.frame_current

    # Preallocate
    rotations = [[ [0.0, 0.0, 0.0] for _ in range(keyframes_count) ] for _ in range(len(lara_skin_names))]
    offsets = [ [0.0, 0.0, 0.0] for _ in range(keyframes_count)]

    try:
        for fi, frame in enumerate(range(frame_start, frame_end + 1)):
            scene.frame_set(frame)

            # Root offset in Blender space (metres), expressed in the armature's axes.
            # IMPORTANT: matrix_basis translation is in the *bone's local axes*.
            # If the dummy root bone's local axes are rotated (common in TR rigs),
            # "lift" can end up in X/Y instead of Blender Z and therefore export sideways.
            # Using pose matrix translation minus rest translation gives a delta in
            # the armature's coordinate system (Blender X/Y/Z).
            rest_t = root_pb.bone.matrix_local.to_translation()
            pose_t = root_pb.matrix.to_translation()
            root_loc = (pose_t - rest_t).copy()
            # Add zoffset in *Blender Z-up* before axis mapping (so it's always "lift")
            root_loc.z += (zoffset / 512.0)

            # Convert to TR units (512 units = 1m)
            loc_tru = root_loc * 512.0

            # Axis mapping (WadTool): user verified WadTool uses Y as UP.
            # Blender axes: X=right, Y=forward, Z=up
            # Map to WadTool offsets:
            #   Offset X =  Blender X   (left/right)
            #   Offset Y =  Blender Z   (up/down lift)
            #   Offset Z = -Blender Y   (forward/back)
            offsets[fi][0] =  loc_tru.x
            offsets[fi][1] =  loc_tru.z
            offsets[fi][2] = -loc_tru.y

            # Bone rotations
            for bi, token in enumerate(lara_skin_names):
                pb = obj.pose.bones[bone_map[token]]
                q = pb.matrix_basis.to_quaternion()

                # Convert to euler in the same order used on import ('ZXY' in this addon)
                euler = q.to_euler("ZXY")
                ang = [math.degrees(a) for a in euler]

                # Angle mapping (legacy wadblender mapping)
                rx = -ang[0]
                ry =  ang[1]
                rz = -ang[2]
                # No special 180° correction (handled by consistent basis conversion)
                # Unwrap angles to avoid 0/360 discontinuities (WadTool interpolates across frames)
                if fi > 0:
                    rx = _unwrap_deg(rotations[bi][fi-1][0], rx)
                    ry = _unwrap_deg(rotations[bi][fi-1][1], ry)
                    rz = _unwrap_deg(rotations[bi][fi-1][2], rz)
                rotations[bi][fi][0] = rx
                rotations[bi][fi][1] = ry
                rotations[bi][fi][2] = rz

    finally:
        scene.frame_set(cur_frame)

    # Write output anim file
    for fi in range(keyframes_count):
        wadkf = ET.SubElement(keyframes_node, 'WadKeyFrame')

        bbox = ET.SubElement(wadkf, 'BoundingBox')

        minimum = ET.SubElement(bbox, 'Minimum')
        ET.SubElement(minimum, 'X').text = "0"
        ET.SubElement(minimum, 'Y').text = "0"
        ET.SubElement(minimum, 'Z').text = "0"

        maximum = ET.SubElement(bbox, 'Maximum')
        ET.SubElement(maximum, 'X').text = "0"
        ET.SubElement(maximum, 'Y').text = "0"
        ET.SubElement(maximum, 'Z').text = "0"

        offset = ET.SubElement(wadkf, 'Offset')
        ET.SubElement(offset, 'X').text = f"{offsets[fi][0]:.6f}"
        ET.SubElement(offset, 'Y').text = f"{offsets[fi][1]:.6f}"
        ET.SubElement(offset, 'Z').text = f"{offsets[fi][2]:.6f}"

        angles = ET.SubElement(wadkf, 'Angles')
        for bi in range(len(lara_skin_names)):
            kfr = ET.SubElement(angles, 'WadKeyFrameRotation')
            rot = ET.SubElement(kfr, 'Rotations')
            ET.SubElement(rot, 'X').text = f"{rotations[bi][fi][0]:.6f}"
            ET.SubElement(rot, 'Y').text = f"{rotations[bi][fi][1]:.6f}"
            ET.SubElement(rot, 'Z').text = f"{rotations[bi][fi][2]:.6f}"

    tree.write(template_anim)



def _read_anim_xml_text(filepath: str) -> str:
    """Read WadTool/Tomb Editor .anim XML robustly even if the XML header lies about encoding."""
    data = open(filepath, "rb").read()

    # BOM-aware decodes first
    if data.startswith(b"\xef\xbb\xbf"):
        text = data.decode("utf-8-sig", errors="replace")
    elif data.startswith(b"\xff\xfe") or data.startswith(b"\xfe\xff"):
        text = data.decode("utf-16", errors="replace")
    else:
        try:
            text = data.decode("utf-8", errors="strict")
        except UnicodeDecodeError:
            # WadTool sometimes writes cp1252-ish bytes
            text = data.decode("cp1252", errors="replace")

    # If the XML declaration contains encoding=..., strip it (we already have unicode text now).
    s = text.lstrip()
    if s.startswith("<?xml"):
        end = text.find("?>")
        if end != -1:
            decl = text[:end+2]
            decl2 = re.sub(r'\sencoding\s*=\s*["\'][^"\']*["\']', "", decl, flags=re.I)
            text = decl2 + text[end+2:]
    return text
def import_anim(filepath: str,
                rig_name: str = "",
                start_frame: int | None = None,
                baseline_offsets: bool = True,
                action_name: str | None = None):
    """Import a Tomb Editor / WadTool .anim (XML) file into the selected armature as a new Action.

    Notes:
      - This importer currently targets the TR4 Lara skeleton order used by the exporter.
      - Root offsets are imported as *deltas* (frame0 treated as baseline) by default, to avoid teleporting.
      - Axis mapping is the inverse of this addon's export mapping (WadTool Y-up).
    """

    # Resolve rig object
    obj = bpy.data.objects.get(rig_name) if rig_name else None
    if obj is None:
        obj = bpy.context.view_layer.objects.active
    if obj is None or obj.type != 'ARMATURE':
        raise RuntimeError("WADBlender: select an Armature (or provide rig_name) before importing .anim")

    scene = bpy.context.scene
    if start_frame is None:
        start_frame = scene.frame_current

    # Parse XML (robust: WadTool sometimes mismatches XML encoding declarations)
    xml_text = _read_anim_xml_text(filepath)
    root = ET.fromstring(xml_text)
    keyframes_node = root.find("KeyFrames")
    if keyframes_node is None:
        raise RuntimeError("WADBlender: invalid .anim (missing <KeyFrames>)")
    keyframes = list(keyframes_node.findall("WadKeyFrame"))
    if not keyframes:
        raise RuntimeError("WADBlender: .anim contains 0 keyframes")

    # Create Action
    if obj.animation_data is None:
        obj.animation_data_create()
    if action_name is None:
        action_name = os.path.splitext(os.path.basename(filepath))[0]
    action = bpy.data.actions.new(name=action_name)
    obj.animation_data.action = action

    # Lara skin order (must match exporter)
    lara_skin_names = [
        'HIPS', 'LEFT_THIGH', 'LEFT_SHIN', 'LEFT_FOOT',
        'RIGHT_THIGH', 'RIGHT_SHIN', 'RIGHT_FOOT',
        'TORSO', 'RIGHT_UPPER_ARM', 'RIGHT_FOREARM', 'RIGHT_HAND',
        'LEFT_UPPER_ARM', 'LEFT_FOREARM', 'LEFT_HAND', 'HEAD'
    ]

    # Bone mapping
    bone_map = {}
    for token in lara_skin_names:
        pb_name = _best_posebone_name(obj, token)
        if pb_name is None:
            raise RuntimeError(f"WADBlender: couldn't find a bone containing '{token}' in the armature '{obj.name}'.")
        bone_map[token] = pb_name

    # Root motion bone
    root_pb = obj.pose.bones.get('mesh_00_anim')
    if root_pb is None:
        root_pb = obj.pose.bones[bone_map['HIPS']]

    def _read_offset(kf):
        off = kf.find('Offset')
        if off is None:
            return (0.0, 0.0, 0.0)
        return (
            float(off.find('X').text),
            float(off.find('Y').text),
            float(off.find('Z').text),
        )

    off0 = _read_offset(keyframes[0]) if baseline_offsets else (0.0, 0.0, 0.0)

    # Save current frame
    cur_frame = scene.frame_current

    # Cache rest matrix for stable root translation in armature axes
    rest_mat = root_pb.bone.matrix_local.copy()
    rest_t = rest_mat.to_translation().copy()

    try:
        for fi, kf in enumerate(keyframes):
            f = int(start_frame + fi)
            scene.frame_set(f)

            # ---- Root offsets ----
            ox, oy, oz = _read_offset(kf)
            ox -= off0[0]; oy -= off0[1]; oz -= off0[2]

            # Inverse of export mapping (WadTool Y-up):
            #   Wad X =  Bl X * 512
            #   Wad Y =  Bl Z * 512
            #   Wad Z = -Bl Y * 512
            # => Bl (metres) = (WadX, -WadZ, WadY) / 512
            delta_bl = Vector((ox, -oz, oy)) / 512.0

            # Apply as a translation in armature axes (robust against bone roll)
            m = rest_mat.copy()
            m.translation = rest_t + delta_bl
            root_pb.matrix = m
            bpy.context.view_layer.update()
            root_pb.keyframe_insert(data_path='location')

            # ---- Bone rotations ----
            angles_node = kf.find('Angles')
            if angles_node is None:
                continue
            rots = list(angles_node.findall('WadKeyFrameRotation'))
            if len(rots) != len(lara_skin_names):
                raise RuntimeError(f"WADBlender: expected {len(lara_skin_names)} rotations, got {len(rots)}")

            # Track previous degrees to unwrap 0/360 discontinuities per bone axis
            if 'prev_deg' not in locals():
                prev_deg = [[None, None, None] for _ in range(len(lara_skin_names))]

            for bi, token in enumerate(lara_skin_names):
                pb = obj.pose.bones[bone_map[token]]

                rnode = rots[bi].find('Rotations')
                rx = float(rnode.find('X').text)
                ry = float(rnode.find('Y').text)
                rz = float(rnode.find('Z').text)
                # No special 180° correction
                # Unwrap per-axis degrees across frames to prevent interpolation flips
                rx = _unwrap_deg(prev_deg[bi][0], rx); prev_deg[bi][0] = rx
                ry = _unwrap_deg(prev_deg[bi][1], ry); prev_deg[bi][1] = ry
                rz = _unwrap_deg(prev_deg[bi][2], rz); prev_deg[bi][2] = rz

                # Undo legacy sign mapping back to Blender Euler(ZXY)
                ang0 = -rx
                ang1 =  ry
                ang2 = -rz

                pb.rotation_mode = 'ZXY'
                pb.rotation_euler = Euler((math.radians(ang0), math.radians(ang1), math.radians(ang2)), 'ZXY')
                pb.keyframe_insert(data_path='rotation_euler')

    finally:
        scene.frame_set(cur_frame)

    return action