import numpy as np
import matplotlib.pyplot as plt
from pydrake.multibody.plant import MultibodyPlant
from pydrake.multibody.tree import (
    DoorHingeConfig, DoorHinge, RevoluteJoint
)

# Build a simple MultibodyPlant that contains two bodies and a revolute joint.
# Then add a DoorHinge ForceElement to the joint.
plant = MultibodyPlant(0.0)

body_a = plant.AddRigidBody(name="body_a")
body_b = plant.AddRigidBody(name="body_b")

revolute_joint = plant.AddJoint(RevoluteJoint(
        name="revolve_joint", frame_on_parent=body_a.body_frame(),
        frame_on_child=body_b.body_frame(), axis=[0, 0, 1],
        damping=0.0))

# Create a default DoorHingeConfig. The user should change the parameters before
# calling AddForceElement.
door_hinge_config = DoorHingeConfig()
door_hinge_config.spring_constant = 3.0

# Add the DoorHinge and finalize the plant.
door_hinge = plant.AddForceElement(DoorHinge(
            joint=revolute_joint, config=door_hinge_config))
plant.Finalize()
# Prepare the data for plotting.
min_q = 0.0
max_q = np.deg2rad(80)
min_v = -0.1
max_v = 0.1
step_size = 0.0005
q_array = np.arange(min_q, max_q, step_size)
v_array = np.arange(min_v, max_v, step_size)

# Plot the torque over angle at zero velocity
torque_zero_velocity = [
    door_hinge.CalcHingeTorque(q, 0.0)
    for q in q_array
]

# Plot the torque over the angular velocity at 30 degree angle
torque_30_angle = [
    door_hinge.CalcHingeTorque(np.deg2rad(30), v)
    for v in v_array
]
plt.plot(v_array, torque_30_angle)
plt.show()