import numpy as np

from pydrake.multibody.tree import SpatialInertia, UnitInertia
from pydrake.multibody.plant import MultibodyPlant_, MultibodyPlant
from pydrake.autodiffutils import AutoDiffXdplant = MultibodyPlant(time_step=0.0)
body1 = plant.AddRigidBody(
    "body1",
    M_BBo_B=SpatialInertia(
        mass=2.0,
        p_PScm_E=[0, 0, 0],
        # N.B. Rotational inertia is unimportant for calculations
        # in this notebook, and thus is arbitrarily chosen.
        G_SP_E=UnitInertia(0.1, 0.1, 0.1),
    ),
)
body2 = plant.AddRigidBody(
    "body2",
    M_BBo_B=SpatialInertia(
        mass=0.5,
        p_PScm_E=[0, 0, 0],
        # N.B. Rotational inertia is unimportant for calculations
        # in this notebook, and thus is arbitrarily chosen.
        G_SP_E=UnitInertia(0.1, 0.1, 0.1),
    ),
)
plant.Finalize()

plant_ad = plant.ToScalarType[AutoDiffXd]()
body1_ad = plant_ad.get_body(body1.index())
body2_ad = plant_ad.get_body(body2.index())
context_ad = plant_ad.CreateDefaultContext()m1 = AutoDiffXd(2.0, [1.0, 0.0])
body1_ad.SetMass(context_ad, m1)

m2 = AutoDiffXd(0.5, [0.0, 1.0])
body2_ad.SetMass(context_ad, m2)def get_z_component(plant, body, v):
    assert body.is_floating()
    # Floating-base velocity dofs are organized as [angular velocity; translation velocity].
    v_start = body.floating_velocities_start_in_v()
    nv_pose = 6
    rxyz_txyz = v[v_start:v_start + nv_pose]
    assert len(rxyz_txyz) == nv_pose
    txyz = rxyz_txyz[-3:]
    z = txyz[2]
    return z@np.vectorize
def ad_to_string(x):
    # Formats an array of AutoDiffXd elements to a string.
    # Note that this implementation is for a scalar, but we use `np.vectorize` to
    # effectively convert our array to `ndarray` of strings.
    return f"AutoDiffXd({x.value()}, derivatives={x.derivatives()})"tau_g = plant_ad.CalcGravityGeneralizedForces(context_ad)
tau_g_z1 = get_z_component(plant_ad, body1_ad, tau_g)
tau_g_z2 = get_z_component(plant_ad, body2_ad, tau_g)
print(ad_to_string(tau_g_z1))
print(ad_to_string(tau_g_z2))