import panel as pn
import numpy as np
import pyvista as pv
pn.extension('vtk', sizing_mode="stretch_width")

alpha = 2
xvals  = np.linspace(-4, 4,101)
yvals  = np.linspace(-4, 4,101)
xs, ys = np.meshgrid(xvals, yvals)

#temporal function to create data on a plane
def time_function(time):
    return np.sin(((ys/alpha)**alpha+time)*xs)

# 3d plane to support the data
mesh_ref = pv.UniformGrid(
    (xvals.size, yvals.size, 1), #dims
    (xvals[1]-xvals[0],yvals[1]-yvals[0],1), #spacing
    (xvals.min(),yvals.min(),0) #origin
)
mesh_ref.point_arrays.append(time_function(0).flatten(order='F'), 'scalars') #add data for time=0
pl_ref = pv.Plotter()
pl_ref.add_mesh(mesh_ref, cmap='rainbow')
pn.panel(pl_ref.ren_win)

mesh_warped = mesh_ref.warp_by_scalar() # warp the mesh using data at time=0
#create the pyvista plotter
pl = pv.Plotter()
pl.add_mesh(mesh_warped, cmap='rainbow')

#initialize panel and widgets
camera = {
    'position': [13.443258285522461, 12.239550590515137, 12.731934547424316],
    'focalPoint': [0, 0, 0],
     'viewUp': [-0.41067028045654297, -0.40083757042884827, 0.8189500570297241]
}
vtkpan = pn.panel(pl.ren_win, orientation_widget=True, sizing_mode='stretch_both', camera=camera)
frame = pn.widgets.Player(value=0, start=0, end=50, interval=100, loop_policy="reflect", height=100)

@pn.depends(frame=frame.param.value)
def update_3d_warp(frame):
    #the player value range in between 0 and 50, howver we want time between 0 and 10
    time = frame/5
    data = time_function(time).flatten(order='F')
    mesh_ref.point_arrays.append(data, 'scalars')
    mesh_warped.point_arrays.append(data, 'scalars')
    mesh_warped.points = mesh_ref.warp_by_scalar(factor=0.5).points
    vtkpan.synchronize()

component = pn.Column(frame, vtkpan, update_3d_warp, height=600)
component

pn.template.FastListTemplate(site="Panel", title="VTK Warp", main=["This app demonstrates the use of `VTK`", component]).servable();