# Customize Trame toolbar
# =======================

# Bring more of the power of trame to the jupyter view.
# This example shows how to add custom tools using the
# jupyter_kwargs option with `show()`.
import asyncio

import pyvista as pv
from trame_pyvista.ui.vuetify3 import button
from trame_pyvista.ui.vuetify3 import divider
from trame_pyvista.ui.vuetify3 import select
from trame_pyvista.ui.vuetify3 import slider
from trame_pyvista.ui.vuetify3 import text_field

# Let's first create the menu items we want to add to the trame's toolbar.
# Here we want a "play" button that will be later connected to a slider
# through the `button_play` function. The slider itself will represent the
# "resolution" of the model we will render, a text field where the value of
# the "resolution" will be displayed.
# We will also add a dropdown menu to toggle the visibility of the model.
# The dividers are the same as already used to divide and organize the toolbar.
def custom_tools():
    divider(vertical=True, classes='mx-1')
    button(
        click=button_play,
        icon='mdi-play',
        tooltip='Play',
    )

    slider(
        model=('resolution', 10),
        tooltip='Resolution slider',
        min=3,
        max=20,
        step=1,
        dense=True,
        hide_details=True,
        style='width: 300px',
        classes='my-0 py-0 ml-1 mr-1',
    )
    text_field(
        model=('resolution', 10),
        tooltip='Resolution value',
        readonly=True,
        type='number',
        dense=True,
        hide_details=True,
        style='min-width: 40px; width: 60px',
        classes='my-0 py-0 ml-1 mr-1',
    )

    divider(vertical=True, classes='mx-1')
    select(
        model=('visibility', 'Show'),
        tooltip='Toggle visibility',
        items=['Visibility', ['Hide', 'Show']],
        hide_details=True,
        dense=True,
    )

# The button callback function `button_play` needs to be created before starting
# the server. This function will toggle the boolean state variable `play`
# and flush the server, i.e. "force" the server to see the change.
# We will see more on the state variables in a bit, but we need to create the
# function here otherwise the server will complain `button_play` does not exist.
def button_play():
    state.play = not state.play
    state.flush()

# We will do a simple rendering of a Cone using ConeSouce.
# When using the `pl.show` method. The function we created `custom_tools`
# should be passed as a `jupyter_kwargs` argument under the key
# `add_menu_items`.
pl = pv.Plotter(notebook=True)
algo = pv.ConeSource()
mesh_actor = pl.add_mesh(algo)

widget = pl.show(jupyter_kwargs=dict(add_menu_items=custom_tools), return_viewer=True)

# To interact with `trame`'s server we need to get the server's state.
# We initialize the `play` variable in the shared state and this will be
# controlled by the play button we created. Note that when creating the
# `slider`, the `text_field` and the `select` tools, we passed something
# like `model=("variable", value). This will automatically create the variable
# "variable" with value ``value` in the server's shared state, so we do not need
# to create `state.resolution` or `state.visibility`.
state, ctrl = widget.viewer.server.state, widget.viewer.server.controller
state.play = False
ctrl.view_update = widget.viewer.update

# Now we can create the callback functions for our menu items.
# The functions are decorated with a `state.change("variable")`. This means
# they will be called when this specific variable has its value changed in the
# server's shared state. When `resolution` changes, we want to update the
# resolution of our cone algorithm. When `visibility` changes, we want to toggle the
# visibility of our cone.
# The `play` variable is a little bit trickier. We want to start something like
# a timer so that an animation can be set to play. To do that with `trame` we need
# to have an asynchronous function so we can continue to do stuff while the
# "timer" function is running. The `_play` function will be called when the `play`
# variable is changed (when we click the play button, through the `button_play`
# callback). While `state.play` is `True` we want to play the animation. We
# change the `state.resolution` value, but to really call the `update_resolution`
# function we need to `flush` the server and force it to see the change in
# the shared variables. When `state.play` changes to `False`, the animation stops.
# Note that using `while play: ...` would not work here because it is not the
# actual state variable, but only an argument value passed to the callback function.
# trame callbacks
@state.change('play')
async def _play(play, **kwargs):
    while state.play:
        state.resolution += 1
        state.flush()
        if state.resolution >= 20:
            state.play = False
        await asyncio.sleep(0.3)


@state.change('resolution')
def update_resolution(resolution, **kwargs):
    algo.resolution = resolution
    ctrl.view_update()


@state.change('visibility')
def set_visibility(visibility, **kwargs):
    toggle = {'Hide': 0, 'Show': 1}
    mesh_actor.visibility = toggle[visibility]
    ctrl.view_update()


widget

# ----------------------------------------------------------------------
# Generated by sphinx-examples-as-code https://github.com/pyvista/sphinx-examples-as-code

