#!/usr/bin/env python
#
# Copyright (C) 2006 Anders Logg
# Licensed under the GNU GPL Version 2
#
# Modified by Garth N. Wells (gmsh function)
# Modified by Alexander H. Jarosch (gmsh fix)
#
# Script for converting between various data formats.
#
# Supported input formats:
#
#     .mesh       (Medit, generated by tetgen with option -g)
#     .gmsh, .msh (Gmsh, version 2.0 file format)
#
# Supported output formats:
#
#     .xml  (DOLFIN XML mesh file)
#
# In addition, a special conversion from the old DOLFIN
# XML format (0.6.1) to the new (0.6.2) is supported.

import getopt
import sys
from commands import getoutput

def main(argv):
    "Main function"

    # Get command-line arguments
    try:
        opts, args = getopt.getopt(argv, "h", ["help"])
    except getopt.GetoptError:
        usage()
        sys.exit(2)
        
    # Get options
    for opt, arg in opts:
        if opt in ("-h", "--help"):
            usage()
            sys.exit()
            
    # Check that we got two filenames
    if not len(args) == 2:
        usage()
        sys.exit(2)

    # Get filenames and suffixes
    ifilename = args[0]
    ofilename = args[1]
    isuffix = ifilename.split(".")[-1]
    osuffix = ofilename.split(".")[-1]

    # Choose conversion
    if isuffix == "mesh" and osuffix == "xml":
        mesh2xml(ifilename, ofilename)
    elif (isuffix == "gmsh" or isuffix == "msh") and osuffix == "xml":
        gmsh2xml(ifilename, ofilename)
    elif isuffix == "old" and osuffix == "xml":
        oldxml2newxml(ifilename, ofilename)
    else:
        print "Sorry, cannot convert between .%s and .%s file formats." % (isuffix, osuffix)
        
def usage():
    "Display usage"
    print """\
Usage: dolfin-convert [OPTION]... input.x output.y

  -h  display this help text and exit

Alternatively, the following long options may be used:

  --help
  """
    return

def error(message):
    "Write an error message"
    for line in message.split("\n"):
        print "*** %s" % line
    sys.exit(2)

def mesh2xml(ifilename, ofilename):
    """Convert between .mesh and .xml, parser implemented as a
    state machine:
    
        0 = read 'Dimension'
        1 = read dimension
        2 = read 'Vertices'
        3 = read number of vertices
        4 = read next vertex
        5 = read 'Triangles' or 'Tetrahedra'
        6 = read number of cells
        7 = read next cell
        8 = done
       
    """
    
    print "Converting from Medit format (.mesh) to DOLFIN XML format"

    # Open files
    ifile = open(ifilename, "r")
    ofile = open(ofilename, "w")

    # Write header
    write_header(ofile)

    # Current state
    state = 0

    # Write data
    num_vertices_read = 0
    num_cells_read = 0
    while 1:

        # Read next line
        line = ifile.readline()
        if not line: break

        # Skip comments
        if line[0] == '#':
            continue

        # Remove newline
        if line[-1] == "\n":
            line = line[:-1]

        if state == 0:
            if line == "Dimension":
                state += 1
        elif state == 1:
            num_dims = int(line)
            state +=1
        elif state == 2:
            if line == "Vertices":
                state += 1
        elif state == 3:
            num_vertices = int(line)
            write_header_vertices(ofile, num_vertices)
            state +=1
        elif state == 4:
            if num_dims == 2:
                (x, y, tmp) = line.split()
                x = float(x)
                y = float(y)
                z = 0.0
            elif num_dims == 3:
                (x, y, z, tmp) = line.split()
                x = float(x)
                y = float(y)
                z = float(z)
            write_vertex(ofile, num_vertices_read, x, y, z)
            num_vertices_read +=1
            if num_vertices == num_vertices_read:
                write_footer_vertices(ofile)
                state += 1
        elif state == 5:
            if line == "Triangles" and num_dims == 2:
                state += 1
            if line == "Tetrahedra" and num_dims == 3:
                state += 1
        elif state == 6:
            num_cells = int(line)
            write_header_cells(ofile, num_cells)
            state +=1
            min_vertex = 1000000
        elif state == 7:
            if num_dims == 2:
                (n0, n1, n2, tmp) = line.split()
                n0 = int(n0) - 1
                n1 = int(n1) - 1
                n2 = int(n2) - 1
                write_cell_triangle(ofile, num_cells_read, n0, n1, n2)
            elif num_dims == 3:
                (n0, n1, n2, n3, tmp) = line.split()
                n0 = int(n0) - 1
                n1 = int(n1) - 1
                n2 = int(n2) - 1
                n3 = int(n3) - 1
                write_cell_tetrahedron(ofile, num_cells_read, n0, n1, n2, n3)
            num_cells_read +=1
            if num_cells == num_cells_read:
                write_footer_cells(ofile)                
                state += 1
        elif state == 8:
            break

        #print line

    # Check that we got all data
    if state == 8:
        print "Conversion done"
    else:
        error("Missing data, unable to convert")

    # Write footer
    write_footer(ofile)

    # Close files
    ifile.close()
    ofile.close()

def gmsh2xml(ifilename, ofilename):
    """Convert between .gmsh v2.0 format (http://www.geuz.org/gmsh/) and .xml, 
    parser implemented as a state machine:
    
        0 = read 'MeshFormat'
        1 = read  mesh format data
        2 = read 'EndMeshFormat'
        3 = read 'Nodes'
        4 = read  number of vertices
        5 = read  vertices
        6 = read 'EndNodes'
        7 = read 'Elements'
        8 = read  number of cells
        9 = read  cells
        10 = done
       
    """
    
    print "Converting from Gmsh format (.msh, .gmsh) to DOLFIN XML format"

    # Open files
    ifile = open(ifilename, "r")
    ofile = open(ofilename, "w")

    # Write header
    write_header(ofile)
    
    # Initialise node list (gmsh does not export all vertexes in order)
    nodelist = {}
    
    # Current state
    state = 0
    
    # Write data
    num_vertices_read = 0
    num_cells_read = 0
    
    while 1:

        # Read next line
        line = ifile.readline()
        if not line: break

        # Skip comments
        if line[0] == '#':
            continue

        # Remove newline
        if line[-1] == "\n":
            line = line[:-1]

        if state == 0:
            if line == "$MeshFormat":
                state += 1
        elif state == 1:
            (version, file_type, data_size) = line.split()
            state += 1
        elif state == 2:
            if line == "$EndMeshFormat":
                state += 1
        elif state == 3:
            if line == "$Nodes":
                state += 1
        elif state == 4:
            num_vertices = int(line)
            write_header_vertices(ofile, num_vertices)
            state +=1
        elif state == 5:
            (node_no, x, y, z) = line.split()
            nodelist[int(node_no)] = num_vertices_read
            x = float(x)
            y = float(y)
            z = float(z)
            write_vertex(ofile, num_vertices_read, x, y, z)
            num_vertices_read +=1
            
            if num_vertices == num_vertices_read:
                write_footer_vertices(ofile)                
                state += 1
        elif state == 6:
            if line == "$EndNodes":
                state += 1
        elif state == 7:
            if line == "$Elements":
                state += 1
        elif state == 8:
            num_cells  = int(line)
            if num_cells == 0:
                print "No cells found in gmsh file."
                break
            else:
                write_header_cells(ofile, num_cells)   
                state += 1
        elif state == 9:
            element = line.split()
            cell_type = int(element[1])
            num_tags  = int(element[2])
            if cell_type == 2:
                nn = [int(node) for node in element[3 + num_tags:8]]
                for node in nn:
                    if not node in nodelist:
                        error("Vertex %d of triangle %d not previously defined." % (node, num_cells_read))
                n0 = nodelist[nn[0]]
                n1 = nodelist[nn[1]]
                n2 = nodelist[nn[2]]
                write_cell_triangle(ofile, num_cells_read, n0, n1, n2)
            elif cell_type == 4:
                nn = [int(node) for node in element[3 + num_tags:9]]
                for node in nn:
                    if not node in nodelist:
                        error("Vertex %d of tetrahedron %d not previously defined." % (node, num_cells_read))
                n0 = nodelist[nn[0]]
                n1 = nodelist[nn[1]]
                n2 = nodelist[nn[2]]
                n3 = nodelist[nn[3]]
                write_cell_tetrahedron(ofile, num_cells_read, n0, n1, n2, n3)

            num_cells_read +=1 
            if num_cells == num_cells_read:
              write_footer_cells(ofile)                
              state += 1
        elif state == 10:
            break

    # Check that we got all data
    if state == 10:
        print "Conversion done"
    else:
        print "*** Missing data, unable to convert"
        print "*** Did you use version 2.0 of the gmsh file format?"
    
    # Write footer
    write_footer(ofile)  

    # Close files
    ifile.close()
    ofile.close()

def oldxml2newxml(ifilename, ofilename):
    "Convert from old DOLFIN XML format to new."

    print "Converting from old to new DOLFIN XML format..."

    # Get dimension
    tris = len(getoutput("grep triangle " + ifilename))
    tets = len(getoutput("grep tetrahedron " + ifilename))
    if tris > 0 and tets == 0:
        dim = 2
        celltype = "triangle"
    elif tris == 0 and tets > 0:
        dim = 3
        celltype = "tetrahedron"
    else:
        error("Inconsistent mesh file, not exactly one type of cells.")

    # Open files
    ifile = open(ifilename, "r")
    ofile = open(ofilename, "w")

    # Read lines and make changes
    while 1:

        # Read next line
        line = ifile.readline()
        if not line: break

        # Modify line
        if "xmlns" in line:
            line = "<dolfin xmlns:dolfin=\"http://www.fenics.org/dolfin/\">\n"
        if "<mesh>" in line:
            line = "  <mesh dim=\"%d\" celltype=\"%s\">\n" % (dim, celltype)
        if dim == 2 and " z=\"0.0\"" in line:
            line = line.replace(" z=\"0.0\"", "")
        if " name=" in line:
            line = line.replace(" name=", " index=")
        if "n0" in line:
            line = line.replace("n0", "v0")
        if "n1" in line:
            line = line.replace("n1", "v1")
        if "n2" in line:
            line = line.replace("n2", "v2")
        if "n3" in line:
            line = line.replace("n3", "v3")
            
        # Write line
        ofile.write(line)

    # Close files
    ifile.close();
    ofile.close();

# Write mesh header
def write_header(ofile):
    ofile.write("""\
<?xml version=\"1.0\" encoding=\"UTF-8\"?>

<dolfin xmlns:dolfin=\"http://www.fenics.org/dolfin/\">
  <mesh>
""")

# Write mesh footer
def write_footer(ofile):
    ofile.write("""\
  </mesh>
</dolfin>
""")

def write_header_vertices(ofile, num_vertices):
    "Write vertices header"
    print "Expecting %d vertices" % num_vertices
    ofile.write("    <vertices size=\"%d\">\n" % num_vertices)

def write_footer_vertices(ofile):
    "Write vertices footer"
    ofile.write("    </vertices>\n")
    print "Found all vertices"

def write_vertex(ofile, vertex, x, y, z):
    "Write vertex"
    ofile.write("      <vertex name=\"%d\" x=\"%g\" y=\"%g\" z=\"%g\"/>\n" % \
        (vertex, x, y, z))

def write_header_cells(ofile, num_cells):
    "Write cells header"
    ofile.write("    <cells size=\"%d\">\n" % num_cells)
    print "Expecting %d cells" % num_cells

def write_footer_cells(ofile):
    "Write cells footer"
    ofile.write("    </cells>\n")
    print "Found all cells"

def write_cell_triangle(ofile, cell, n0, n1, n2):
    "Write cell (triangle)"
    ofile.write("      <triangle name=\"%d\" n0=\"%d\" n1=\"%d\" n2=\"%d\"/>\n" % \
        (cell, n0, n1, n2))

def write_cell_tetrahedron(ofile, cell, n0, n1, n2, n3):
    "Write cell (tetrahedron)"
    ofile.write("      <tetrahedron name=\"%d\" n0=\"%d\" n1=\"%d\" n2=\"%d\" n3=\"%d\"/>\n" % \
        (cell, n0, n1, n2, n3))

if __name__ == "__main__":
    main(sys.argv[1:])

