#!/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, lutfile="gauss_120.lut", datatype=None):
    from dolfin import Mesh, UnitSquare, File, plot
    from dolfin import MeshFunction
    if infile:
        title = "Mesh plot"
        mesh = Mesh(infile)
    else:
        title = "Unit Square plot"
        mesh = UnitSquare(10,10)
    if data is None:
        p = Viper(mesh=mesh, title=title)
        p.update()
    else:
        assert  os.path.isfile(lutfile) or \
os.path.isfile(os.path.join(os.path.dirname(viper.__file__), "data", lutfile)), \
"Could not locate lutfile \"%s\" in current directory or default data dir." % (lutfile,)
        file = File(data)
        
        dim = datatype["dim"]
        if datatype["type"] == "i":
            mf = MeshFunction("int", mesh, dim)
            file >> mf
        elif datatype["type"] == "u":
            mf = MeshFunction("uint", mesh, dim)
            file >> mf
        elif datatype["type"] == "r":
            mf = MeshFunction("real", mesh, dim)
            file >> mf
        else:
            print _help()
            sys.exit()
        p = plot(mf, lutfile=lutfile, title="Data plot")
    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, lutfile="gauss_120.lut", datatype={}):
    if infile:
        suffix = os.path.splitext(infile)[-1]
        if suffix == ".vtk":
            return _vtk_plotter(infile)

    return _dolfin_mesh_plotter(infile, outfile, data=data, lutfile=lutfile, datatype=datatype) 
        
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) -tdatatype -l lutfile -o file ]
Alternatively: viper file (mesh file in dolfin mesh format)

datatype is a format string specifying i or r for integer or real valued data,
and the dimensionality of the data set. E.g. -tr0 (real vertex data), -ti2
(integer cell data).

""" % __version__

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

Usage:
viper file [--lut=lut_file --mode=mode --title=title]
or
viper file [-l lut_file -m mode -t title]

where 'file' is is a dolfin data file, like a Mesh, FunctionPlotData, or a file
containing both a Mesh and either a cell or vertex valued MeshFunction of type
uint or double. The 'mode' argument can be either 'auto', 'scalar', or
'vector', 'lut_file' a color lookup table, and 'title' the title field of the
plotting window.
""" % (__version__,)

def _plot_function_plot_data(infile, lutfile, mode, title):
    from dolfin.cpp import FunctionPlotData, File
    from dolfin import plot
    d = FunctionPlotData()
    f = File(infile)
    f >> d
    return plot(d, lutfile=lutfile, mode=mode, title=title)

def _plot_mesh(infile, lutfile, mode, title):
    import dolfin
    f = dolfin.File(infile)
    d = dolfin.Mesh()
    f >> d
    return dolfin.plot(d, title=title)

def _plot_mesh_function(infile, lutfile, mode, title):
    import dolfin
    f = dolfin.File(infile)
    mesh = dolfin.Mesh()
    f >> mesh
    try:
        f = dolfin.File(infile)
        mf = dolfin.MeshFunction("double", mesh, mesh.topology().dim())
        f >> mf
        return dolfin.plot(mf, title=title)
    except:
        pass
    try:
        f = dolfin.File(infile)
        mf = dolfin.MeshFunction("uint", mesh, mesh.topology().dim())
        f >> mf
        return dolfin.plot(mf, title=title)
    except:
        raise ValueError, "No meshfunction found."
    

def __main(infile, lutfile, mode, title):
    try:
        return _plot_function_plot_data(infile, lutfile, mode, title)
    except:
        pass
    try:
        return _plot_mesh_function(infile, lutfile, mode, title)
    except:
        pass
    return _plot_mesh(infile, lutfile, mode, title)
    print "Cannot plot data in file %s" % (infile,)
        

def _oldmain():
    import sys
    (opts, args) = getopt.getopt(sys.argv[1:], "i:d:o:l:h:t:", ["input=", "data=", "output=", "lut=", "help", "datatype="])
    infile = None
    outfile = None
    datafile = None
    datatype = {"type": "r", "dim": 0}
    lutfile = "gauss_120.lut"
    for (opt, val) in opts:
        if (opt == "--input" or opt == "-i"):
            infile = val
        elif (opt =="--output" or opt == "-o"):
            outfile = val
        elif (opt =="--lut" or opt == "-l"):
            lutfile = val
        elif (opt =="--datatype" or opt == "-t"):
            import re
            if re.search("i", val):
                datatype["type"] = "i"
            if re.search("u", val):
                datatype["type"] = "u"
            match = re.search("(\d+)", val)
            if match:
                datatype["dim"] = int(match.groups()[0])
        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, lutfile=lutfile, datatype=datatype)

def _newmain():
    import sys, os
    (opts, args) = getopt.getopt(sys.argv[1:], "o:l:hm:", ["output=", "lut=", "mode=", "title=", "help",])
    outfile = None
    lutfile = "gauss_120.lut"
    mode = "auto"
    title = "FEniCS Viper"

    for (opt, val) in opts:
        if (opt in ("--output", "-o")):
            outfile = val
        elif (opt in ("--lut", "-l")):
            lutfile = val
        elif (opt in ("--help", "-h")):
            _newhelp()
        elif (opt in ("--mode", "-m")):
            mode = val
        elif (opt in ("--title", "-t")):
            title = val
    
    if len(args) > 0:
        infile = args[-1]
        if not os.path.isfile(infile):
            raise RuntimeError, "No such file %s" % (infile,)
    else:
        raise RuntimeError, _usage()
    
    if mode == "":
        mode = "auto"
    p = __main(infile, lutfile, mode, title)
    p.interactive()
    
if __name__ == '__main__':
    #_oldmain()
    _newmain()

