#!/usr/bin/env python
"""
Viper: A simple mesh plotter and run--time visualization module for plotting
and saving simulation data. The class C{Viper} can visualize solutions given
as numpy arrays, and meshes that provide the two methods C{cells()} and
C{coordinates()}. These methods should return numpy arrays specifying the
node-element ordering and coordinates of the nodes, respectively.

To use the module as a script, the syntax is:
viper.py -i <infile.xml[.gz]> [-o <savefile>]

or,
viper.py --input <infile.xml[.gz]> [--output <savefile>]

If the optional command line argument -o is given, the mesh is stored in vtk
format.
"""

from viper import *
from viper import __version__
import sys
import getopt
import numpy
import os

def _dolfin_mesh_plotter(infile, outfile, data=None):
    from dolfin import Mesh, UnitSquare, File, plot
    from dolfin.dolfin import MeshFunctionInt
    if infile:
        mesh = Mesh(infile)
    else:
        mesh = UnitSquare(10,10)
    if data is None:
        n = mesh.numVertices()
        x = numpy.zeros((n,), dtype='d')
        p = Viper(mesh, x, 0, 1)
        p.update()
    else:
        file = File(data)
        mf = MeshFunctionInt(mesh, mesh.topology().dim())
        file >> mf
        p = plot(mf)
    if outfile: 
        p.init_writer("")
        p.write_vtk(os.path.splitext(os.path.splitext(outfile)[0])[0]+".vtk")
    p.interactive()
    return p

def _vtk_plotter(infile):
    p = Viper()
    p.init_from_file(infile)
    p.interactive()
    return p

def _main(infile, outfile, data=None):
    if infile:
        suffix = os.path.splitext(infile)[-1]
        if suffix == ".vtk":
            return _vtk_plotter(infile)

    return _dolfin_mesh_plotter(infile, outfile, data=data) 
        
def _help():
    print """This is Viper, the simple FEniCS run-time plotter, version %s.
For further information, go to http://www.fenics.org/wiki/Viper

Usage: viper -i file [-d data (xml meshfunction file)] -o file

""" % __version__

if __name__ == '__main__':
    import sys
    (opts, args) = getopt.getopt(sys.argv[1:], "i:d:o:h", ["input=", "data=", "output=", "help"])
    infile = None
    outfile = None
    datafile = None
    for (opt, val) in opts:
        print opt
        if (opt == "--input" or opt == "-i"):
            infile = val
        elif (opt =="--output" or opt == "-o"):
            outfile = val
        elif (opt =="--data" or opt == "-d"):
            datafile = val
        elif (opt == "--help" or opt == "-h"):
            _help()
            sys.exit()
    if infile == outfile == None:
        if len(args) > 0:
            _main(args[-1], None)
            sys.exit()

    p = _main(infile, outfile, data=datafile)
