import numpy as np
from mayavi.scripts import mayavi2
from tvtk.api import tvtk
 
class voxels(object):
 
        def __init__(self, positions):
            self.points = []
            self.polygons = []
            self.positions = positions
            self.mesh = None
 
        def createmesh(self, length):
            for i in self.positions:
                self.cube(i, length)
            points = np.asarray(self.points)
            polygons = np.asarray(self.polygons)
            mesh = tvtk.PolyData(points=points, polys=polygons)
            return(mesh)
 
 
        def cube(self, start, length):
            '''
            Fuegt der Liste der Punkte und der Liste der Dreiecke eine
            Pyramide um den Ursprung mit Grundlaenge length und hoehe height hinzu.
            '''
            a = len(self.points)
            self.points.append([start[0], start[1], start[2]])
            self.points.append([start[0]+length, start[1], start[2]])
            self.points.append([start[0], start[1]+length, start[2]])
            self.points.append([start[0], start[1], start[2]+length])
            self.points.append([start[0]+length, start[1]+length, start[2]])
            self.points.append([start[0], start[1]+length, start[2]+length])
            self.points.append([start[0]+length, start[1], start[2]+length])
            self.points.append([start[0]+length, start[1]+length, start[2]+length])
            self.polygons.append([a, a+1, a+2])
            self.polygons.append([a, a+2, a+3])
            self.polygons.append([a, a+1, a+3])
            self.polygons.append([a+5, a+2, a+3])
            self.polygons.append([a+2, a+4, a+5])
            self.polygons.append([a+3, a+5, a+6])
            self.polygons.append([a+1, a+2, a+4])
            self.polygons.append([a+4, a+5, a+7])
            self.polygons.append([a+6, a+4, a+7])
            self.polygons.append([a+5, a+6, a+7])
            self.polygons.append([a+1, a+4, a+6])
            self.polygons.append([a+1, a+3, a+6])
 
 
@mayavi2.standalone
def view(mesh):
    '''Stellt ein tvtk.Polydata mesh im mayavi2 Standalone dar.'''
    from mayavi.sources.vtk_data_source import VTKDataSource
    from mayavi.modules.surface import Surface
 
    mayavi.new_scene()
    src = VTKDataSource(data=mesh)
    mayavi.add_source(src)
    s = Surface()
    mayavi.add_module(s)
 
def main():
    '''Zeichnet eine Pyramide und 4 Kreise'''
    points, polygons = [], []
    cube((0, 0, 0), 10, points, polygons)
    points = np.asarray(points)
    polygons = np.asarray(polygons)
    mesh = tvtk.PolyData(points=points, polys=polygons)
    temperature = np.array([10., 20., 30., 40.])
    mesh.point_data.scalars = temperature
    view(mesh)
 
 
if __name__ == '__main__':
    positions = [(x, x, x) for x in range(10)]
    cube = voxels(positions)
    view(cube.createmesh(0.5))