running_as_notebook = False  # Manually set this to True if you are a human (see #13862).

# Imports
import numpy as np

from IPython.display import display
from ipywidgets import Layout, Text, ToggleButton

from pydrake.geometry import DrakeVisualizer, SceneGraph
from pydrake.lcm import DrakeLcm
from pydrake.multibody.plant import AddMultibodyPlantSceneGraph, MultibodyPlant
from pydrake.multibody.parsing import Parser
from pydrake.systems.analysis import Simulator                                      
from pydrake.systems.framework import DiagramBuilder
from pydrake.systems.rendering import MultibodyPositionToGeometryPose

from pydrake.multibody.jupyter_widgets import JointSliders, MakeJointSlidersThatPublishOnCallbackbuilder = DiagramBuilder()
scene_graph = builder.AddSystem(SceneGraph())
plant = MultibodyPlant(time_step=0.0)
plant.RegisterAsSourceForSceneGraph(scene_graph)
Parser(plant, scene_graph).AddModelsFromUrl(
    url="package://drake_models/iiwa_description/sdf/iiwa7_no_collision.sdf")
plant.Finalize()

# Add sliders to set positions of the joints.
sliders = builder.AddSystem(JointSliders(robot=plant))
to_pose = builder.AddSystem(MultibodyPositionToGeometryPose(plant))
builder.Connect(sliders.get_output_port(0), to_pose.get_input_port())
builder.Connect(
    to_pose.get_output_port(),
    scene_graph.get_source_pose_port(plant.get_source_id()))

# Note: We can't use AddDefaultVisualization because there is no plant
# in the Diagram.
DrakeVisualizer.AddToBuilder(builder, scene_graph)

# Make the diagram and run it.
diagram = builder.Build()
simulator = Simulator(diagram)

if running_as_notebook:  
    simulator.set_target_realtime_rate(1.0)
    stop_button = ToggleButton(value=False, description='Stop Simulation')
    display(stop_button)
    while not stop_button.value:
        simulator.AdvanceTo(simulator.get_context().get_time() + 5.0)
    stop_button.value = False
else:  # running as a test.
    simulator.AdvanceTo(0.1)
builder = DiagramBuilder()
plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)
Parser(plant).AddModelsFromUrl(
    url="package://drake_models/iiwa_description/sdf/iiwa7_no_collision.sdf")
plant.Finalize()

viz = DrakeVisualizer.AddToBuilder(builder, scene_graph)

# Make the diagram and run it.
diagram = builder.Build()
context = diagram.CreateDefaultContext()

callback_output = Text(value="Not called yet", description="Callback output: ", 
                       layout=Layout(width='500px'), 
                       style={'description_width':'initial'})
num_calls = 0
def callback_test(context):
    global num_calls
    num_calls += 1
    callback_output.value = f"Called {num_calls} times"

sliders = MakeJointSlidersThatPublishOnCallback(plant, viz, context, my_callback=callback_test)
display(callback_output)
