#!/usr/bin/env python
# coding: utf-8

# # Taking Derivatives of `MultibodyPlant` Computations w.r.t. Mass
# For instructions on how to run these tutorial notebooks, please see the [README](https://github.com/RobotLocomotion/drake/blob/master/tutorials/README.md).
# 
# We create a simple `MultibodyPlant` (`MultibodyPlant_[float]`). This means you *could*
# parse it from a URDF / SDFormat file.
# We then convert the system to `AutoDiffXd`, and set the parameters and partial derivatives
# we desire. In this case, we want $\frac{\partial{\boldsymbol{f}}}{\partial{\boldsymbol{m}}}$, where $\boldsymbol{f}$ is just some arbitrary expression.
# 
# In our case, we choose $\boldsymbol{f}$ to be generalized forces at the default / home configuration.
# Also in this case, we choose $\boldsymbol{m} = \left[ m_1, m_2 \right] \in \mathbb{R}_+^2$ just to show how to choose gradients for
# independent values.
# 
# For related reading, see also:
# - [Underactuated: System Identification](http://underactuated.csail.mit.edu/sysid.html) - at present, this only presents the symbolic approach for `MultibodyPlant`.
# - [`Modeling Dynamical Systems`](./dynamical_systems.ipynb)
# - [`Mathematical Program MultibodyPlant Tutorial`](./mathematical_program_multibody_plant.ipynb)

# In[ ]:


import numpy as np

from pydrake.multibody.tree import SpatialInertia, UnitInertia
from pydrake.multibody.plant import MultibodyPlant_, MultibodyPlant
from pydrake.autodiffutils import AutoDiffXd


# In[ ]:


plant = 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()


# We need to take gradients with respect to particular parameters, so we populate those parameters now.
# 
# For forward-mode automatic differentiation, we must ensure that specify our gradients
# according to our desired independent variables. In this case, we want $m_1$ and $m_2$
# to be independent, so we ensure their gradients are distinct unit vectors. (You could
# use [`InitializeAutoDiff`](https://drake.mit.edu/pydrake/pydrake.autodiffutils.html#pydrake.autodiffutils.InitializeAutoDiff) to do this, but
# we do it "by hand" here for illustration.)
# 
# If we wanted, we could set more parameters (e.g. center-of-mass), but we'll stick to just mass for simplicity.
# 
# It is important to note that every other "autodiff-able" quantity in this case (state,
# other parameters) are considered constant w.r.t. our independent values (mass), thus
# they will have 0-valued gradients.

# In[ ]:


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)


# The generalized force, in the z-translation direction for each body $i$, should just be $(-m_i \cdot g)$ with derivative $(-g)$.

# In[ ]:


def get_z_component(plant, body, v):
    assert body.is_floating()
    # N.B. This method returns position w.r.t. *state* [q; v]. We only have v (or vdot).
    x_start = body.floating_velocities_start()
    # Floating-base velocity dofs are organized as [angular velocity; translation velocity].
    v_start = x_start - plant.num_positions()
    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


# In[ ]:


@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()})"


# In[ ]:


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))


# In[ ]: