import numpy as np
from mayavi.scripts import mayavi2
from tvtk.api import tvtk
 
def pyramid(length, height, points, polygons):
    '''
    Fuegt der Liste der Punkte und der Liste der Dreiecke eine
    Pyramide um den Ursprung mit Grundlaenge length und hoehe height hinzu.
    '''
    a = len(points)
    points.append([length, length, 0])
    points.append([-length, length, 0])
    points.append([length, -length, 0])
    points.append([-length, -length, 0])
    points.append([0, 0, height])
    polygons.append([a, a+1, a+2])
    polygons.append([a+1, a+2, a+3])
    polygons.append([a, a+1, a+4])
    polygons.append([a, a+2, a+4])
    polygons.append([a+3, a+1, a+4])
    polygons.append([a+3, a+2, a+4])
    return (points, polygons)
 
 
def circle(center, radius, step, points, polygons):
    '''
    Approximiert einen Kreis Entlang der x-y Achse mithilfe
    eines gleichseitigen Vielecks.
    '''
    a = len(points)
    DEGREE = 2./360.
    tridegree = 360./step
    points.append(list(center))
    for i in range(0, step):
        x = center[0] + np.sin(DEGREE*tridegree*(i+1)*np.pi)*radius
        y = center[1] + np.cos(DEGREE*tridegree*(i+1)*np.pi)*radius
        triangle = [x, y, center[2]]
        points.append(triangle)
        if i != step-1:
            polygons.append([a, a+i+1, a+i+2])
    polygons.append([a, a+i+1, a+1])
 
@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 = [], []
    pyramid(10, 10, points, polygons)
    for i in range(1, 6):
        circle((0, 0, i*20), 10, 4**i, 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__':
    main()