#!/usr/bin/env python

# This is an example script to generate a stereoscopic visualization
# We'll import some libraries
import random
import os

from enthought import mayavi # the application
from enthought.mayavi import mlab # a procedural facade to the library which uses the same names as matlab
from enthought.mayavi.modules.api import Outline, GridPlane 
from enthought.mayavi.sources.vtk_data_source import VTKDataSource
from enthought.tvtk.api import tvtk


import numpy
from numpy import mgrid, empty, sin, pi, array, meshgrid, arange, prod


# some path to your local netcdf file
path = 'rhinerofi.nc'
path = '/Users/fedorbaart/Documents/checkouts/baart_f/programs/matlab/gerben/gerbenvankees.nc'

# try and import the dataset, prefer netcdf4, you might want to use pydap also here if you access a dap server.
ds = None 
try:
    import netCDF4 
    ds = netCDF4.Dataset(path)
except:
    import pupynere 
    # use mmap = False if it doesn't fit ...
    ds = pupynere.NetCDFFile(path, mmap=True)

# read the variables (these are just links to the objects for now)
sal = ds.variables['sal'] # salinity
x = ds.variables['x'] # x location (projected)
y = ds.variables['y'] # y location
z = ds.variables['z'] # sigma level
zwl = ds.variables['zwl']  # sea surface elevation
u = ds.variables['u'] # velocity in u direction
v = ds.variables['v'] # v direction
w = ds.variables['w'] # w direction

# we'll slice some data for better performance and less memory. 
# If the grid is not changing (as it does with a sigma layer) this can be done with an extractgrid in the pipeline
# But now we have to update the whole grid and this can be time consuming, so we make use of netcdf slicing to reduce the 
# amount of data we read in
SLICE = True
if SLICE:
    mslice = slice(0,165,3) # from 0 to 165, each 3rd value
    nslice = slice(40,150,3)
    kslice = slice(0,16,2)
else:
    mslice = slice(0,-1,1) # from 0 to end each value (no slice)
    nslice = slice(0,-1,1)
    kslice = slice(0,-1,1)
# read the data into matrices, raise the seawater level a bit (this should be done in the pipeline)
X, Y, Z = x[mslice,nslice], y[mslice,nslice], zwl[:, mslice, nslice] + 5
# we don't want nan's in our coordinates. You could use the mask cell function to hide these cells
if not SLICE:
    X[:,-3:] = [-1750., -1250.,  -750.]

# set the data to a value
if not SLICE:
    Y[:,-3:] = numpy.c_[Y[:,-4], Y[:,-4], Y[:,-4]] # this can be done simpler

# strip of a vector because here the sigma layer is constant over time. 
z = z[1,kslice]

# Now we'll create an n by 3 vector of xyz data
# create an empty matrix and fill in:
pts = empty((prod(X.shape + (len(z),)), 3), dtype=float)
# x
pts[:,0] = numpy.tile(X.T.ravel(), z.shape[0])
# y
pts[:,1] = numpy.tile(Y.T.ravel(), z.shape[0])
# z (also multiply by 3000, this should be done in the pipeline)
pts[:,2] = numpy.repeat(z, prod(Z[1,:,:].shape)) * numpy.tile(Z[1,:,:].T.ravel(), z.shape[0]) * 3000

# Now we create the main data object that stores coordinates together with scalar and vector values
sg = tvtk.StructuredGrid(dimensions=X.shape + (len(z),))
# set the coordinates
sg.points = pts
# set the scalars to the salinity (the convention in netcdf is to use zyx for performance reasons)
sg.point_data.scalars= sal[1,kslice, mslice,nslice].swapaxes(1,2).ravel()
sg.point_data.scalars.name = 'scalars'
# compute the vectors and multiply the w component with the same factor as used for the points
vectors = numpy.c_[u[1,kslice,mslice,nslice].swapaxes(1,2).ravel(), 
                   v[1,kslice,mslice,nslice].swapaxes(1,2).ravel(),
                   w[1,kslice,mslice,nslice].swapaxes(1,2).ravel()*3000]
sg.point_data.vectors = vectors
sg.point_data.vectors.name = 'vectors' 


# We have the data, now we're creating a pipeline for visualization
# start with the source
d = mlab.pipeline.add_dataset(sg)

# we'll extract the grid, so we can tweak the performance. This is only used of slice is not active. If it is
# we slice at the source, not in the pipeline
eg = mlab.pipeline.extract_grid(sg)
# We'll create an extra extract grid to extract the top layer for the water surface
egsurf = mlab.pipeline.extract_grid(sg)
if not SLICE:
    # slice
    eg.set(x_max=165, x_ratio=3, y_min=40, y_max=150, y_ratio=3, z_ratio=3)
    # slice and only extract the top layer
    egsurf.set(x_max=165, x_ratio=3, y_min=40, y_max=150, z_min = 0, z_max=0, y_ratio=3, z_ratio=3)
else:
    # extract the top layer
    egsurf.set(z_min = 0, z_max=0)

# create the water surface using a green to blue color (depending on salinity) and make it a bit transparent
sfsurf = mlab.pipeline.surface(egsurf, opacity=0.3, colormap='GnBu', representation='wireframe' ) 

# create the salinity contours
sf = mlab.pipeline.iso_surface(eg, opacity=0.5) 
sf.contour.auto_contours = False
sf.contour.auto_update_range = False
# show the most interesting contours
sf.contour.contours = [20, 30.0, 31.7, 32.5] 
# show every 20th vector, make it a factor 5000 longer and show as 3D arrows
vf = mlab.pipeline.vectors(eg, mask_points=20, scale_factor=5000, mode='arrow') 

# create a streamline
sl = mlab.pipeline.streamline(eg, linetype='tube', opacity=0.8)
# make it stream longer
sl.stream_tracer.maximum_propagation = 50000.0
# make the tubes thicker
sl.tube_filter.radius = 200.0
# set the size of the sphere
sl.seed.widget.radius = 5608.0572051878689
# make the tubes a bit rounder
sl.tube_filter.number_of_sides = 5
# set the location of the sphere
sl.seed.widget.center = array([-12359.64662114,  51292.88341463,  25157.53124445])
sl.actor.visible = True
# hide it for now...
sl.visible = False

# create a set of line streamlines
sp = mlab.pipeline.streamline(eg, opacity=0.8, seedtype='plane')
# make it stream north/southward
sp.seed.widget.normal_to_y_axis = True
sp.stream_tracer.maximum_propagation = 50000.0
sp.actor.visible = True
# hide it for now
sp.visible = False

# show the outline of the grid
ol = mlab.pipeline.outline(eg)

# show the resolution of the grid in all dimensions
gx = GridPlane()
gx.grid_plane.axis = 'x'
eg.add_module(gx)
gy = GridPlane()
gy.grid_plane.axis = 'y'
eg.add_module(gy)
gz = GridPlane()
gz.grid_plane.axis = 'z'
eg.add_module(gz)

# set some colors and a title
f = mlab.gcf()
f.scene.foreground=(0,0,0)
f.scene.background=(1,1,1)
title = mlab.title('North Sea salinity example')
title.width=0.3

# The image is now done. Now we'll add some functions for interaction

def timestep(i):
    """set the data from timestep(i)"""
    # set a new title
    title.text=str(i)
    title.width=0.03
    title.x_position = 0.1
    title.y_position = 0.9
    # update the grid with new values
    sg.point_data.scalars = sal[i,kslice,mslice,nslice].swapaxes(1,2).ravel()
    sg.point_data.scalars.name = 'scalars'
    # and new locations (because pts is referenced from sg.points, we can just update the source)
    pts[:,2] = numpy.repeat(z, prod(Z[i,:,:].shape)) * numpy.tile(Z[i,:,:].T.ravel(), z.shape[0]) * 7000
    # and new vectors
    vectors = numpy.c_[u[i,kslice,mslice,nslice].swapaxes(1,2).ravel(), v[i,kslice,mslice,nslice].swapaxes(1,2).ravel(), w[i,kslice,mslice,nslice].swapaxes(1,2).ravel()*300]
    sg.point_data.vectors = vectors
    sg.point_data.vectors.name = 'vectors' 
    # flush the pipe by notifying the object is updated
    sg.modified()

# we could do without this
engine = mlab.get_engine()
# now we'll add a method for an animation
# an animation is just an iterator
@mlab.animate(delay=100)
def anim(i=0): 
    scene = engine.scenes[0]
    while True: # keep playing
        # improve performance by temporary disabling rendering
        scene.scene.disable_render = True 
        # make the animation start at the beginning
        i+=1 
        if i >= 77:
            i=1
        # update to timestep i
        timestep(i)
        # start rendering again
        scene.scene.disable_render = False
        value = (yield i) # allows to rewind the animation (not used)
        if value:
            i = value
# now for some more utility functions (all use global variables which is a bit ugly but since they're used during presentation
# they also allow for less typing... 
# switch grid on and off
def gridon():
    """show computational grid"""
    gx.visible = True
    gy.visible = True
    gz.visible = True
def gridoff():
    """turn off computational grid"""
    gx.visible = False
    gy.visible = False
    gz.visible = False
# switches for streamlines
def slon():
    """enable streamlines and disable vector fields"""
    sl.visible = True
    sp.visible = True
    sp.seed.visible = False
    vf.visible = False
def sloff():
    """enable vector field and disable streamline"""
    sl.visible = False
    sp.visible = False
    sp.seed.visible = False
    vf.visible = True


# set the camera in the correct angle
scene = engine.scenes[0]
mlab.view(40, 50)

# scale colors a bit
module_manager = vf.module_manager
module_manager.vector_lut_manager.use_default_range = False
module_manager.vector_lut_manager.data_range = array([ 0. ,  0.7])
module_manager.scalar_lut_manager.use_default_range = False
module_manager.scalar_lut_manager.data_range = array([ 31.0 ,  34.5])

#turn on stereo
renderer = engine.scenes[0].scene.render_window
renderer.stereo_type = 'anaglyph' # 'crystal_eyes' for quad buffered stereo
renderer.stereo_render = 1

mlab.show()