#!/usr/bin/env python

import numpy as np

from ase import molecule

atoms = molecule('H2O')
atoms.center(vacuum=4.0)

import vtk
from ase.visualize.vtk.atoms import vtkAtoms

usewx = False
try:
    import wx
    usewx = True
except ImportError:
    pass

if usewx:
    from vtk.wx.wxVTKRenderWindow import wxVTKRenderWindow

    app = wx.PySimpleApp()
    frame = wx.Frame(None, -1, 'wxVTKRenderWindow', size=(800,600))
    widget = wxVTKRenderWindow(frame, -1)
    win = widget.GetRenderWindow()
    ren = vtk.vtkRenderer()
    win.AddRenderer(ren)
else:
    ren = vtk.vtkRenderer()
    win = vtk.vtkRenderWindow()
    win.AddRenderer(ren)
    win.SetSize(800,600)

    iren = vtk.vtkRenderWindowInteractor()
    iren.SetRenderWindow(win)

    style = vtk.vtkInteractorStyleTrackballCamera()
    iren.SetInteractorStyle(style)

va = vtkAtoms(atoms)
va.add_cell()
va.add_axes()
#va.add_forces()
va.add_actors_to_renderer(ren)

if usewx:
    frame.Show()
    app.MainLoop()
else:
    iren.Initialize()
    win.OffScreenRenderingOff()
    win.Render()
    iren.Start()

w2i = vtk.vtkWindowToImageFilter()
w2i.SetInput(win)
pw = vtk.vtkPNGWriter()
pw.SetFileName("dump.png")
pw.SetInputConnection(w2i.GetOutputPort())
pw.Write()