import numpy as np
from mayavi import mlab
from traits.api import HasTraits, Range, Instance, on_trait_change
from traitsui.api import View, Item, Group
from mayavi.core.api import PipelineBase
from mayavi.core.ui.api import MayaviScene, SceneEditor, MlabSceneModel

def radialsine(x,y,delta=0):
    return np.sin(np.sqrt(x**2+y**2)+delta)

class bartest(HasTraits):
    delta = Range(-np.pi,np.pi,0.)
    scene = Instance(MlabSceneModel, ())
    plot = Instance(PipelineBase)
    xs,ys=np.mgrid[-2:2:10j,-2:2:10j]
    zs=np.zeros_like(xs)

    def __init__(self):
        HasTraits.__init__(self)
        self.scene.background, self.scene.foreground = (1,1,1), (0,0,0)
        #axes
        self.scene.mlab.plot3d([-1.2, 1.2], [0, 0], [0, 0], color=(0,0,0), tube_radius=.005)
        self.scene.mlab.plot3d([0, 0], [-1.2, 1.2], [0, 0], color=(0,0,0), tube_radius=.005)
        self.scene.mlab.plot3d([0, 0], [0, 0], [-1.2, 1.2], color=(0,0,0), tube_radius=.005)   
    
    @on_trait_change('scene.activated')
    def anntonate(self):
        #labels
        coords, labels = [(2,0,0),(0,2,0),(0,0,2)], ['x','y','z']
        for coord, label in zip(coords, labels):
            self.scene.mlab.text3d(*coord, label, color=(0,0,0), scale=0.5)

    @on_trait_change('delta, scene.activated')
    def update(self):
        if self.plot is None:     
            self.plot = self.scene.mlab.barchart(self.xs, self.ys, self.zs, radialsine(x=self.xs,y=self.ys,delta=self.delta), vmax=1,vmin=-1,lateral_scale=0.3,colormap='viridis',figure=self.scene.mayavi_scene)
        else:
            self.plot.mlab_source.trait_set(x=self.xs, y=self.ys, z=radialsine(self.xs,self.ys,self.delta), s=self.xs)

    view = View(Item('scene', editor=SceneEditor(scene_class=MayaviScene), height=250, width=300, show_label=False),
                Item('delta'))#, format_str='%.2f', editor_args={'low_label':'-π', 'high_label':' π'}))

my_model = bartest()
my_model.configure_traits()