#!/usr/bin/env python

"Script for simple command-line plotting."

# Copyright (C) 2010 Anders Logg
#
# This file is part of DOLFIN.
#
# DOLFIN is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# DOLFIN is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with DOLFIN. If not, see <http://www.gnu.org/licenses/>.
#
# Modified by Marie E. Rognes, 2010
#
# First added:  2010-12-08
# Last changed: 2010-12-22

import sys

# Try importing DOLFIN
try:
    from dolfin import *
except:
    print """\
Unable to import DOLFIN. The DOLFIN Python module is required
to run this script. Check that DOLFIN is in your PYTHONPATH."""
    sys.exit(1)

def usage():
    "Print usage of this script."

    # Build list of supported elements
    element_list = ""
    for e in supported_elements_for_plotting:
        if e in supported_elements:
            element_list += ("   %s\n" % e)
        else:
            element_list += ("   %s (*)\n" % e)

    print """\
Usage:

1. dolfin-plot <meshfile>

   where <meshfile> is a mesh stored in DOLFIN XML format
   with suffix .xml or .xml.gz.

2. dolfin-plot <family> <domain> [degree] [rotate=1/0]

   where <family> is the name of a finite element family,
   <domain> is the domain type ('triangle' or 'tetrahedron'),
   and <degree> is an optional degree for elements that
   support variable degree.

Examples:

   dolfin-plot mesh.xml
   dolfin-plot Lagrange triangle 3
   dolfin-plot BDM tetrahedron 5
   dolfin-plot Hermite triangle

List of supported element families:

%s
A (*) indicates that the element is not supported by DOLFIN/FFC,
but the element may still be plotted.
""" % element_list

# Check command-line arguments
if len(sys.argv) < 2:
    usage()
    sys.exit(1)

# Check for help text
if "-h" in sys.argv[1:] or "--help" in sys.argv[1:]:
    usage()
    sys.exit(0)

# Extract arguments and options
args = [arg for arg in sys.argv[1:] if not "=" in arg]
options = dict([arg.split("=") for arg in sys.argv[1:] if "=" in arg])

# Check for plotting of mesh
if len(args) == 1:

    # Read mesh
    print "Reading mesh from file '%s'." % args[0]
    try:
        mesh = Mesh(args[0])
    except:
        print "Unable to read mesh from file."
        sys.exit(1)

    # Plot mesh
    print "Plotting mesh."
    plot(mesh, title="Mesh", interactive=True)

    sys.exit(0)

# Check for plotting of element
if len(args) in (2, 3, 4):

    # Get family and domain
    family, domain = args[:2]

    # Get degree
    if len(args) == 2:
        degree = None
    else:
        degree = int(args[2])


    # Create element
    print "Creating finite element."

    if len(args) == 3:
        element = FiniteElement(family, domain, degree)
    elif len(args) == 4:
        assert family == "P Lambda" or family == "P- Lambda"
        form_degree = int(args[3])
        element = FiniteElement(family, domain, degree, form_degree)
    else:
        pass

    # Rotate by default in 3D
    rotate = 1 if str(domain) == "tetrahedron" else 0

    # Check for rotate option
    if "rotate" in options:
        rotate = int(options["rotate"])

    # Plot element
    print "Plotting finite element."
    plot(element, rotate=rotate)

    sys.exit(0)

# Catch all
usage()
sys.exit(1)
