pygame_model_anim_v1.py

import numpy as np
import time

def rotation_matrix(axis, theta):
    axis = axis / np.linalg.norm(axis)
    a = np.cos(theta / 2.0)
    b, c, d = -axis*np.sin(theta / 2.0)
    aa, bb, cc, dd = a*a, b*b, c*c, d*d
    bc, ad, ac, ab, bd, cd = b*c, a*d, a*c, a*b, b*d, c*d
    return np.array([[aa+bb-cc-dd, 2*(bc+ad),   2*(bd-ac), 0],
                     [2*(bc-ad),   aa+cc-bb-dd, 2*(cd+ab), 0],
                     [2*(bd+ac),   2*(cd-ab),   aa+dd-bb-cc, 0],
                     [0, 0, 0, 1]])

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

def lerp_matrix(m1, m2, t):
    """Simple element-wise linear interpolation between two 4x4 matrices."""
    return m1 * (1 - t) + m2 * t

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

        self.rest_transform = np.identity(4)  # Default pose
        self.local_transform = np.identity(4) # Current local transform

        if parent:
            parent.children.append(self)

    def get_global_transform(self):
        if self.parent:
            return np.dot(self.parent.get_global_transform(), self.local_transform)
        else:
            return self.local_transform

    def reset_pose(self):
        self.local_transform = self.rest_transform.copy()
        for c in self.children:
            c.reset_pose()

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

    def add_bone(self, name, parent_name=None):
        parent = self.bones.get(parent_name, None)
        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 get_pose_matrices(self):
        """Returns a dict of global matrices for all bones."""
        return {name: bone.get_global_transform() for name, bone in self.bones.items()}

class Animation:
    def __init__(self, model):
        self.model = model
        self.keyframes = []  # list of (time, pose) where pose is dict of bone_name->matrix
        self.duration = 0

    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 interpolate_pose(self, t):
        """Interpolate between keyframes at time t."""
        if not self.keyframes:
            return

        if t <= self.keyframes[0][0]:
            return self.keyframes[0][1]

        if t >= self.keyframes[-1][0]:
            return self.keyframes[-1][1]

        # Find keyframes surrounding t
        for i in range(len(self.keyframes) - 1):
            t0, pose0 = self.keyframes[i]
            t1, pose1 = self.keyframes[i+1]
            if t0 <= t <= t1:
                alpha = (t - t0) / (t1 - t0)
                interpolated_pose = {}
                for bone_name in pose0.keys():
                    m0 = pose0[bone_name]
                    m1 = pose1[bone_name]
                    interpolated_pose[bone_name] = lerp_matrix(m0, m1, alpha)
                return interpolated_pose

    def apply_pose(self, pose):
        """Apply a pose (dict of bone_name->matrix) to the model."""
        for bone_name, matrix in pose.items():
            if bone_name in self.model.bones:
                self.model.bones[bone_name].local_transform = matrix

    def play(self, start_time, current_time):
        """Apply the interpolated pose based on time elapsed since start_time."""
        t = current_time - start_time
        # Loop animation for demonstration
        t = t % self.duration
        pose = self.interpolate_pose(t)
        if pose:
            self.apply_pose(pose)

# Helper to create a rotation matrix around an axis and translation combined
def create_transform(axis, angle_rad, translation):
    rot = rotation_matrix(axis, angle_rad)
    trans = translation_matrix(translation)
    return np.dot(trans, rot)

# === Usage Example ===
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')

    # Set rest poses for all bones (identity here, but could be any)
    model.reset_pose()

    # Create two poses: rest pose and salute pose
    rest_pose = {name: bone.rest_transform.copy() for name, bone in model.bones.items()}

    salute_pose = {}
    for name, bone in model.bones.items():
        salute_pose[name] = bone.rest_transform.copy()

    # Example: rotate right_arm 70 degrees up to salute
    salute_pose['right_arm'] = create_transform(axis=np.array([1,0,0]), angle_rad=np.radians(-70), translation=np.array([0,0,0]))
    # Bend right_forearm 45 degrees
    salute_pose['right_forearm'] = create_transform(axis=np.array([1,0,0]), angle_rad=np.radians(-45), translation=np.array([0,0,0]))
    # Slightly rotate right_hand for salute
    salute_pose['right_hand'] = create_transform(axis=np.array([0,0,1]), angle_rad=np.radians(15), translation=np.array([0,0,0]))

    # Build animation with 3 keyframes: rest -> salute -> rest
    anim = Animation(model)
    anim.add_keyframe(0.0, rest_pose)
    anim.add_keyframe(2.0, salute_pose)  # salute pose at 2 sec
    anim.add_keyframe(4.0, rest_pose)    # back to rest at 4 sec

    start_time = time.time()

    print("Animating soldier salute and return to rest (simulated)...")

    # Simulate animation playback for 5 seconds
    while True:
        current_time = time.time()
        elapsed = current_time - start_time
        anim.play(start_time, current_time)

        # For demonstration, print out the right_arm global transform
        right_arm_global = model.bones['right_arm'].get_global_transform()
        print(f"Time {elapsed:.2f}s - right_arm global transform:\n{right_arm_global}\n")

        time.sleep(0.1)
        if elapsed > 5.0:
            break