import sys

from mayavi.api import OffScreenEngine
from mayavi.tools.tools import _typical_distance
from mayavi.modules.api import Glyph, Surface
from mayavi.core.api import Engine

from simphony.cuds.abstractmesh import ABCMesh
from simphony.cuds.abstractparticles import ABCParticles
from simphony.cuds.abstractlattice import ABCLattice

from simphony_mayavi.sources.api import (
    ParticlesSource, MeshSource, LatticeSource)


def snapshot(cuds, filename):
    """ Shave a snapshot of the cuds object using the default visualization.

     Parameters
     ----------
     cuds :
         A top level cuds object (e.g. a mesh). The method will detect
         the type of object and create the appropriate visualisation.

     filename : string
         The filename to use for the output file.

    """
    size = 800, 600

    if sys.platform == 'win32':
        engine = OffScreenEngine()
    else:
        # The OffScreenEngine does not reliably work for linux/mac-os
        engine = Engine()

    engine.start()
    try:
        if isinstance(cuds, ABCMesh):
            source = MeshSource.from_mesh(cuds)
            window = engine.new_scene()
            engine.add_source(source)
            engine.add_module(Surface())
        elif isinstance(cuds, ABCParticles):
            source = ParticlesSource.from_particles(cuds)
            engine = OffScreenEngine()
            engine.start()
            window = engine.new_scene()
            engine.add_source(source)
            module = Glyph()
            glyph_source = module.glyph.glyph_source
            glyph_source.glyph_source = \
                glyph_source.glyph_dict['sphere_source']
            module.glyph.glyph.scale_factor = \
                _typical_distance(source.data) * 0.5
            module.glyph.scale_mode = 'data_scaling_off'
            engine.add_module(module)
            module = Surface()
            module.actor.mapper.scalar_visibility = False
            engine.add_module(module)
        elif isinstance(cuds, ABCLattice):
            source = LatticeSource.from_lattice(cuds)
            engine = OffScreenEngine()
            engine.start()
            window = engine.new_scene()
            engine.add_source(source)
            module = Glyph()
            glyph_source = module.glyph.glyph_source
            glyph_source.glyph_source = \
                glyph_source.glyph_dict['sphere_source']
            module.glyph.glyph.scale_factor = \
                _typical_distance(source.data) * 0.5
            module.glyph.scale_mode = 'data_scaling_off'
            engine.add_module(module)
        else:
            msg = 'Provided object {} is not of any known cuds type'
            raise TypeError(msg.format(type(cuds)))

        # Use some default values for the camera
        camera = window.scene.camera
        camera.azimuth(45)
        camera.elevation(30)
        window.scene.reset_zoom()
        window.scene.save(filename, size)
    finally:
            engine.stop()

@@ -524,17 +524,17 @@
  visuali
-z
+s
 ation.%0A%0A
@@ -842,17 +842,17 @@
 latform 
-=
+!
 = 'win32
@@ -866,33 +866,46 @@
     
-engine = OffScreenEngine(
+_snapshot_windows(cuds, filename, size
 )%0A  
@@ -990,25 +990,123 @@
         
-engine = 
+_snapshot_linux(cuds, filename, size)%0A%0A%0Adef _snapshot_windows(cuds, filename, size):%0A    engine = OffScreen
 Engine()
@@ -1106,17 +1106,16 @@
 ngine()%0A
-%0A
     engi
@@ -3075,34 +3075,30 @@
 inally:%0A        
-    
 engine.stop()%0A
@@ -3094,11 +3094,1226 @@
 ine.stop
+()%0A%0A%0Adef _snapshot_linux(cuds, filename, size):%0A    width, height = size%0A    from mayavi import mlab%0A    if isinstance(cuds, ABCMesh):%0A        source = MeshSource.from_mesh(cuds)%0A        mlab.options.offscreen = True%0A        mlab.pipeline.surface(source, name=cuds.name)%0A    elif isinstance(cuds, ABCParticles):%0A        source = ParticlesSource.from_particles(cuds)%0A        scale_factor = _typical_distance(source.data) * 0.5%0A        mlab.options.offscreen = True%0A        mlab.pipeline.glyph(%0A            source, name=cuds.name,%0A            scale_factor=scale_factor, scale_mode='none')%0A        surface = mlab.pipeline.surface(source)%0A        surface.actor.mapper.scalar_visibility = False%0A    elif isinstance(cuds, ABCLattice):%0A        source = LatticeSource.from_lattice(cuds)%0A        scale_factor = _typical_distance(source.data) * 0.5%0A        mlab.options.offscreen = True%0A        mlab.pipeline.glyph(%0A            source, name=cuds.name,%0A            scale_factor=scale_factor, scale_mode='none')%0A    else:%0A        msg = 'Provided object %7B%7D is not of any known cuds type'%0A        raise TypeError(msg.format(type(cuds)))%0A    mlab.savefig(filename, (width, height), magnification=1.0)%0A    mlab.clf()%0A    mlab.close
 ()%0A

make sure that we discard the figure when we finish on linux