pygame_model_anim.py

import numpy as np
import time
from math import cos, sin, acos, sqrt

# --- Quaternion utility functions for rotations ---

def quat_from_axis_angle(axis, angle_rad):
    axis = axis / np.linalg.norm(axis)
    s = sin(angle_rad / 2)
    return np.array([cos(angle_rad/2), axis[0]*s, axis[1]*s, axis[2]*s])

def quat_multiply(q1, q2):
    w1, x1, y1, z1 = q1
    w2, x2, y2, z2 = q2
    w = w1*w2 - x1*x2 - y1*y2 - z1*z2
    x = w1*x2 + x1*w2 + y1*z2 - z1*y2
    y = w1*y2 - x1*z2 + y1*w2 + z1*x2
    z = w1*z2 + x1*y2 - y1*x2 + z1*w2
    return np.array([w, x, y, z])

def quat_to_matrix(q):
    w, x, y, z = q
    return np.array([
        [1-2*(y**2+z**2), 2*(x*y - z*w),   2*(x*z + y*w), 0],
        [2*(x*y + z*w),   1-2*(x**2+z**2), 2*(y*z - x*w), 0],
        [2*(x*z - y*w),   2*(y*z + x*w),   1-2*(x**2+y**2),0],
        [0, 0, 0, 1]
    ])

def quat_slerp(q0, q1, t):
    # Spherical linear interpolation of quaternions
    dot = np.dot(q0, q1)
    if dot < 0.0:
        q1 = -q1
        dot = -dot
    DOT_THRESHOLD = 0.9995
    if dot > DOT_THRESHOLD:
        result = q0 + t*(q1 - q0)
        return result / np.linalg.norm(result)
    theta_0 = acos(dot)
    theta = theta_0 * t
    q2 = q1 - q0*dot
    q2 = q2 / np.linalg.norm(q2)
    return q0*cos(theta) + q2*sin(theta)

def translation_matrix(v):
    m = np.identity(4)
    m[:3, 3] = v
    return m

def compose_transform(translation, rotation_quat):
    rot_m = quat_to_matrix(rotation_quat)
    trans_m = translation_matrix(translation)
    return np.dot(trans_m, rot_m)

# --- Bone Class ---

class Bone:
    def __init__(self, name, parent=None):
        self.name = name
        self.parent = parent
        self.children = []

        self.rest_translation = np.zeros(3)
        self.rest_rotation = np.array([1,0,0,0])  # Quaternion w,x,y,z

        # Current pose
        self.translation = self.rest_translation.copy()
        self.rotation = self.rest_rotation.copy()

        # Override for manual repositioning (optional blending)
        self.override_translation = None
        self.override_rotation = None
        self.override_blend = 0.0  # 0 = no override, 1 = full override

        if parent:
            parent.children.append(self)

    def get_local_transform(self):
        # Compose translation + rotation (with override blended if active)
        trans = self.translation
        rot = self.rotation
        if self.override_blend > 0:
            # Blend translation
            if self.override_translation is not None:
                trans = (1 - self.override_blend)*trans + self.override_blend*self.override_translation
            # Blend rotation with slerp
            if self.override_rotation is not None:
                rot = quat_slerp(rot, self.override_rotation, self.override_blend)
        return compose_transform(trans, rot)

    def get_global_transform(self):
        local = self.get_local_transform()
        if self.parent:
            parent_global = self.parent.get_global_transform()
            return np.dot(parent_global, local)
        return local

    def reset_pose(self):
        self.translation = self.rest_translation.copy()
        self.rotation = self.rest_rotation.copy()
        self.override_translation = None
        self.override_rotation = None
        self.override_blend = 0.0
        for c in self.children:
            c.reset_pose()

    # Helpers to apply rotation/translation in local space
    def set_translation(self, v):
        self.translation = np.array(v)

    def set_rotation_quat(self, q):
        self.rotation = np.array(q)

    def rotate_axis_angle(self, axis, angle_rad):
        q_rot = quat_from_axis_angle(axis, angle_rad)
        self.rotation = quat_multiply(q_rot, self.rotation)

    def translate(self, v):
        self.translation += np.array(v)

    # Override control
    def set_override(self, translation=None, rotation=None, blend=1.0):
        if translation is not None:
            self.override_translation = np.array(translation)
        if rotation is not None:
            self.override_rotation = np.array(rotation)
        self.override_blend = blend

    def clear_override(self):
        self.override_translation = None
        self.override_rotation = None
        self.override_blend = 0.0

# --- Pose class for easy pose management ---

class Pose:
    def __init__(self, model):
        self.model = model
        # Dict bone_name -> (translation, rotation_quat)
        self.bone_transforms = {}

    def set_bone(self, bone_name, translation=None, rotation=None):
        self.bone_transforms[bone_name] = (translation, rotation)

    def apply_to_model(self, blend=1.0):
        for bone_name, (trans, rot) in self.bone_transforms.items():
            bone = self.model.bones.get(bone_name)
            if bone is None:
                continue
            if trans is not None:
                if blend == 1.0:
                    bone.set_translation(trans)
                else:
                    bone.translation = (1-blend)*bone.translation + blend*np.array(trans)
            if rot is not None:
                if blend == 1.0:
                    bone.set_rotation_quat(rot)
                else:
                    bone.rotation = quat_slerp(bone.rotation, rot, blend)

    @classmethod
    def from_rest_pose(cls, model):
        p = cls(model)
        for name, bone in model.bones.items():
            p.set_bone(name, bone.rest_translation, bone.rest_rotation)
        return p

# --- Animation Player for smooth playback and blending ---

class AnimationPlayer:
    def __init__(self, model):
        self.model = model
        self.keyframes = []  # (time, Pose)
        self.duration = 0.0
        self.playing = False
        self.loop = True
        self.start_time = 0.0
        self.speed = 1.0
        self.current_pose = None

    def add_keyframe(self, time_sec, pose):
        self.keyframes.append((time_sec, pose))
        self.keyframes.sort(key=lambda kf: kf[0])
        self.duration = max(self.duration, time_sec)

    def play(self):
        self.playing = True
        self.start_time = time.time()

    def stop(self):
        self.playing = False

    def get_pose_at(self, t):
        if not self.keyframes:
            return None
        if t <= self.keyframes[0][0]:
            return self.keyframes[0][1]
        if t >= self.keyframes[-1][0]:
            return self.keyframes[-1][1]

        for i in range(len(self.keyframes)-1):
            t0, p0 = self.keyframes[i]
            t1, p1 = self.keyframes[i+1]
            if t0 <= t <= t1:
                alpha = (t - t0) / (t1 - t0)
                return self.blend_poses(p0, p1, alpha)

    def blend_poses(self, p0, p1, alpha):
        blended = Pose(self.model)
        for bone_name in self.model.bones.keys():
            t0, r0 = p0.bone_transforms.get(bone_name, (None,None))
            t1, r1 = p1.bone_transforms.get(bone_name, (None,None))
            if t0 is None: t0 = self.model.bones[bone_name].translation
            if r0 is None: r0 = self.model.bones[bone_name].rotation
            if t1 is None: t1 = t0
            if r1 is None: r1 = r0
            trans = (1-alpha)*np.array(t0) + alpha*np.array(t1)
            rot = quat_slerp(r0, r1, alpha)
            blended.set_bone(bone_name, trans, rot)
        return blended

    def update(self):
        if not self.playing:
            return
        t = (time.time() - self.start_time) * self.speed
        if self.loop and self.duration > 0:
            t = t % self.duration
        pose = self.get_pose_at(t)
        if pose:
            pose.apply_to_model()
            self.current_pose = pose

# --- Model class with bone dictionary and helpers ---

class Model:
    def __init__(self):
        self.bones = {}
        self.root_bones = []

    def add_bone(self, name, parent_name=None):
        parent = self.bones.get(parent_name)
        bone = Bone(name, parent)
        self.bones[name] = bone
        if parent is None:
            self.root_bones.append(bone)
        return bone

    def reset_pose(self):
        for bone in self.bones.values():
            bone.reset_pose()

    def print_hierarchy(self):
        def recurse(bone, indent=0):
            print("  "*indent + bone.name)
            for c in bone.children:
                recurse(c, indent+1)
        for root in self.root_bones:
            recurse(root)

# --- Example Usage ---

if __name__ == "__main__":
    model = Model()
    model.add_bone("root")
    model.add_bone("spine", "root")
    model.add_bone("right_arm", "spine")
    model.add_bone("right_forearm", "right_arm")
    model.add_bone("right_hand", "right_forearm")

    model.reset_pose()

    # Create rest pose
    rest_pose = Pose.from_rest_pose(model)

    # Create salute pose
    salute_pose = Pose(model)
    salute_pose.set_bone("right_arm", rotation=quat_from_axis_angle(np.array([1,0,0]), np.radians(-70)))
    salute_pose.set_bone("right_forearm", rotation=quat_from_axis_angle(np.array([1,0,0]), np.radians(-45)))
    salute_pose.set_bone("right_hand", rotation=quat_from_axis_angle(np.array([0,0,1]), np.radians(15)))

    # Setup animation player
    anim_player = AnimationPlayer(model)
    anim_player.add_keyframe(0.0, rest_pose)
    anim_player.add_keyframe(2.0, salute_pose)
    anim_player.add_keyframe(4.0, rest_pose)
    anim_player.play()

    print("Animating soldier salute, press Ctrl+C to stop...")

    try:
        while True:
            anim_player.update()

            # Example: print right_arm global transform
            right_arm_global = model.bones["right_arm"].get_global_transform()
            print(f"Right arm global transform:\n{right_arm_global}\n")

            time.sleep(0.1)
    except KeyboardInterrupt:
        print("Animation stopped.")