import sys
import time
import vtk

vtk_m = vtk.vtkMath()

def random_pts(n, r=1.0):
    global vtk_m
    ret = []
    for i in range(n):
        ret.append(r*vtk_m.Random())
    return ret

def associate_color(obj):
    f = vtk.vtkFloatArray()
    obj.Update()
    val = vtk_m.Random()
    for i in range(obj.GetOutput().GetNumberOfPoints()):
        f.InsertNextValue(val)
    obj.GetOutput().GetPointData().SetScalars(f)


def create_spheres(n, rad=0.1, region=1.0):
    sp = []
    for i in range(n):
        val = random_pts(3, region)
        s = vtk.vtkSphereSource()
        s.SetCenter(val)
        s.SetRadius(rad)
        associate_color(s)
        m = vtk.vtkPolyDataMapper()
        m.SetInput(s.GetOutput())
        a = vtk.vtkActor()
        a.SetMapper(m)
        sp.append(a)        
    return sp


def create_sphere_assy(n, rad=0.1, region=1.0):
    sp = create_spheres(n, rad, region)
    ass = vtk.vtkAssembly()
    for i in sp:
        ass.AddPart(i)
    return ass


def create_spheres_append(n, rad=0.1, region=1.0):
    app = vtk.vtkAppendPolyData()
    for i in range(n):
        val = random_pts(3, region)
        s = vtk.vtkSphereSource()
        s.SetCenter(val)
        s.SetRadius(rad)
        associate_color(s)
        app.AddInput(s.GetOutput())
        
    m = vtk.vtkPolyDataMapper()
    m.SetInput(app.GetOutput())
    a = vtk.vtkActor()
    a.SetMapper(m)
    return a

def viewer():
    ren = vtk.vtkRenderer()
    rw = vtk.vtkRenderWindow()
    rw.AddRenderer(ren)
    rwi = vtk.vtkRenderWindowInteractor()
    rwi.SetRenderWindow(rw)
    return ren, rwi

def show_fps(obj, evt):
    print "FPS:", 1.0/obj.GetLastRenderTimeInSeconds()


def main(option=1):
    ren, rwi = viewer()
    t1 = time.clock()
    if option == 1:
        a = create_spheres(2000, 0.1, 10)
        for x in a:
            ren.AddActor(x)
    elif option == 2:
        a = create_sphere_assy(2000, 0.1, 10)
        ren.AddActor(a)
    elif option == 3:
        a = create_spheres_append(2000, 0.1, 10)
        ren.AddActor(a)
    t2 = time.clock()
    print "Time taken to create and add actors:", t2 - t1, "secs"

    rwi.Render()
    ren.AddObserver('ModifiedEvent', show_fps)
    rwi.Start()


if __name__ == "__main__":
    if len(sys.argv) > 1:
        main(int(sys.argv[1]))
    else:
        main()