import os
from pathlib import Path

from IPython.display import Code

from pydrake.geometry import StartMeshcat
from pydrake.math import RigidTransform
from pydrake.multibody.meshcat import ContactVisualizer, ContactVisualizerParams
from pydrake.multibody.parsing import PackageMap, Parser
from pydrake.multibody.plant import AddMultibodyPlant, MultibodyPlantConfig
from pydrake.systems.analysis import Simulator
from pydrake.systems.framework import DiagramBuilder
from pydrake.visualization import ApplyVisualizationConfig, ModelVisualizer, VisualizationConfig# 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()# Show the contents of the SDFormat file.
bell_pepper_url = "package://drake_models/veggies/yellow_bell_pepper_no_stem_low.sdf"
bell_pepper_str = Path(PackageMap().ResolveUrl(bell_pepper_url)).read_text(encoding="utf-8")
Code(bell_pepper_str, language="xml")# Visualize the SDFormat file you just defined.
visualizer = ModelVisualizer(meshcat=meshcat)
visualizer.AddModels(url=bell_pepper_url)
visualizer.Run(loop_once=True)# Show the contents of the URDF file.
bowl_url = "package://drake_models/dishes/evo_bowl_compliant.urdf"
bowl_str = Path(PackageMap().ResolveUrl(bowl_url)).read_text(encoding="utf-8")
Code(bowl_str, language="xml")# (Drake users can ignore this cell. When Drake regression testing runs in CI,
# we need to use a coarser tetrahedral mesh to improve debug performance.)
test_mode = "TEST_SRCDIR" in os.environ
if test_mode:
    bowl_str = bowl_str.replace("evo_bowl_fine44k.vtk", 
                                "evo_bowl_coarse3k.vtk")# Visualize the URDF file.
visualizer = ModelVisualizer(meshcat=meshcat)
visualizer.parser().AddModels(file_contents=bowl_str, file_type="urdf")
visualizer.Run(loop_once=True)# Show the contents of the URDF file.
table_top_url = "package://drake_models/dishes/table_top.urdf"
table_top_str = Path(PackageMap().ResolveUrl(table_top_url)).read_text(encoding="utf-8")
Code(table_top_str, language="xml")# Visualize the URDF file.
visualizer = ModelVisualizer(meshcat=meshcat)
visualizer.AddModels(url=table_top_url)
visualizer.Run(loop_once=True)def add_scene(time_step):
    builder = DiagramBuilder()
    plant, scene_graph = AddMultibodyPlant(
        MultibodyPlantConfig(
            time_step=time_step,
            discrete_contact_approximation="lagged"),
        builder)
    parser = Parser(plant)

    # Load the assets that we created.
    parser.AddModels(url=bell_pepper_url)
    parser.AddModels(file_contents=bowl_str, file_type="urdf")
    parser.AddModels(url=table_top_url)

    # Weld the table top 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()

    # Place the bowl on top of the table.
    X_WB = RigidTransform(p=[0, 0, 0.03])
    plant.SetDefaultFreeBodyPose(plant.GetBodyByName("bowl"), X_WB)
    
    # Drop the bell pepper from above the rim of the bowl. 
    X_WC = RigidTransform(p=[-0.06, 0, 0.30])
    plant.SetDefaultFreeBodyPose(plant.GetBodyByName("yellow_bell_pepper_no_stem"), X_WC)

    return builder, plantdef add_viz(builder, plant):
    ApplyVisualizationConfig(
        builder=builder, meshcat=meshcat,
        config=VisualizationConfig(
                 publish_contacts=False))    
    ContactVisualizer.AddToBuilder(
        builder=builder, plant=plant, meshcat=meshcat,
        params=ContactVisualizerParams(
                 newtons_per_meter=20,
                 newton_meters_per_meter= 0.1))# Clear MeshCat window from the previous blocks.
meshcat.Delete()
meshcat.DeleteAddedControls()

time_step = 1e-2
builder, plant = add_scene(time_step)
add_viz(builder, plant)

diagram = builder.Build()

simulator = Simulator(diagram)

# In interactive mode, simulate for longer time.
# In test mode, simulate for shorter time.
sim_time = 2 if not test_mode else 0.01

meshcat.StartRecording()
simulator.set_target_realtime_rate(1)
simulator.AdvanceTo(sim_time)
meshcat.StopRecording()
meshcat.PublishRecording()print(f"{meshcat.web_url()}/download")