# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0. See the LICENSE file in this repository.
#
# Feature encoding, decoding and endpoint correction adapted from NVIDIA MotionBricks
# (https://github.com/NVlabs/motionbricks), Apache-2.0. See NOTICE.

import functools
import math

import torch
import torch.nn.functional as F


FPS = 30.0
CONTEXT_FRAMES = 4
TARGET_FRAMES = 4
BOUNDARY_FRAMES = 3
FRAMES_PER_TOKEN = 4
# Heading angle of +X, which the last context pose is turned to face.
ALIGNED_HEADING = math.pi / 2

# SOMA30 joint names in model order; the root is Hips.
SOMA30_JOINTS = (
    "Hips", "Spine1", "Spine2", "Chest", "Neck1", "Neck2", "Head", "Jaw", "LeftEye", "RightEye",
    "LeftShoulder", "LeftArm", "LeftForeArm", "LeftHand", "LeftHandThumbEnd", "LeftHandMiddleEnd",
    "RightShoulder", "RightArm", "RightForeArm", "RightHand", "RightHandThumbEnd", "RightHandMiddleEnd",
    "LeftLeg", "LeftShin", "LeftFoot", "LeftToeBase", "RightLeg", "RightShin", "RightFoot", "RightToeBase",
)  # fmt: skip
SOMA30_PARENTS = (
    -1, 0, 1, 2, 3, 4, 5, 6, 6, 6, 3, 10, 11, 12, 13, 13, 3, 16, 17, 18, 19, 19, 0, 22, 23, 24, 0, 26, 27, 28,
)  # fmt: skip
# SOMA77 index of each SOMA30 joint.
SOMA77_TO_SOMA30 = (
    0, 1, 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 18, 28, 39, 40, 41, 42, 46, 56, 67, 68, 69, 70, 72, 73, 74, 75,
)  # fmt: skip
RIGHT_LEG, LEFT_LEG = SOMA30_JOINTS.index("RightLeg"), SOMA30_JOINTS.index("LeftLeg")
FOOT_CONTACT_JOINTS = ("LeftFoot", "LeftToeBase", "RightFoot", "RightToeBase")


@functools.cache
def soma77_index(device: torch.device) -> torch.Tensor:
    return torch.tensor(SOMA77_TO_SOMA30, device=device)


def on_ground(positions: torch.Tensor) -> torch.Tensor:
    """`(..., 3)` positions with their height set to 0."""
    return torch.stack([positions[..., 0], torch.zeros_like(positions[..., 1]), positions[..., 2]], dim=-1)


@functools.cache
def following_frames(num_frames: int, device: torch.device) -> torch.Tensor:
    """Index of the frame after each guide frame; the last context frame and the last frame are their own."""
    following = torch.arange(1, num_frames + 1).clamp(max=num_frames - 1)
    following[CONTEXT_FRAMES - 1] = CONTEXT_FRAMES - 1
    return following.to(device)


def soma30(values: torch.Tensor, joint_dim: int) -> torch.Tensor:
    """Select the SOMA30 joints of SOMA30 or SOMA77 values along `joint_dim`."""
    count = values.shape[joint_dim]
    if count == len(SOMA30_JOINTS):
        return values
    if count == 77:
        return values.index_select(joint_dim, soma77_index(values.device))
    raise ValueError(f"Expected SOMA30 or SOMA77 joints, got {count}")


def yaw_matrix(angle: torch.Tensor) -> torch.Tensor:
    """Rotation about +Y by `angle` radians, `(..., 3, 3)`."""
    c, s = torch.cos(angle), torch.sin(angle)
    zero, one = torch.zeros_like(angle), torch.ones_like(angle)
    return torch.stack([c, zero, s, zero, one, zero, -s, zero, c], dim=-1).unflatten(-1, (3, 3))


def heading_from_legs(right: torch.Tensor, left: torch.Tensor) -> torch.Tensor:
    """Heading angle about +Y of the forward direction, from the right and left leg joint positions."""
    across = right - left
    length = torch.hypot(across[..., 0], across[..., 2])
    angle = torch.atan2(across[..., 2], -across[..., 0])
    return torch.where(length > 1e-6, angle, torch.zeros_like(angle))


def place_target(
    positions: torch.Tensor,
    rotations: torch.Tensor,
    anchor: torch.Tensor,
    offset: torch.Tensor,
    turn: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Move and turn a motion so its first pose stands relative to an anchor pose.

    Args:
        positions: `(batch_size, frames, joints, 3)` SOMA30 or SOMA77 joint positions of the motion.
        rotations: `(batch_size, frames, joints, 3, 3)` global joint rotations of the motion.
        anchor: `(batch_size, 30, 3)` SOMA30 joint positions of the anchor pose.
        offset: `(batch_size, 2)` distance `(forward, left)` in meters from the anchor's ground point, along the
            direction the anchor faces.
        turn: `(batch_size,)` heading of the motion's first pose relative to the anchor's, in radians, positive to
            the left.

    Returns:
        The moved `positions` and `rotations`.
    """
    first = soma30(positions[:, 0], 1)
    anchor_angle = heading_from_legs(anchor[:, RIGHT_LEG], anchor[:, LEFT_LEG])
    first_angle = heading_from_legs(first[:, RIGHT_LEG], first[:, LEFT_LEG])
    zero = torch.zeros_like(anchor_angle)
    forward = torch.stack([torch.sin(anchor_angle), zero, torch.cos(anchor_angle)], dim=-1)
    left = torch.stack([torch.cos(anchor_angle), zero, -torch.sin(anchor_angle)], dim=-1)
    destination = on_ground(anchor[:, 0]) + offset[:, :1] * forward + offset[:, 1:] * left
    rotation = yaw_matrix(anchor_angle + turn - first_angle)[:, None, None]
    origin = on_ground(first[:, 0])[:, None, None]
    moved = (rotation @ (positions - origin)[..., None])[..., 0] + destination[:, None, None]
    return moved, rotation @ rotations


def rotation_6d(rotations: torch.Tensor) -> torch.Tensor:
    """The first two columns of `(..., 3, 3)` rotation matrices, `(..., 6)`."""
    return torch.cat([rotations[..., :, 0], rotations[..., :, 1]], dim=-1)


def rotation_from_6d(features: torch.Tensor) -> torch.Tensor:
    """Rotation matrices `(..., 3, 3)` from two `(..., 6)` column vectors, orthonormalized."""
    x = F.normalize(features[..., :3], dim=-1)
    z = F.normalize(torch.linalg.cross(x, features[..., 3:]), dim=-1)
    y = torch.linalg.cross(z, x)
    return torch.stack([x, y, z], dim=-1)


def encode_guides(positions: torch.Tensor, rotations: torch.Tensor):
    """Encode the context and target poses as the model's guide features.

    Args:
        positions: `(batch_size, 8, 30, 3)` SOMA30 joint positions, four context poses then four target poses, in
            meters with +Y up.
        rotations: `(batch_size, 8, 30, 3, 3)` SOMA30 global joint rotations, in the convention where the identity is
            the neutral pose of the skeleton.

    Returns:
        `(global_root, local_root, pose, alignment)`: `(batch_size, 8, 5)` root position and heading cosine and sine,
        `(batch_size, 8, 4)` heading rate, ground velocity and height, `(batch_size, 8, 267)` root-relative joint
        positions and 6D rotations, and the `(origin, yaw)` that maps aligned poses back to the input space.
    """
    last = CONTEXT_FRAMES - 1
    origin = on_ground(positions[:, last, 0])
    current = heading_from_legs(positions[:, last, RIGHT_LEG], positions[:, last, LEFT_LEG])
    yaw = current - ALIGNED_HEADING
    inverse = yaw_matrix(yaw).transpose(-1, -2)[:, None, None]
    aligned = (inverse @ (positions - origin[:, None, None])[..., None])[..., 0]
    aligned_rotations = inverse @ rotations

    num_frames = positions.shape[1]
    following = following_frames(num_frames, positions.device)
    hips, next_hips = aligned[:, :, 0], aligned[:, following, 0]
    angle = heading_from_legs(aligned[:, :, RIGHT_LEG], aligned[:, :, LEFT_LEG])
    next_angle = heading_from_legs(aligned[:, following, RIGHT_LEG], aligned[:, following, LEFT_LEG])
    delta = torch.atan2(torch.sin(next_angle - angle), torch.cos(next_angle - angle))

    global_root = torch.cat([hips, torch.cos(angle)[..., None], torch.sin(angle)[..., None]], dim=-1)
    velocity = (next_hips - hips) * FPS
    local_root = torch.stack([delta * FPS, velocity[..., 0], velocity[..., 2], hips[..., 1]], dim=-1)
    ground = on_ground(hips)
    joint_positions = (aligned[:, :, 1:] - ground[:, :, None]).flatten(2)
    pose = torch.cat([joint_positions, rotation_6d(aligned_rotations).flatten(2)], dim=-1)
    return global_root, local_root, pose, (origin, yaw)


def guide_masks(num_targets: torch.Tensor):
    """`(has_global_root, has_local_root, has_pose)` masks of the guide frames, `(batch_size, 8)` each.

    The context poses and the first `num_targets` target poses are given. The local root, a derivative, is unknown
    at the last context pose and the last target pose.
    """
    index = torch.arange(CONTEXT_FRAMES + TARGET_FRAMES, device=num_targets.device)[None]
    given = index < CONTEXT_FRAMES + num_targets[:, None]
    last = (index == CONTEXT_FRAMES - 1) | (index == CONTEXT_FRAMES + num_targets[:, None] - 1)
    return given, given & ~last, given


def root_to_local(root: torch.Tensor, num_tokens: torch.Tensor) -> torch.Tensor:
    """Local root `(heading rate, x velocity, z velocity, height)` of a `(batch_size, num_frames, 5)` global root.

    The last frame of the motion repeats the derivatives of the frame before it.
    """
    angle = torch.atan2(root[..., 4], root[..., 3])
    cos, sin = torch.cos(angle), torch.sin(angle)
    rate = torch.atan2(
        sin[:, 1:] * cos[:, :-1] - cos[:, 1:] * sin[:, :-1], cos[:, 1:] * cos[:, :-1] + sin[:, 1:] * sin[:, :-1]
    )
    rate = F.pad(rate * FPS, (0, 1))
    velocity = F.pad((root[:, 1:, :3] - root[:, :-1, :3]) * FPS, (0, 0, 0, 1))
    local = torch.cat([rate[..., None], velocity[..., 0:3:2], root[..., 1:2]], dim=-1)
    batch = torch.arange(root.shape[0], device=root.device)
    last = num_tokens * FRAMES_PER_TOKEN - 1
    local[batch, last] = local[batch, last - 1]
    local[batch, last, 3] = root[batch, last, 1]
    return local


def integrate_local_root(local_root: torch.Tensor) -> torch.Tensor:
    """Global root `(x, y, z, heading cosine, heading sine)` from a local root, starting at the origin."""
    rate, vx, vz, height = local_root.unbind(-1)
    angle = F.pad(torch.cumsum(rate[:, :-1] / FPS, dim=1), (1, 0))
    x = F.pad(torch.cumsum(vx[:, :-1] / FPS, dim=1), (1, 0))
    z = F.pad(torch.cumsum(vz[:, :-1] / FPS, dim=1), (1, 0))
    return torch.stack([x, height, z, torch.cos(angle), torch.sin(angle)], dim=-1)


def select_root(
    decoded_root: torch.Tensor,
    predicted_root: torch.Tensor,
    target_height: torch.Tensor,
    has_target: torch.Tensor,
    num_tokens: torch.Tensor,
) -> torch.Tensor:
    """Keep the decoded root, unless its height misses the targets by over 5 cm and the root model's is closer."""
    frames = (
        num_tokens[:, None] * FRAMES_PER_TOKEN - TARGET_FRAMES + torch.arange(TARGET_FRAMES, device=num_tokens.device)
    )
    batch = torch.arange(decoded_root.shape[0], device=decoded_root.device)[:, None]
    weight = has_target.to(decoded_root.dtype)
    count = weight.sum(-1)
    decoded_error = ((decoded_root[batch, frames, 1] - target_height).abs() * weight).sum(-1) / count
    predicted_error = ((predicted_root[batch, frames, 1] - target_height).abs() * weight).sum(-1) / count
    use_predicted = (predicted_error < decoded_error) & (decoded_error > 0.05)
    return torch.where(use_predicted[:, None, None], predicted_root, decoded_root)


def decode_poses(
    motion: torch.Tensor,
    num_tokens: torch.Tensor,
    alignment: tuple[torch.Tensor, torch.Tensor],
    neutral_joints: torch.Tensor,
    endpoints: dict | None = None,
):
    """World-space SOMA30 poses of the valid frames of each motion.

    Args:
        motion: `(batch_size, num_frames, 366)` features: global root, child joint positions, 6D global rotations,
            joint velocities and foot contacts, in the aligned space.
        num_tokens: `(batch_size,)` motion lengths in tokens; the first and last three frames are dropped.
        alignment: The `(origin, yaw)` from `encode_guides`.
        neutral_joints: `(30, 3)` joint positions of the neutral pose, which give the bone offsets.
        endpoints: The `endpoints` from `encode_guides`, to correct the endpoints with `correct_endpoints`.

    Returns:
        A list with one dict per motion: `posed_joints` `(frames, 30, 3)`, `global_rot_mats` `(frames, 30, 3, 3)`,
        `local_rot_mats` `(frames, 30, 3, 3)`, `root_positions` `(frames, 3)` and `foot_contacts` `(frames, 4)`.
    """
    origin, yaw = alignment
    rotation = yaw_matrix(yaw)
    joints = len(SOMA30_JOINTS)
    position_end = 5 + 3 * (joints - 1)
    global_rots = rotation[:, None, None] @ rotation_from_6d(
        motion[..., position_end : position_end + 6 * joints].unflatten(-1, (joints, 6))
    )
    roots = (rotation[:, None] @ motion[..., :3, None])[..., 0] + origin[:, None]
    contacts = motion[..., -len(FOOT_CONTACT_JOINTS) :].clamp(0, 1)
    local_rots = local_rotations(global_rots)
    if endpoints is not None:
        local_rots, roots = correct_endpoints(local_rots, roots, num_tokens, endpoints)
        global_rots = global_rotations(local_rots)
    positions = joint_positions(global_rots, roots, neutral_joints)
    outputs = []
    for i, count in enumerate((num_tokens * FRAMES_PER_TOKEN).tolist()):
        frames = slice(BOUNDARY_FRAMES, count - BOUNDARY_FRAMES)
        outputs.append(
            {
                "posed_joints": positions[i, frames],
                "global_rot_mats": global_rots[i, frames],
                "local_rot_mats": local_rots[i, frames],
                "root_positions": roots[i, frames],
                "foot_contacts": contacts[i, frames],
            }
        )
    return outputs


@functools.cache
def skeleton(device: torch.device) -> tuple[torch.Tensor, list[torch.Tensor], torch.Tensor]:
    """SOMA30 parent indices (the root its own parent), the non-root joints grouped by depth, and the `(30, 30)`
    matrix whose row `j` marks the non-root joints on the path from `j` to the root."""
    parents = torch.tensor(SOMA30_PARENTS).clamp(min=0)
    depth = [0] * len(SOMA30_PARENTS)
    ancestors = torch.zeros(len(SOMA30_PARENTS), len(SOMA30_PARENTS))
    for joint in range(1, len(SOMA30_PARENTS)):
        depth[joint] = depth[SOMA30_PARENTS[joint]] + 1
        ancestors[joint] = ancestors[SOMA30_PARENTS[joint]]
        ancestors[joint, joint] = 1
    levels = [torch.tensor([j for j, d in enumerate(depth) if d == level]) for level in range(1, max(depth) + 1)]
    return parents.to(device), [level.to(device) for level in levels], ancestors.to(device)


def local_rotations(global_rots: torch.Tensor) -> torch.Tensor:
    """Local rotations of `(..., 30, 3, 3)` global rotations; the root's local rotation is its global one."""
    parents, _, _ = skeleton(global_rots.device)
    local = global_rots[..., parents, :, :].transpose(-1, -2) @ global_rots
    return torch.cat([global_rots[..., :1, :, :], local[..., 1:, :, :]], dim=-3)


def joint_positions(global_rots: torch.Tensor, root_positions: torch.Tensor, neutral_joints: torch.Tensor):
    """Joint positions from global rotations and a root trajectory, with the bone offsets of the neutral pose."""
    parents, _, ancestors = skeleton(global_rots.device)
    offsets = (neutral_joints - neutral_joints[parents]).to(global_rots.dtype)
    bones = (global_rots[..., parents, :, :] @ offsets[:, :, None])[..., 0]
    return root_positions[..., None, :] + torch.einsum("jk,...kc->...jc", ancestors.to(bones.dtype), bones)


def global_rotations(local_rots: torch.Tensor) -> torch.Tensor:
    """Global rotations of `(..., 30, 3, 3)` local rotations."""
    parents, levels, _ = skeleton(local_rots.device)
    global_rots = local_rots.clone()
    for joints in levels:
        global_rots[..., joints, :, :] = global_rots[..., parents[joints], :, :] @ local_rots[..., joints, :, :]
    return global_rots


def axis_angle(rotations: torch.Tensor) -> torch.Tensor:
    """Axis-angle vectors of `(..., 3, 3)` rotation matrices."""
    quaternion = matrix_to_quaternion(rotations)
    quaternion = quaternion * torch.where(quaternion[..., :1] < 0, -1.0, 1.0)
    vector = quaternion[..., 1:]
    sin_half = vector.norm(dim=-1, keepdim=True)
    angle = 2 * torch.atan2(sin_half, quaternion[..., :1])
    scale = torch.where(sin_half > 1e-8, angle / sin_half.clamp(min=1e-8), torch.full_like(angle, 2.0))
    return vector * scale


def axis_angle_matrix(vector: torch.Tensor) -> torch.Tensor:
    angle = vector.norm(dim=-1, keepdim=True)[..., None]
    axis = F.normalize(vector, dim=-1, eps=1e-12)
    x, y, z = axis.unbind(-1)
    zero = torch.zeros_like(x)
    skew = torch.stack([zero, -z, y, z, zero, -x, -y, x, zero], dim=-1).unflatten(-1, (3, 3))
    eye = torch.eye(3, dtype=vector.dtype, device=vector.device).expand_as(skew)
    return eye + torch.sin(angle) * skew + (1 - torch.cos(angle)) * (skew @ skew)


def matrix_to_quaternion(rotations: torch.Tensor) -> torch.Tensor:
    """`(w, x, y, z)` unit quaternions of `(..., 3, 3)` rotation matrices."""
    m = rotations
    trace = m[..., 0, 0] + m[..., 1, 1] + m[..., 2, 2]
    candidates = torch.stack(
        [
            torch.stack(
                [1 + trace, m[..., 2, 1] - m[..., 1, 2], m[..., 0, 2] - m[..., 2, 0], m[..., 1, 0] - m[..., 0, 1]], -1
            ),
            torch.stack(
                [
                    m[..., 2, 1] - m[..., 1, 2],
                    1 + m[..., 0, 0] - m[..., 1, 1] - m[..., 2, 2],
                    m[..., 0, 1] + m[..., 1, 0],
                    m[..., 0, 2] + m[..., 2, 0],
                ],
                -1,
            ),
            torch.stack(
                [
                    m[..., 0, 2] - m[..., 2, 0],
                    m[..., 0, 1] + m[..., 1, 0],
                    1 - m[..., 0, 0] + m[..., 1, 1] - m[..., 2, 2],
                    m[..., 1, 2] + m[..., 2, 1],
                ],
                -1,
            ),
            torch.stack(
                [
                    m[..., 1, 0] - m[..., 0, 1],
                    m[..., 0, 2] + m[..., 2, 0],
                    m[..., 1, 2] + m[..., 2, 1],
                    1 - m[..., 0, 0] - m[..., 1, 1] + m[..., 2, 2],
                ],
                -1,
            ),
        ],
        dim=-2,
    )
    best = torch.stack([trace, m[..., 0, 0], m[..., 1, 1], m[..., 2, 2]], -1).argmax(-1)
    quaternion = torch.gather(candidates, -2, best[..., None, None].expand(*best.shape, 1, 4))[..., 0, :]
    return F.normalize(quaternion, dim=-1)


def scaled_rotation(rotations: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
    """The rotations scaled by `weight` along their geodesic from the identity, as slerp(identity, q, weight)."""
    return axis_angle_matrix(axis_angle(rotations) * weight[..., None])


def correct_endpoints(
    local_rots: torch.Tensor, roots: torch.Tensor, num_tokens: torch.Tensor, endpoints: dict
) -> tuple[torch.Tensor, torch.Tensor]:
    """Blend corrections over each motion so its first valid frame is the last context pose and its last valid frame
    the first target pose.

    Args:
        local_rots: `(batch_size, num_frames, 30, 3, 3)` local rotations.
        roots: `(batch_size, num_frames, 3)` root positions.
        num_tokens: `(batch_size,)` motion lengths in tokens.
        endpoints: `context_local` and `target_local` `(batch_size, 30, 3, 3)` local rotations of the last context
            pose and the first target pose, `context_root` and `target_root` `(batch_size, 3)` their root positions.

    Returns:
        The corrected local rotations and root positions.
    """
    batch = torch.arange(len(local_rots), device=local_rots.device)
    first = BOUNDARY_FRAMES
    last = num_tokens * FRAMES_PER_TOKEN - BOUNDARY_FRAMES - 1
    frames = torch.arange(local_rots.shape[1], dtype=local_rots.dtype, device=local_rots.device)
    end_weight = ((frames - first) / (last - first).to(local_rots.dtype)[:, None]).clamp(0, 1)
    start_weight = 1 - end_weight
    ends = torch.stack([local_rots[batch, first], local_rots[batch, last]], dim=1)
    deltas = ends.transpose(-1, -2) @ torch.stack([endpoints["context_local"], endpoints["target_local"]], dim=1)
    weights = torch.stack([start_weight, end_weight], dim=1)[..., None]
    start, end = scaled_rotation(deltas[:, :, None], weights).unbind(1)
    local_rots = local_rots @ start @ end
    roots = (
        roots
        + (endpoints["context_root"] - roots[batch, first])[:, None] * start_weight[..., None]
        + (endpoints["target_root"] - roots[batch, last])[:, None] * end_weight[..., None]
    )
    return local_rots, roots