from pydrake.geometry import StartMeshcat

# Start the visualizer. The cell will output an HTTP link after the execution.
# Click the link and a MeshCat tab should appear in your browser.
meshcat = StartMeshcat()from pydrake.visualization import ModelVisualizer

visual_box_sdf = """<?xml version="1.0"?>
<sdf version="1.7">
  <model name="CompliantBox">
    <pose>0 0 0 0 0 0</pose>
    <link name="compliant_box">
      <visual name="visual">
        <geometry>
          <box>
            <size>0.10 0.20 0.40</size>
          </box>
        </geometry>
        <material>
          <diffuse>1.0 1.0 1.0 0.5</diffuse>
        </material>
      </visual>
    </link>
  </model>
</sdf>
"""

# Visualize the SDFormat string you just defined.
visualizer = ModelVisualizer(meshcat=meshcat, visualize_frames=True)
visualizer.parser().AddModelsFromString(visual_box_sdf, "sdf")
visualizer.Run(loop_once=True)from pydrake.visualization import ModelVisualizer

collision_box_sdf = """<?xml version="1.0"?>
<sdf version="1.7">
  <model name="CompliantBox">
    <pose>0 0 0 0 0 0</pose>
    <link name="compliant_box">
      <collision name="collision">
        <geometry>
          <box>
            <size>0.10 0.20 0.40</size>
          </box>
        </geometry>
        <drake:proximity_properties>
          <drake:compliant_hydroelastic/>
          <drake:hydroelastic_modulus>1e7</drake:hydroelastic_modulus>
          <drake:mu_dynamic>0.5</drake:mu_dynamic>
          <drake:hunt_crossley_dissipation>1.25</drake:hunt_crossley_dissipation>
        </drake:proximity_properties>
      </collision>
      <visual name="visual">
        <geometry>
          <box>
            <size>0.10 0.20 0.40</size>
          </box>
        </geometry>
        <material>
          <diffuse>1.0 1.0 1.0 0.5</diffuse>
        </material>
      </visual>
    </link>
  </model>
</sdf>
"""

# Visualize the SDFormat string you just defined.
visualizer = ModelVisualizer(meshcat=meshcat, visualize_frames=True)
visualizer.parser().AddModelsFromString(collision_box_sdf, "sdf")
visualizer.Run(loop_once=True)from pydrake.multibody.tree import SpatialInertia

spatial_inertia = SpatialInertia.SolidBoxWithMass(
          mass=1, lx=0.1, ly=0.2, lz=0.4)
mass = spatial_inertia.get_mass()
center_of_mass = spatial_inertia.get_com()
matrix6 = spatial_inertia.CopyToFullMatrix6()

print(f"mass = {mass}\n")
print(f"p_PScm = center of mass = {center_of_mass}\n")
print(f"I_SP = rotational inertia = \n{matrix6[:3, :3]}\n")from pydrake.visualization import ModelVisualizer

# Define a compliant-hydroelastic box from
# a given hydroelastic modulus.
def compliant_box(hydroelastic_modulus='1e7'):
    return f"""<?xml version="1.0"?>
    <sdf version="1.7">
      <model name="CompliantBox">
        <pose>0 0 0 0 0 0</pose>
        <link name="compliant_box">
          <inertial>
            <mass>1.0</mass>
            <inertia>
              <ixx>0.016666</ixx> <ixy>0.0</ixy> <ixz>0.0</ixz>
              <iyy>0.014166</iyy> <iyz>0.0</iyz>
              <izz>0.004166</izz>
            </inertia>
          </inertial>
          <collision name="collision">
            <geometry>
              <box>
                <size>0.10 0.20 0.40</size>
              </box>
            </geometry>
            <drake:proximity_properties>
              <drake:compliant_hydroelastic/>
              <drake:hydroelastic_modulus>{hydroelastic_modulus}</drake:hydroelastic_modulus>
              <drake:mu_dynamic>0.5</drake:mu_dynamic>
              <drake:hunt_crossley_dissipation>1.25</drake:hunt_crossley_dissipation>
            </drake:proximity_properties>
          </collision>
          <visual name="visual">
            <geometry>
              <box>
                <size>0.10 0.20 0.40</size>
              </box>
            </geometry>
            <material>
              <diffuse>1.0 1.0 1.0 0.5</diffuse>
            </material>
          </visual>
        </link>
      </model>
    </sdf>
    """

compliant_box_sdf = compliant_box()
print(compliant_box_sdf)

# Visualize the SDFormat string you just defined.
visualizer = ModelVisualizer(meshcat=meshcat, visualize_frames=True)
visualizer.parser().AddModelsFromString(compliant_box_sdf, "sdf")
visualizer.Run(loop_once=True)from pydrake.visualization import ModelVisualizer

# Create a rigid-hydroelastic table top
rigid_box_sdf = """<?xml version="1.0"?>
<sdf version="1.7">
  <model name="RigidBox">
    <link name="rigid_box_link">
      <visual name="visual">
        <pose>0 0 0 0 0 0</pose>
        <geometry>
          <box>
            <size>0.6 1.0 0.05</size>
          </box>
        </geometry>
        <material>
         <diffuse>0.9 0.8 0.7 0.5</diffuse>
        </material>
      </visual>
      <collision name="collision">
        <pose>0 0 0 0 0 0</pose>
        <geometry>
          <box>
            <size>0.6 1.0 0.05</size>
          </box>
        </geometry>
        <drake:proximity_properties>
          <drake:rigid_hydroelastic/>
          <drake:mu_dynamic>0.5</drake:mu_dynamic>
          <drake:hunt_crossley_dissipation>1.25</drake:hunt_crossley_dissipation>
        </drake:proximity_properties>
      </collision>
    </link>
    <frame name="top_surface">
      <pose relative_to="rigid_box_link">0 0 0.025 0 0 0</pose>
    </frame>
  </model>
</sdf>
"""

# Visualize the SDFormat string you just defined.
visualizer = ModelVisualizer(meshcat=meshcat, visualize_frames=True)
visualizer.parser().AddModelsFromString(rigid_box_sdf, "sdf")
visualizer.Run(loop_once=True)from pydrake.math import RigidTransform
from pydrake.multibody.parsing import Parser
from pydrake.multibody.plant import AddMultibodyPlant, MultibodyPlantConfig
from pydrake.systems.framework import DiagramBuilder

def clear_meshcat():
    # Clear MeshCat window from the previous blocks.
    meshcat.Delete()
    meshcat.DeleteAddedControls()

def add_scene(time_step=1e-3):
    builder = DiagramBuilder()
    plant, scene_graph = AddMultibodyPlant(
        MultibodyPlantConfig(
            time_step=time_step,
            discrete_contact_approximation="similar"),
        builder)
    parser = Parser(plant)

    # Load the table top and the box we created.
    parser.AddModelsFromString(compliant_box_sdf, "sdf")
    parser.AddModelsFromString(rigid_box_sdf, "sdf")

    # Weld the rigid box to the world so that it's fixed during simulation.
    # The top surface passes the world's origin.
    plant.WeldFrames(plant.world_frame(), 
                     plant.GetFrameByName("top_surface"))

    # Finalize the plant after loading the scene.
    plant.Finalize()

    # Set how high the center of the compliant box is from the world's origin. 
    # W = the world's frame
    # C = frame at the center of the compliant box
    X_WC = RigidTransform(p=[0, 0, 1])
    plant.SetDefaultFreeBodyPose(plant.GetBodyByName("compliant_box"), X_WC)

    return builder, plantfrom pydrake.visualization import ApplyVisualizationConfig, VisualizationConfig

def add_viz(builder, plant):
    ApplyVisualizationConfig(
        config=VisualizationConfig(
                   publish_period = 1 / 256.0,
                   publish_contacts = False),
        builder=builder, meshcat=meshcat)
    
    return builder, plant


from pydrake.systems.analysis import Simulator

# Test creation of the diagram by simulating for 0 second.
# For now, use only the DiagramBuilder from the first return value and
# ignore the other return value. We will use it later.
clear_meshcat()
builder, plant = add_scene()
add_viz(builder, plant)
simulator = Simulator(builder.Build())
simulator.AdvanceTo(0)from pydrake.systems.analysis import Simulator

def run_simulation(sim_time, time_step=1e-3):
    clear_meshcat()
    
    builder, plant = add_scene(time_step)
    add_viz(builder, plant)
    
    diagram = builder.Build()
    
    simulator = Simulator(diagram)
    simulator.set_target_realtime_rate(1.)
    
    meshcat.StartRecording(frames_per_second=256.0)
    simulator.AdvanceTo(sim_time)
    meshcat.StopRecording()


run_simulation(sim_time=6)meshcat.PublishRecording()from pydrake.common.value import Value
from pydrake.multibody.plant import ContactResults
from pydrake.systems.framework import LeafSystem

class ContactReporter(LeafSystem):
    def __init__(self):
        super().__init__()  # Don't forget to initialize the base class.
        self.DeclareAbstractInputPort(
            name="contact_results",
            model_value=Value(
                # Input port will take ContactResults from MultibodyPlant
                ContactResults()))
        # Calling `ForcedPublish()` will trigger the callback.
        self.DeclareForcedPublishEvent(self.Publish)
        
    def Publish(self, context):
        print()
        print(f"ContactReporter::Publish() called at time={context.get_time()}")
        contact_results = self.get_input_port().Eval(context)
        
        num_hydroelastic_contacts = contact_results.num_hydroelastic_contacts()
        print(f"num_hydroelastic_contacts() = {num_hydroelastic_contacts}")
        
        for c in range(num_hydroelastic_contacts):
            print()
            print(f"hydroelastic_contact_info({c}): {c}-th hydroelastic contact patch")
            hydroelastic_contact_info = contact_results.hydroelastic_contact_info(c)
            
            spatial_force = hydroelastic_contact_info.F_Ac_W()
            print("F_Ac_W(): spatial force (on body A, at centroid of contact surface, in World frame) = ")
            print(f"{spatial_force}")
                        
            print("contact_surface()")
            contact_surface = hydroelastic_contact_info.contact_surface()
            num_faces = contact_surface.num_faces()
            total_area = contact_surface.total_area()
            centroid = contact_surface.centroid()
            print(f"total_area(): area of contact surface in m^2 = {total_area}")
            print(f"num_faces(): number of polygons or triangles = {num_faces}")
            print(f"centroid(): centroid (in World frame) = {centroid}")        
        
        print()

def add_contact_report(builder, plant):   
    contact_reporter = builder.AddSystem(ContactReporter())    
    builder.Connect(plant.get_contact_results_output_port(),
                    contact_reporter.get_input_port(0))
        
    return builder, plantfrom pydrake.systems.analysis import Simulator

def run_simulation_with_contact_report(sim_time, time_step=1e-3):
    clear_meshcat()
    
    builder, plant = add_scene(time_step)
    add_viz(builder, plant)
    add_contact_report(builder, plant)
    
    diagram = builder.Build()
    
    simulator = Simulator(diagram)
    simulator.set_target_realtime_rate(1.)
    
    meshcat.StartRecording(frames_per_second=256.0)
    simulator.AdvanceTo(sim_time)
    meshcat.StopRecording()

    # Forced publish after the simulation has finished.
    diagram.ForcedPublish(simulator.get_context())
    

run_simulation_with_contact_report(sim_time=0)
meshcat.PublishRecording()run_simulation_with_contact_report(sim_time=1)
meshcat.PublishRecording()run_simulation_with_contact_report(sim_time=6)
meshcat.PublishRecording()from pydrake.multibody.meshcat import ContactVisualizer, ContactVisualizerParams


def add_contact_viz(builder, plant):
    contact_viz = ContactVisualizer.AddToBuilder(
        builder, plant, meshcat,
        ContactVisualizerParams(
            publish_period= 1.0 / 256.0,
            newtons_per_meter= 2e1,
            newton_meters_per_meter= 1e-1))

    return builder, plantfrom pydrake.systems.analysis import Simulator

def run_simulation_with_contact_report_and_viz(sim_time, time_step=1e-3):
    clear_meshcat()
    
    builder, plant = add_scene(time_step)
    add_viz(builder, plant)
    add_contact_report(builder, plant)
    add_contact_viz(builder, plant)
    
    diagram = builder.Build()
    
    simulator = Simulator(diagram)
    simulator.set_target_realtime_rate(1.0)
    
    meshcat.StartRecording(frames_per_second=256.0)
    simulator.AdvanceTo(sim_time)
    meshcat.StopRecording()

    # Numerically report contact results at the end of simulation.
    diagram.ForcedPublish(simulator.get_context())

    
run_simulation_with_contact_report_and_viz(sim_time=6)# In the current version, we can playback contact forces and torques;
# however, contact surfaces are not recorded properly.
# For now, we delete contact surfaces to prevent confusion.
# See issue https://github.com/RobotLocomotion/drake/issues/19142
meshcat.Delete("/drake/contact_forces/hydroelastic/compliant_box+rigid_box_link/contact_surface")

meshcat.PublishRecording()print(f"{meshcat.web_url()}/download")# run_simulation_with_contact_report_and_viz(sim_time=0.405)
# meshcat.Delete("/drake/contact_forces/hydroelastic/compliant_box+rigid_box_link/contact_surface")
# meshcat.PublishRecording()# new_hydroelastic_modulus = '1e1'
# compliant_box_sdf = compliant_box(hydroelastic_modulus = new_hydroelastic_modulus)
# 
# run_simulation_with_contact_report_and_viz(sim_time=2)
# meshcat.Delete("/drake/contact_forces/hydroelastic/compliant_box+rigid_box_link/contact_surface")
# meshcat.PublishRecording()# new_hydroelastic_modulus = '1e15'
# compliant_box_sdf = compliant_box(hydroelastic_modulus = new_hydroelastic_modulus)
# 
# run_simulation_with_contact_report_and_viz(sim_time=3)
# meshcat.Delete("/drake/contact_forces/hydroelastic/compliant_box+rigid_box_link/contact_surface")
# meshcat.PublishRecording()