# Copyright (C) 2026 Embedl AB

# This module's whole job is surgery on private torch and upstream-model
# internals, so the private-member-access lint is disabled file-wide.
# ruff: noqa: SLF001

"""Export utilities for the SAM 3D Body tutorial (``sam_3d_body.py``).

This module is one half of the "SAM 3D Body on TensorRT" tutorial. The tutorial
itself lives in ``sam_3d_body.py``, which imports this module by name; to run
it, download both files into the same directory.

The decoder + MHR head of ``facebookresearch/sam-3d-body`` cannot be exported
as-is; this module contains the graph surgery that fixes that, so the tutorial
can focus on the deployment story.

The tutorial's ``apply_export_patches(model)`` binds the ONNX-safe method
patches defined here — ``_safe_compact_cont_to_model_params_body``,
``_safe_compact_cont_to_model_params_hand``,
``_camera_encoder_fwd_nodownsample``, ``_patched_embed_keypoints``,
``_patched_mhr_head_forward``, ``_patched_perspective_projection``, and
``_patched_replace_hands_in_pose`` — onto the model.

``flatten_rig_for_export(decoder, head_pose, mhr_path, args)``
    Replace the TorchScript MHR body rig with a flat aten FX graph that
    the ONNX exporter can trace through (self-checked on probe inputs).

``register_ts_exporter_symbolics()``
    Register ONNX symbolics for the aten ops the flattened rig contains
    but the TorchScript exporter does not know.

Importing this module requires the ``sam-3d-body`` repository on ``sys.path``;
the tutorial clones it next to itself and inserts the path before importing.
"""

import contextlib
import math

import numpy as np  # type: ignore[import-not-found]
import numpy.typing as npt  # type: ignore[import-not-found]
import torch
import torch.onnx.symbolic_helper as sym_help
from sam_3d_body.models.modules import (  # type: ignore[import-not-found]
    mhr_utils as MU,
)
from sam_3d_body.models.modules import rot6d_to_rotmat
from sam_3d_body.models.modules.geometry_utils import (  # type: ignore[import-not-found]
    perspective_projection,
)
from sam_3d_body.models.modules.mhr_utils import (  # type: ignore[import-not-found]
    mhr_param_hand_mask,
)
from torch import nn
from torch.func import functionalize
from torch.fx.experimental.proxy_tensor import make_fx

# not re-exported by the torch.onnx stubs, but public API
from torch.onnx import (  # type: ignore[attr-defined]
    register_custom_op_symbolic,
)
from torch.onnx._internal.torchscript_exporter import _type_utils
from torch.onnx._internal.torchscript_exporter.jit_utils import GraphContext

__all__ = [
    "flatten_rig_for_export",
    "register_ts_exporter_symbolics",
]


def _densify_sparse_constants(gm: torch.fx.GraphModule) -> int:
    """Replace constant sparse_coo tensors in `gm` with dense buffers.

    The rig's pose-correctives layer builds a sparse weight from constant
    index/value buffers; ONNX has no sparse tensors, so fold it to dense (the
    downstream ``aten.mm`` is dense-compatible).

    :param gm:
        The flattened rig graph, modified in place.
    :returns:
        The number of sparse constructors replaced.
    """
    count = 0
    for node in list(gm.graph.nodes):
        if node.op == "call_function" and "sparse_coo" in str(node.target):
            _sd, _dd, shape, idx_n, val_n = node.args
            assert idx_n.op == "get_attr", "sparse index is not a constant"
            assert val_n.op == "get_attr", "sparse values are not constants"
            idx = getattr(gm, idx_n.target)
            val = getattr(gm, val_n.target)
            dense = torch.sparse_coo_tensor(idx, val, shape).to_dense()
            name = f"_densified_sparse_{count}"
            gm.register_buffer(name, dense)
            with gm.graph.inserting_before(node):
                new_node = gm.graph.get_attr(name)
            node.replace_all_uses_with(new_node)
            gm.graph.erase_node(node)
            count += 1
    if count:
        gm.graph.lint()  # type: ignore[no-untyped-call]
        gm.recompile()
    return count


def _replace_functional_copy(gm: torch.fx.GraphModule) -> int:
    """Rewrite ``aten.copy`` and dynamic expands to static broadcasts.

    ``aten.copy(self, src)`` becomes an expand/cast of `src` with the static
    shape recorded at trace time: the ONNX symbolic route (Expand + Shape)
    confuses ORT/TRT shape inference when ``self`` is a strided view, while the
    FX metadata carries the exact shapes, so bake them in. ``aten.expand`` with
    ``-1`` sizes gets the same static-shape treatment.

    :param gm:
        The flattened rig graph, modified in place.
    :returns:
        The number of nodes rewritten.
    """
    aten = torch.ops.aten
    count = 0
    for node in list(gm.graph.nodes):
        if node.op != "call_function":
            continue
        if node.target in (aten.copy.default, aten.copy_.default):
            src_n = node.args[1]
            val = node.meta.get("val")
            shape, dtype = list(val.shape), val.dtype
            src_val = (
                src_n.meta.get("val")
                if isinstance(src_n, torch.fx.Node)
                else None
            )
            with gm.graph.inserting_before(node):
                new = src_n
                if src_val is not None and src_val.dtype != dtype:
                    new = gm.graph.call_function(
                        aten._to_copy.default, (new,), {"dtype": dtype}
                    )
                new = gm.graph.call_function(
                    aten.broadcast_to.default, (new, shape)
                )
            node.replace_all_uses_with(new)
            gm.graph.erase_node(node)
            count += 1
        elif node.target in (aten.expand.default, aten.broadcast_to.default):
            sizes = node.args[1]
            if isinstance(sizes, (list, tuple)) and any(
                isinstance(s, int) and s == -1 for s in sizes
            ):
                shape = list(node.meta["val"].shape)
                node.args = (node.args[0], shape)
                count += 1
    if count:
        gm.graph.lint()  # type: ignore[no-untyped-call]
        gm.recompile()
    return count


def _replace_index_copy_with_scatter(
    gm: torch.fx.GraphModule, sample_args: tuple[torch.Tensor, ...]
) -> int:
    """Rewrite ``aten.index_copy``/``index_add`` to constant scatters.

    The exporter's ``index_copy`` lowering emits a dynamic
    ConstantOfShape/Slice/Shape/Expand chain that fails ORT/TRT shape
    inference. The FK indices are constants, so run the graph once on
    `sample_args` to get their concrete values and bake the expanded scatter
    index in as a buffer.

    :param gm:
        The flattened rig graph, modified in place.
    :param sample_args:
        Inputs used to record every node's concrete value.
    :returns:
        The number of nodes rewritten.
    """
    aten = torch.ops.aten
    copy_targets = (aten.index_copy.default, aten.index_copy_.default)
    add_targets = (aten.index_add.default, aten.index_add_.default)
    targets = copy_targets + add_targets
    if not any(
        n.op == "call_function" and n.target in targets for n in gm.graph.nodes
    ):
        return 0
    env: dict[torch.fx.Node, torch.Tensor] = {}

    class Recorder(torch.fx.Interpreter):
        def run_node(self, n: torch.fx.Node) -> torch.Tensor:
            out = super().run_node(n)
            env[n] = out
            return out  # type: ignore[no-any-return]

    with torch.no_grad():
        Recorder(gm).run(*sample_args)
    count = 0
    for node in list(gm.graph.nodes):
        if node.op == "call_function" and node.target in targets:
            self_n, dim, index_n, src_n = node.args[:4]
            assert len(node.args) == 4, "index_add alpha != 1 not handled"
            index_val = env[index_n]
            src_val = env[src_n]
            dim_i = dim % src_val.ndim
            shape = [1] * src_val.ndim
            shape[dim_i] = -1
            expanded = index_val.reshape(shape).expand_as(src_val).contiguous()
            buf = f"_index_copy_scatter_idx_{count}"
            gm.register_buffer(buf, expanded)
            op = (
                aten.scatter.src
                if node.target in copy_targets
                else aten.scatter_add.default
            )
            with gm.graph.inserting_before(node):
                idx_node = gm.graph.get_attr(buf)
                new = gm.graph.call_function(
                    op, (self_n, dim_i, idx_node, src_n)
                )
            node.replace_all_uses_with(new)
            gm.graph.erase_node(node)
            count += 1
    gm.graph.lint()  # type: ignore[no-untyped-call]
    gm.recompile()
    return count


def _flatten_mhr_rig(
    head_pose: nn.Module,
    mhr_path: str,
    sample_args: tuple[torch.Tensor, ...],
) -> None:
    """Replace the TorchScript MHR rig with a flat aten FX graph.

    The ONNX TorchScript tracer cannot trace through a ``jit.load``-ed
    ScriptModule ("not part of the active trace"); ``make_fx`` executes the
    TorchScript interpreter with real tensors and records the dispatched aten
    ops, unrolling the rig's control flow. The result is self-checked against
    the original rig on five probe input sets before it is accepted.

    :param head_pose:
        The MHR head whose ``mhr`` attribute holds the rig; replaced in
        place.
    :param mhr_path:
        Path to the TorchScript ``mhr_model.pt`` file.
    :param sample_args:
        Real rig inputs recorded from a full decoder run.
    """
    rig = head_pose.mhr
    if isinstance(rig, torch.fx.GraphModule):
        return
    dev = sample_args[0].device
    # JIT executor optimizations would specialize the rig during the tracing
    # run below; disable them so make_fx records the plain aten ops
    with torch.jit.optimized_execution(False):  # type: ignore[attr-defined]
        fresh = torch.jit.load(mhr_path, map_location=dev)  # type: ignore[no-untyped-call]
        gen = torch.Generator(device=dev).manual_seed(0)
        generic_args = tuple(
            torch.randn(a.shape, generator=gen, device=dev, dtype=a.dtype)
            for a in sample_args
        )
        with torch.no_grad():
            # functionalize: the rig mutates select/slice views in place
            # (copy_ into select), which the ONNX exporter cannot express;
            # functionalization rewrites those to select/slice_scatter
            # the lambda gives functionalize a plain python callable
            # instead of the ScriptModule's dispatching __call__
            gm = make_fx(
                functionalize(lambda s, p, e: fresh(s, p, e)),  # noqa: PLW0108
                tracing_mode="real",
            )(*generic_args)
    _densify_sparse_constants(gm)
    _replace_functional_copy(gm)
    _replace_index_copy_with_scatter(gm, generic_args)
    with torch.no_grad():
        probes = {
            "generic": generic_args,
            "real sample": sample_args,
            "zeros": tuple(torch.zeros_like(a) for a in sample_args),
            "small": tuple(0.1 * torch.randn_like(a) for a in sample_args),
            "big": tuple(3 * torch.randn_like(a) for a in sample_args),
        }
        for tag, probe in probes.items():
            for w, g in zip(rig(*probe), gm(*probe), strict=True):  # type: ignore[operator]
                err = (w - g).abs().max().item()
                # outputs are in cm x100 scale; a wrong branch errs by >>1
                assert err < 1e-2, f"flattened rig mismatch on {tag}: {err}"
    head_pose.mhr = gm


class _RigArgRecorder(nn.Module):
    """Pass-through wrapper that records the rig's first call arguments."""

    def __init__(self, rig: nn.Module) -> None:
        super().__init__()
        self.rig = rig
        self.args: tuple[torch.Tensor, ...] | None = None

    def forward(self, *args: torch.Tensor) -> torch.Tensor:
        if self.args is None:
            self.args = tuple(a.detach().clone() for a in args)
        return self.rig(*args)  # type: ignore[no-any-return]


def flatten_rig_for_export(
    decoder: nn.Module,
    head_pose: nn.Module,
    mhr_path: str,
    args: tuple[torch.Tensor, ...],
) -> None:
    """Capture real rig inputs from one decoder run, then flatten the rig.

    Runs `decoder` once with a recorder wrapped around the rig to obtain
    realistic sample inputs, then replaces the TorchScript rig on `head_pose`
    with an export-friendly flat FX graph.

    :param decoder:
        A module whose forward drives `head_pose` (the tutorial's
        decoder wrapper); called once under ``torch.no_grad``.
    :param head_pose:
        The MHR head owning the rig, modified in place.
    :param mhr_path:
        Path to the TorchScript ``mhr_model.pt`` file.
    :param args:
        Positional inputs for one `decoder` forward pass.
    """
    recorder = _RigArgRecorder(head_pose.mhr)  # type: ignore[arg-type]
    head_pose.mhr = recorder
    with torch.no_grad():
        decoder(*args)
    head_pose.mhr = recorder.rig
    assert recorder.args is not None, "decoder never called the rig"
    _flatten_mhr_rig(head_pose, mhr_path, recorder.args)


# Upstream scatters three continuous-parameter blocks (3-DoF joints,
# 1-DoF joints, translation-like extras) into a zeros(133) tensor with
# constant index assignments (a broken index_put export); the indices
# form a complete permutation of 0..132, so gathering the concatenated
# blocks with the inverse permutation is identical.  Only the 3-DoF
# slot table carries information -- the 1-DoF slots are the sorted
# complement and the extras are the tail block.

_BODY_IDX3 = np.array(
    [
        [0, 2, 4],
        [6, 8, 10],
        [12, 13, 14],
        [15, 16, 17],
        [18, 19, 20],
        [21, 22, 23],
        [24, 25, 26],
        [27, 28, 29],
        [34, 35, 36],
        [37, 38, 39],
        [44, 45, 46],
        [53, 54, 55],
        [64, 65, 66],
        [85, 69, 73],
        [86, 70, 79],
        [87, 71, 82],
        [88, 72, 76],
        [91, 92, 93],
        [112, 96, 100],
        [113, 97, 106],
        [114, 98, 109],
        [115, 99, 103],
        [130, 131, 132],
    ]
).flatten()
_BODY_IDX_TAIL = np.arange(124, 130)
_BODY_IDX1 = np.setdiff1d(
    np.arange(133), np.concatenate([_BODY_IDX3, _BODY_IDX_TAIL])
)
_BODY_INV_PERM = np.argsort(
    np.concatenate([_BODY_IDX3, _BODY_IDX1, _BODY_IDX_TAIL])
)
assert np.array_equal(
    np.sort(np.concatenate([_BODY_IDX3, _BODY_IDX1, _BODY_IDX_TAIL])),
    np.arange(133),
), "body slot tables are not a permutation of 0..132"


def _safe_compact_cont_to_model_params_body(
    body_pose_cont: torch.Tensor,
) -> torch.Tensor:
    """``compact_cont_to_model_params_body`` via inverse-perm gather.

    :param body_pose_cont:
        Continuous body-pose block, shape ``(..., 236)``.
    :returns:
        Euler body-pose parameters, shape ``(..., 133)``.
    """
    n3, n1 = len(_BODY_IDX3), len(_BODY_IDX1)
    assert body_pose_cont.shape[-1] == 2 * n3 + 2 * n1 + len(_BODY_IDX_TAIL)
    c3 = body_pose_cont[..., : 2 * n3].unflatten(-1, (-1, 6))  # type: ignore[no-untyped-call]
    p3 = MU.batchXYZfrom6D(c3).flatten(-2, -1)
    c1 = body_pose_cont[..., 2 * n3 : 2 * n3 + 2 * n1].unflatten(-1, (-1, 2))  # type: ignore[no-untyped-call]
    p1 = torch.atan2(c1[..., -2], c1[..., -1])
    pt = body_pose_cont[..., 2 * n3 + 2 * n1 :]
    inv = torch.from_numpy(_BODY_INV_PERM).to(body_pose_cont.device)
    return torch.cat([p3, p1, pt], dim=-1)[..., inv]


def _dof_indices(
    dofs: npt.NDArray[np.int64],
    match: tuple[int, ...],
    widths: npt.NDArray[np.int64],
) -> npt.NDArray[np.intp]:
    """Return the flat slot indices of joints whose DoF count matches.

    :param dofs:
        Per-joint degree-of-freedom counts.
    :param match:
        DoF counts to select.
    :param widths:
        Per-joint slot widths (`dofs` for parameter slots, ``2 * dofs``
        for continuous-representation slots).
    :returns:
        The selected flat slot indices, in slot order.
    """
    mask = np.repeat([k in match for k in dofs], widths)
    return np.where(mask)[0]


# The two DoF masks partition the 27 hand-pose outputs completely, so
# the same inverse-permutation-gather trick applies.
_HAND_DOFS = np.array([3, 1, 1, 3, 1, 1, 3, 1, 1, 3, 1, 1, 2, 3, 1, 1])
_HAND_CONT3 = _dof_indices(_HAND_DOFS, (3,), 2 * _HAND_DOFS)
_HAND_CONT1 = _dof_indices(_HAND_DOFS, (1, 2), 2 * _HAND_DOFS)
_HAND_INV_PERM = np.argsort(
    np.concatenate(
        [
            _dof_indices(_HAND_DOFS, (3,), _HAND_DOFS),
            _dof_indices(_HAND_DOFS, (1, 2), _HAND_DOFS),
        ]
    )
)


def _safe_compact_cont_to_model_params_hand(
    hand_cont: torch.Tensor,
) -> torch.Tensor:
    """``compact_cont_to_model_params_hand`` via inverse-perm gather.

    :param hand_cont:
        Continuous hand-pose block, shape ``(..., 54)``.
    :returns:
        Euler hand-pose parameters, shape ``(..., 27)``.
    """
    assert hand_cont.shape[-1] == 2 * int(_HAND_DOFS.sum())
    dev = hand_cont.device
    c3 = hand_cont[..., torch.from_numpy(_HAND_CONT3).to(dev)]
    p3 = MU.batchXYZfrom6D(c3.unflatten(-1, (-1, 6))).flatten(-2, -1)  # type: ignore[no-untyped-call]
    c1 = hand_cont[..., torch.from_numpy(_HAND_CONT1).to(dev)]
    c1 = c1.unflatten(-1, (-1, 2))  # type: ignore[no-untyped-call]
    p1 = torch.atan2(c1[..., -2], c1[..., -1])
    inv = torch.from_numpy(_HAND_INV_PERM).to(dev)
    return torch.cat([p3, p1], dim=-1)[..., inv]


def _safe_rotmat_to_euler_zyx(
    rotmat: torch.Tensor, epsilon: float = 1e-7
) -> torch.Tensor:
    """Branchless drop-in for ``roma.rotmat_to_euler("ZYX", R)``.

    roma's implementation builds the quaternion and handles gimbal-lock and
    angle-wrap cases with data-dependent ``nonzero()``-indexed assignments,
    whose element counts get baked into an ONNX trace (the graph then crashes
    for poses that take a different case). Same algorithm, but every case is a
    ``torch.where`` and the quaternion candidate is picked with a shape-static
    gather, so the export is valid for every input.

    :param rotmat:
        Batch of rotation matrices, shape ``(B, 3, 3)``.
    :param epsilon:
        Gimbal-lock detection threshold on the middle angle.
    :returns:
        ZYX Euler angles, shape ``(B, 3)``.
    """
    B = rotmat.shape[0]
    m = rotmat.reshape(B, 9)
    m00, m01, m02, m10, m11, m12, m20, m21, m22 = [m[:, i] for i in range(9)]
    # rotmat -> unit quaternion (xyzw): all 4 candidates, pick the stablest
    t0 = 1 + m00 - m11 - m22
    t1 = 1 - m00 + m11 - m22
    t2 = 1 - m00 - m11 + m22
    t3 = 1 + m00 + m11 + m22
    c0 = torch.stack([t0, m10 + m01, m20 + m02, m21 - m12], 1)
    c1 = torch.stack([m10 + m01, t1, m21 + m12, m02 - m20], 1)
    c2 = torch.stack([m20 + m02, m21 + m12, t2, m10 - m01], 1)
    c3 = torch.stack([m21 - m12, m02 - m20, m10 - m01, t3], 1)
    cands = torch.stack([c0, c1, c2, c3], 1)  # (B, 4, 4)
    best = torch.stack([t0, t1, t2, t3], 1).argmax(1)  # (B,)
    quat = torch.gather(cands, 1, best.view(B, 1, 1).expand(B, 1, 4))[:, 0]
    quat = quat / quat.norm(dim=1, keepdim=True)
    qx, qy, qz, qw = quat[:, 0], quat[:, 1], quat[:, 2], quat[:, 3]
    # unitquat -> euler, intrinsic ZYX == extrinsic xyz (roma's algorithm,
    # scipy-style): i=0, j=1, k=2, sign=+1, asymmetric, angle order reversed
    a, b = qw - qy, qx + qz
    c, d = qy + qw, qz - qx
    ang1 = 2 * torch.atan2(torch.hypot(c, d), torch.hypot(a, b))
    case1 = ang1.abs() <= epsilon
    case2 = (ang1 - math.pi).abs() <= epsilon
    half_sum = torch.atan2(b, a)
    half_diff = torch.atan2(d, c)
    # intrinsic: angle_first slot is index 2, angle_third slot is index 0
    ang_first = half_sum - half_diff
    ang_third = torch.where(
        case1,
        2 * half_sum,
        torch.where(case2, 2 * half_diff, half_sum + half_diff),
    )
    ang_first = torch.where(
        case1 | case2, torch.zeros_like(ang_first), ang_first
    )
    ang1 = ang1 - math.pi / 2

    def wrap(x: torch.Tensor) -> torch.Tensor:
        return torch.where(
            x < -math.pi,
            x + 2 * math.pi,
            torch.where(x > math.pi, x - 2 * math.pi, x),
        )

    return torch.stack([wrap(ang_third), wrap(ang1), wrap(ang_first)], 1)


def _camera_encoder_fwd_nodownsample(
    self: nn.Module, img_embeddings: torch.Tensor, rays_ds: torch.Tensor
) -> torch.Tensor:
    """``CameraEncoder.forward`` without the antialiased interpolate.

    TensorRT's Resize has no antialias, so the tutorial downsamples the rays on
    the host during preprocessing and passes them in at feature resolution.

    :param self:
        The ``CameraEncoder`` instance this forward is bound to.
    :param img_embeddings:
        Backbone features, shape ``(B, D, h, w)``.
    :param rays_ds:
        Pre-downsampled camera rays, shape ``(B, 2, h, w)``.
    :returns:
        Ray-conditioned features, shape ``(B, D, h, w)``.
    """
    B, _D, h, w = img_embeddings.shape
    rays = rays_ds.permute(0, 2, 3, 1).contiguous()  # [B, h, w, 2]
    rays = torch.cat([rays, torch.ones_like(rays[..., :1])], dim=-1)
    rays_embeddings = self.camera(pos=rays.reshape(B, -1, 3))  # type: ignore[operator]  # [B, h*w, 99]
    rays_embeddings = (
        rays_embeddings.reshape(B, h, w, -1).permute(0, 3, 1, 2).contiguous()
    )
    z = torch.cat([img_embeddings, rays_embeddings], dim=1)
    return self.norm(self.conv(z))  # type: ignore[no-any-return, operator]


def _patched_replace_hands_in_pose(
    self: nn.Module,
    full_pose_params: torch.Tensor,
    hand_pose_params: torch.Tensor,
) -> torch.Tensor:
    """``replace_hands_in_pose`` without advanced-index assignment.

    ``x[:, idx_tensor] = v`` exports through a broken index_put lowering that
    drops the index constant; scatter exports as a clean ScatterElements.

    :param self:
        The MHR head this method is bound to.
    :param full_pose_params:
        Full body-pose parameters, shape ``(B, 136)``.
    :param hand_pose_params:
        Compact hand-pose parameters for both hands.
    :returns:
        `full_pose_params` with both hand blocks replaced.
    """
    assert full_pose_params.shape[1] == 136
    left_hand_params, right_hand_params = torch.split(
        hand_pose_params,
        [self.num_hand_pose_comps, self.num_hand_pose_comps],  # type: ignore[list-item]
        dim=1,
    )
    left_model = _safe_compact_cont_to_model_params_hand(
        self.hand_pose_mean  # type: ignore[operator]
        + torch.einsum("da,ab->db", left_hand_params, self.hand_pose_comps)
    )
    right_model = _safe_compact_cont_to_model_params_hand(
        self.hand_pose_mean  # type: ignore[operator]
        + torch.einsum("da,ab->db", right_hand_params, self.hand_pose_comps)
    )
    B = full_pose_params.shape[0]
    left_idx = self.hand_joint_idxs_left.view(1, -1).expand(B, -1)  # type: ignore[operator]
    right_idx = self.hand_joint_idxs_right.view(1, -1).expand(B, -1)  # type: ignore[operator]
    full_pose_params = full_pose_params.scatter(1, left_idx, left_model)
    return full_pose_params.scatter(1, right_idx, right_model)


def _patched_mhr_head_forward(
    self: nn.Module,
    x: torch.Tensor,
    init_estimate: torch.Tensor | None = None,
    do_pcblend: bool = True,
    slim_keypoints: bool = False,  # noqa: ARG001 -- upstream passes it
) -> dict[str, torch.Tensor | None]:
    """``MHRHead.forward`` with ONNX-safe zeroing and axis flips.

    Mask multiplies replace advanced-index assignment; numerically identical to
    upstream.

    :param self:
        The MHR head this forward is bound to.
    :param x:
        Decoder pose tokens, shape ``(B, token_dim)``.
    :param init_estimate:
        Optional initial estimate added to the projected parameters.
    :param do_pcblend:
        Forwarded to the rig.
    :param slim_keypoints:
        Unused; kept for signature compatibility.
    :returns:
        The upstream output dict (vertices, keypoints, pose blocks).
    """
    batch_size = x.shape[0]
    pred = self.proj(x)  # type: ignore[operator]
    if init_estimate is not None:
        pred = pred + init_estimate

    count = 6
    global_rot_6d = pred[:, :count]
    global_rot_rotmat = rot6d_to_rotmat(global_rot_6d)
    global_rot_euler = _safe_rotmat_to_euler_zyx(global_rot_rotmat)
    global_trans = torch.zeros_like(global_rot_euler)

    pred_pose_cont = pred[:, count : count + self.body_cont_dim]  # type: ignore[operator]
    count += self.body_cont_dim  # type: ignore[assignment, operator]
    pred_pose_euler = _safe_compact_cont_to_model_params_body(pred_pose_cont)
    # zero hands + jaw via a constant keep-mask built OUTSIDE tracing
    # (torch index assignment here would reintroduce broken index_put nodes)
    keep_np = np.ones(pred_pose_euler.shape[1], dtype=np.float32)
    keep_np[mhr_param_hand_mask.cpu().numpy()] = 0
    keep_np[-3:] = 0
    keep = torch.from_numpy(keep_np).to(device=pred.device, dtype=pred.dtype)
    pred_pose_euler = pred_pose_euler * keep

    pred_shape = pred[:, count : count + self.num_shape_comps]  # type: ignore[operator]
    count += self.num_shape_comps  # type: ignore[assignment, operator]
    pred_scale = pred[:, count : count + self.num_scale_comps]  # type: ignore[operator]
    count += self.num_scale_comps  # type: ignore[assignment, operator]
    pred_hand = pred[:, count : count + self.num_hand_comps * 2]  # type: ignore[operator]
    count += self.num_hand_comps * 2  # type: ignore[assignment, operator]
    pred_face = pred[:, count : count + self.num_face_comps] * 0  # type: ignore[operator]
    count += self.num_face_comps  # type: ignore[assignment, operator]

    output = self.mhr_forward(  # type: ignore[operator]
        global_trans=global_trans,
        global_rot=global_rot_euler,
        body_pose_params=pred_pose_euler,
        hand_pose_params=pred_hand,
        scale_params=pred_scale,
        shape_params=pred_shape,
        expr_params=pred_face,
        do_pcblend=do_pcblend,
        return_keypoints=True,
        return_joint_coords=True,
        return_model_params=True,
        return_joint_rotations=True,
    )
    verts, j3d, jcoords, mhr_model_params, joint_global_rots = output
    j3d = j3d[:, :70]

    flip_yz = torch.tensor(
        [1.0, -1.0, -1.0], device=pred.device, dtype=verts.dtype
    )
    verts = verts * flip_yz
    j3d = j3d * flip_yz
    jcoords = jcoords * flip_yz

    return {
        "pred_pose_raw": torch.cat([global_rot_6d, pred_pose_cont], dim=1),
        "pred_pose_rotmat": None,
        "global_rot": global_rot_euler,
        "body_pose": pred_pose_euler,
        "shape": pred_shape,
        "scale": pred_scale,
        "hand": pred_hand,
        "face": pred_face,
        "pred_keypoints_3d": j3d.reshape(batch_size, -1, 3),
        "pred_vertices": verts.reshape(batch_size, -1, 3),
        "pred_joint_coords": jcoords.reshape(batch_size, -1, 3),
        "joint_global_rots": joint_global_rots,
        "mhr_model_params": mhr_model_params,
    }


def _patched_perspective_projection(
    self: nn.Module,
    points_3d: torch.Tensor,
    pred_cam: torch.Tensor,
    bbox_center: torch.Tensor,
    bbox_size: torch.Tensor,
    img_size: torch.Tensor,
    cam_int: torch.Tensor,
    use_intrin_center: bool = False,
) -> dict[str, torch.Tensor]:
    """``PerspectiveHead.perspective_projection`` with a constant flip.

    The sign flip is a constant multiply instead of the in-place
    ``pred_cam[..., [0, 2]] *= -1``.

    :param self:
        The perspective head this method is bound to.
    :param points_3d:
        Predicted 3D points, shape ``(B, N, 3)``.
    :param pred_cam:
        Weak-perspective camera prediction, shape ``(B, 3)``.
    :param bbox_center:
        Person bbox center in image pixels, shape ``(B, 2)``.
    :param bbox_size:
        Person bbox size in image pixels, shape ``(B,)``.
    :param img_size:
        Original image size, shape ``(B, 2)``.
    :param cam_int:
        Camera intrinsics, shape ``(B, 3, 3)``.
    :param use_intrin_center:
        Center the bbox offset on the intrinsic principal point instead
        of the image center.
    :returns:
        2D keypoints, camera translation, focal length, and depths.
    """
    batch_size = points_3d.shape[0]
    flip = torch.tensor(
        [-1.0, 1.0, -1.0], device=pred_cam.device, dtype=pred_cam.dtype
    )
    pred_cam = pred_cam * flip
    s, tx, ty = pred_cam[:, 0], pred_cam[:, 1], pred_cam[:, 2]
    bs = bbox_size * s * self.default_scale_factor + 1e-8  # type: ignore[operator]
    focal_length = cam_int[:, 0, 0]
    tz = 2 * focal_length / bs
    if not use_intrin_center:
        cx = 2 * (bbox_center[:, 0] - (img_size[:, 0] / 2)) / bs
        cy = 2 * (bbox_center[:, 1] - (img_size[:, 1] / 2)) / bs
    else:
        cx = 2 * (bbox_center[:, 0] - (cam_int[:, 0, 2])) / bs
        cy = 2 * (bbox_center[:, 1] - (cam_int[:, 1, 2])) / bs
    pred_cam_t = torch.stack([tx + cx, ty + cy, tz], dim=-1)
    j3d_cam = points_3d + pred_cam_t.unsqueeze(1)
    j2d = perspective_projection(j3d_cam, cam_int)
    return {
        "pred_keypoints_2d": j2d.reshape(batch_size, -1, 2),
        "pred_cam_t": pred_cam_t,
        "focal_length": focal_length,
        "pred_keypoints_2d_depth": j3d_cam.reshape(batch_size, -1, 3)[:, :, 2],
    }


def _patched_embed_keypoints(
    self: nn.Module, points: torch.Tensor, labels: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """``PromptEncoder._embed_keypoints`` with ``where()`` selection.

    Replaces boolean-mask index assignment; same math for the label set used
    here.

    :param self:
        The prompt encoder this method is bound to.
    :param points:
        Keypoint prompt coordinates, shape ``(B, N, 2)``.
    :param labels:
        Keypoint prompt labels, shape ``(B, N)``.
    :returns:
        The point embeddings and the valid-point mask.
    """
    point_embedding = self.pe_layer._pe_encoding(points.to(torch.float))  # type: ignore[operator, union-attr]
    lab = labels.unsqueeze(-1)
    point_embedding = torch.where(
        lab == -2,
        self.invalid_point_embed.weight.to(point_embedding),  # type: ignore[arg-type, union-attr]
        point_embedding,
    )
    point_embedding = torch.where(
        lab == -1,
        self.not_a_point_embed.weight.to(point_embedding),  # type: ignore[arg-type, union-attr]
        point_embedding,
    )
    for i in range(self.num_body_joints):  # type: ignore[arg-type]
        point_embedding = point_embedding + (
            (lab == i).to(point_embedding) * self.point_embeddings[i].weight  # type: ignore[index, union-attr]
        )
    point_mask = labels > -2
    return point_embedding, point_mask


def register_ts_exporter_symbolics() -> None:
    """Register ONNX symbolics for aten ops the exporter doesn't know.

    The flattened rig graph contains aten ops without a TorchScript exporter
    symbolic. All shapes in this graph are static, so each one reduces to a
    simple constant-shape formulation.
    """

    # aten::hypot (rotation math) -- values are O(1) rotation components so
    # the naive sqrt(x^2+y^2) form is numerically fine
    def _hypot(
        g: GraphContext, self: torch.Value, other: torch.Value
    ) -> torch.Value:
        return g.op(  # type: ignore[no-any-return]
            "Sqrt",
            g.op("Add", g.op("Mul", self, self), g.op("Mul", other, other)),
        )

    register_custom_op_symbolic("aten::hypot", _hypot, 17)

    # aten::_to_copy comes from the make_fx-flattened rig (dtype casts for
    # the float64 FK section); ONNX-wise it is just a Cast
    def _to_copy(
        g: GraphContext,
        self: torch.Value,
        dtype: torch.Value | None = None,
        _layout: torch.Value | None = None,
        _device: torch.Value | None = None,
        _pin_memory: torch.Value | None = None,
        _non_blocking: torch.Value | None = None,
        _memory_format: torch.Value | None = None,
    ) -> torch.Value:
        if dtype is None or sym_help._is_none(dtype):
            return g.op("Identity", self)  # type: ignore[no-any-return]
        dtype_i = sym_help._get_const(dtype, "i", "dtype")  # type: ignore[no-untyped-call]
        onnx_ty = _type_utils.JitScalarType(dtype_i).onnx_type()
        return g.op("Cast", self, to_i=onnx_ty)  # type: ignore[no-any-return]

    register_custom_op_symbolic("aten::_to_copy", _to_copy, 17)

    # aten::mT (matrix transpose of the affine transform)
    def _mat_transpose(g: GraphContext, self: torch.Value) -> torch.Value:
        rank = sym_help._get_tensor_rank(self)
        assert rank is not None, "mT needs a known static rank"
        assert rank >= 2, "mT needs rank >= 2"
        perm = list(range(rank))
        perm[-2], perm[-1] = perm[-1], perm[-2]
        return g.op("Transpose", self, perm_i=perm)  # type: ignore[no-any-return]

    register_custom_op_symbolic("aten::mT", _mat_transpose, 17)

    # Functionalizing the rig introduces select_scatter / slice_scatter /
    # functional copy; with static shapes these are Slice+Concat / Expand.
    def _const_i64(g: GraphContext, vals: list[int]) -> torch.Value:
        return g.op("Constant", value_t=torch.tensor(vals, dtype=torch.int64))  # type: ignore[no-any-return]

    IMAX = 9223372036854775807

    def _slice_along(
        g: GraphContext, x: torch.Value, dim_i: int, start_i: int, end_i: int
    ) -> torch.Value:
        return g.op(  # type: ignore[no-any-return]
            "Slice",
            x,
            _const_i64(g, [start_i]),
            _const_i64(g, [end_i]),
            _const_i64(g, [dim_i]),
        )

    def _select_scatter(
        g: GraphContext,
        self: torch.Value,
        src: torch.Value,
        dim: torch.Value,
        index: torch.Value,
    ) -> torch.Value:
        # self with src written at `index` along `dim`:
        # concat(self[..:idx], unsqueeze(src), self[idx+1:..])
        dim_i = sym_help._get_const(dim, "i", "dim")  # type: ignore[no-untyped-call]
        idx_i = sym_help._get_const(index, "i", "index")  # type: ignore[no-untyped-call]
        left = _slice_along(g, self, dim_i, 0, idx_i)
        right_start = IMAX if idx_i == -1 else idx_i + 1
        right = _slice_along(g, self, dim_i, right_start, IMAX)
        src_u = g.op("Unsqueeze", src, _const_i64(g, [dim_i]))
        return g.op("Concat", left, src_u, right, axis_i=dim_i)  # type: ignore[no-any-return]

    def _slice_scatter(
        g: GraphContext,
        self: torch.Value,
        src: torch.Value,
        dim: torch.Value,
        start: torch.Value,
        end: torch.Value,
        step: torch.Value,
    ) -> torch.Value:
        dim_i = sym_help._get_const(dim, "i", "dim")  # type: ignore[no-untyped-call]
        start_i = (
            0
            if sym_help._is_none(start)
            else sym_help._get_const(start, "i", "start")  # type: ignore[no-untyped-call]
        )
        end_i = (
            IMAX
            if sym_help._is_none(end)
            else sym_help._get_const(end, "i", "end")  # type: ignore[no-untyped-call]
        )
        step_i = (
            1
            if sym_help._is_none(step)
            else sym_help._get_const(step, "i", "step")  # type: ignore[no-untyped-call]
        )
        assert step_i == 1, "slice_scatter with step != 1 not handled"
        left = _slice_along(g, self, dim_i, 0, start_i)
        right = _slice_along(g, self, dim_i, end_i, IMAX)
        return g.op("Concat", left, src, right, axis_i=dim_i)  # type: ignore[no-any-return]

    def _copy(
        g: GraphContext,
        self: torch.Value,
        src: torch.Value,
        _non_blocking: torch.Value | None = None,
    ) -> torch.Value:
        # dtype unknown -> same-dtype copy, no cast needed
        with contextlib.suppress(Exception):
            src = g.op(
                "Cast",
                src,
                to_i=_type_utils.JitScalarType.from_value(self).onnx_type(),
            )
        return g.op("Expand", src, g.op("Shape", self))  # type: ignore[no-any-return]

    register_custom_op_symbolic("aten::select_scatter", _select_scatter, 17)
    register_custom_op_symbolic("aten::slice_scatter", _slice_scatter, 17)
    register_custom_op_symbolic("aten::copy", _copy, 17)