#!/usr/bin/env python
# -*- coding: utf-8 -*-
# run within virtual environment (uses Python 3 syntax)

import argparse
import os
import re
from array import array
from collections import Counter
from subprocess import check_output

def meshFile(param):
    global base
    base, ext = os.path.splitext(param)
    if ext.lower() != '.pyfrm':
        raise argparse.ArgumentTypeError('Mesh file must have a .pyfrm extension')
    return param

def writeHeader(xdmfFile):
    "write XDMF header"
    xdmfFile.write('<?xml version="1.0" ?>\n')
    xdmfFile.write('<!DOCTYPE Xdmf SYSTEM "Xdmf.dtd" []>\n')
    xdmfFile.write('<Xdmf xmlns:xi="http://www.w3.org/2003/XInclude" Version="2.2">\n')
    xdmfFile.write(' <Domain>\n')
    return

def writeTopology(xdmfFile, tType, nCells, connFile):
    "write Topology element"
    xdmfFile.write('   <Topology TopologyType="{}" NumberOfElements="{}">\n'.format(tType, nCells))
    xdmfFile.write('    <xi:include href="{}"/>\n'.format(connFile))
    xdmfFile.write('   </Topology>\n')
    return

def writeGeometry(xdmfFile, nDims, nCells, nodeArray, pyfrm, dataset):
    "write Geometry element"
    nVerts = len(nodeArray)
    if nDims == 2:
        xdmfFile.write('   <Geometry GeometryType="X_Y">\n')   # co-ordinates in separate arrays
    else:
        xdmfFile.write('   <Geometry GeometryType="X_Y_Z">\n') # co-ordinates in separate arrays
    for coord in range(nDims):
        xdmfFile.write('    <DataItem ItemType="Function" Dimensions="{}"\n'.format(nVerts*nCells)) # 1D-array
        xdmfFile.write('      Function="JOIN({})">\n'.format(' ; '.join("$" + str(k) for k in range(nVerts))))
        writeHyperSlab(xdmfFile, coord, nDims, nCells, nodeArray, pyfrm, dataset)
        xdmfFile.write('    </DataItem>\n')
    xdmfFile.write('   </Geometry>\n')
    return

def writeHyperSlab(xdmfFile, coord, nDims, nCells, nodeArray, pyfrm, dataset):
    "write HyperSlab element"
    nVerts = len(nodeArray)
    for vert in nodeArray:
        xdmfFile.write('      <DataItem ItemType="HyperSlab"\n')
        xdmfFile.write('        Dimensions="{} 1 1"\n'.format(nCells))
        xdmfFile.write('        Type="HyperSlab">\n')
        xdmfFile.write('        <DataItem\n') # start, stride and count of hyperslab region
        xdmfFile.write('         Dimensions="3 3"\n')
        xdmfFile.write('         Format="XML">\n')
        xdmfFile.write('         {:<3} 0   {}\n'.format(vert, coord)) # select vertex and co-ordinate (format is vertex, cell, co-ordinate)
        xdmfFile.write('         1   1   1\n')             # select every cell, for this vertex and co-ordinate
        xdmfFile.write('         1   {} 1\n'.format(nCells)) # loop over cells
        xdmfFile.write('         </DataItem>\n')
        xdmfFile.write('         <DataItem\n')
        xdmfFile.write('         Name="Points" \n')
        xdmfFile.write('         Dimensions="{} {} {}"\n'.format(nVerts, nCells, nDims))
        xdmfFile.write('         Format="HDF">\n')
        xdmfFile.write('         {}:/{}\n'.format(pyfrm, dataset))
        xdmfFile.write('        </DataItem>\n')
        xdmfFile.write('      </DataItem>\n')
    return

def writeAttribute(xdmfFile, tag):
    "write Attribute element"
    xdmfFile.write('   <Attribute Name="Partition" Center="Grid">\n')
    xdmfFile.write('    <DataItem\n')
    xdmfFile.write('     Dimensions="1"\n')
    xdmfFile.write('     Format="XML">\n')
    xdmfFile.write('     {}\n'.format(tag)) # tag with partition number
    xdmfFile.write('    </DataItem>\n')
    xdmfFile.write('   </Attribute>\n')
    return

def writeConnectivities(connFile, nCells, nVerts):
    "write connectivities to xml file"
    cf = open(connFile, 'w')
    cf.write('<DataItem DataType="Int"\n')
    cf.write('  Dimensions="{} {}"\n'.format(nCells, nVerts))
    cf.write('  Format="XML">\n')
    
    for i in range (0, nCells):
        cf.write(' ')
        for j in range(nVerts):
            cf.write(' ' + repr(j*nCells+i).ljust(1))
        cf.write('\n')
    
    cf.write('</DataItem>\n')
    cf.close()
    print('connectivities written to ' + connFile)
    return

def writeFooter(xdmfFile):
    "write XDMF footer"
    xdmfFile.write(' </Domain>\n')
    xdmfFile.write('</Xdmf>\n')
    return

def readH5lsOutput(os_output):
    "read output from 'h5ls' command"
    tet4,  tet10, tet20,  tet35,  tet56,  tet84  = ({} for i in range(6)) # tet cells
    pri6,  pri18, pri40,  pri75,  pri126, pri196 = ({} for i in range(6)) # pri cells
    pyr5,  pyr14, pyr30,  pyr55,  pyr91,  pyr140 = ({} for i in range(6)) # pyr cells
    hex8,  hex27, hex64,  hex125, hex216  = ({} for i in range(5))        # hex cells
    tri3,  tri6,  tri10,  tri15,  tri21   = ({} for i in range(5))        # tri cells
    quad4, quad9, quad16, quad25, quad36  = ({} for i in range(5))        # quad cells

    for line in os_output.splitlines():
        spt = re.search('spt', line.decode())     # restrict to 'spt' arrays
        if spt:
            chunk = line.decode().split()
            partno = int(re.search('\d+', chunk[0]).group())
            nnodes = int(re.search('(\d+),', line.decode()).group(1))
            ncells = int(re.search(' (\d+),', line.decode()).group(1))
            if re.search('quad', line.decode()):  # check whether cell is quadrangular
                if nnodes == 4:
                    quad4[partno] = ncells
                elif nnodes == 9:
                    quad9[partno] = ncells
                elif nnodes == 16:
                    quad16[partno] = ncells
                elif nnodes == 25:
                    quad25[partno] = ncells
                elif nnodes == 36:
                    quad36[partno] = ncells
                else:
                    print("unknown cell order")
            elif re.search('tri', line.decode()): # check whether cell is triangular
                if nnodes == 3:
                    tri3[partno] = ncells
                elif nnodes == 6:
                    tri6[partno] = ncells
                elif nnodes == 10:
                    tri10[partno] = ncells
                elif nnodes == 15:
                    tri15[partno] = ncells
                elif nnodes == 21:
                    tri21[partno] = ncells
                else:
                    print("unknown cell order")
            elif re.search('hex', line.decode()): # check whether cell is hexahedral
                if nnodes == 8:
                    hex8[partno] = ncells
                elif nnodes == 27:
                    hex27[partno] = ncells
                elif nnodes == 64:
                    hex64[partno] = ncells
                elif nnodes == 125:
                    hex125[partno] = ncells
                elif nnodes == 216:
                    hex216[partno] = ncells
                else:
                    print("unknown cell order")
            else:
                print("unknown cell type")
                break

    return (tet4,  tet10, tet20,  tet35,  tet56,  tet84,
            pri6,  pri18, pri40,  pri75,  pri126, pri196,
            pyr5,  pyr14, pyr30,  pyr55,  pyr91,  pyr140,
            hex8,  hex27, hex64,  hex125, hex216,
            tri3,  tri6,  tri10,  tri15,  tri21,
            quad4, quad9, quad16, quad25, quad36)

# read command line arguments
parser = argparse.ArgumentParser(description="extract connectivities from mesh file")
parser.add_argument("mesh", help="mesh file (.pyfrm)", type=meshFile)
args = parser.parse_args()

# use 'h5ls' command to provide array dimensions
h5lsOutput = check_output(["h5ls", args.mesh])

(ntet4,  ntet10, ntet20,  ntet35,  ntet56,  ntet84,
npri6,  npri18, npri40,  npri75,  npri126, npri196,
npyr5,  npyr14, npyr30,  npyr55,  npyr91,  npyr140,
nhex8,  nhex27, nhex64,  nhex125, nhex216,
ntri3,  ntri6,  ntri10,  ntri15,  ntri21,
nquad4, nquad9, nquad16, nquad25, nquad36) = readH5lsOutput(h5lsOutput)

# cell types
firstOrderCellType = ['tet', 'tet', 'tet', 'tet', 'tet', 'tet',
                      'pri', 'pri', 'pri', 'pri', 'pri', 'pri',
                      'pyr', 'pyr', 'pyr', 'pyr', 'pyr', 'pyr',
                      'hex', 'hex', 'hex', 'hex', 'hex',
                      'tri', 'tri', 'tri', 'tri', 'tri',
                      'quad', 'quad', 'quad', 'quad', 'quad']
ndims              = {'quad': 2, 'tri': 2, 'hex': 3, 'pyr': 3, 'pri': 3, 'tet':3}
xdmfTopologyType   = {'quad': 'Quadrilateral', 'tri': 'Triangle', 'hex': 'Hexahedron', 
                      'pyr' : 'Pyramid',       'pri': 'Wedge',    'tet': 'Tetrahedron'}

# node identification (see pyfr/readers/nodemaps.py): reduce high-order cells to first order
nodeIDs = {}
nodeIDs[29] = [0, 2, 8, 6]                 # quad9
nodeIDs[28] = [0, 1, 3, 2]                 # quad4
nodeIDs[23] = [0, 1, 2]                    # tri3
nodeIDs[19] = [0, 2, 8, 6, 18, 20, 26, 24] # hex27
nodeIDs[18] = [0, 1, 3, 2, 4, 5, 7, 6]     # hex8

# sort datasets
hex8Keys  = list(nhex8.keys())    # keys are partition numbers
hex27Keys = list(nhex27.keys())
quad4Keys = list(nquad4.keys())
tri3Keys  = list(ntri3.keys())
allKeys   = quad4Keys + tri3Keys + hex8Keys + hex27Keys   # concatenate keys
numTypes  = Counter(allKeys)        # number of types present in each partition
partKeys  = list(numTypes.keys())   # partition keys
partitions = [ntet4,  ntet10, ntet20,  ntet35,  ntet56,  ntet84,
              npri6,  npri18, npri40,  npri75,  npri126, npri196,
              npyr5,  npyr14, npyr30,  npyr55,  npyr91,  npyr140,
              nhex8,  nhex27, nhex64,  nhex125, nhex216,
              ntri3,  ntri6,  ntri10,  ntri15,  ntri21,
              nquad4, nquad9, nquad16, nquad25, nquad36]   # list of cell dictionaries
numCellTypes = len(partitions)

if partKeys:
    # write files
    g = open(os.path.join(base + '.xdmf'), 'w')
    writeHeader(g)
    
    for part in partKeys:
        if numTypes[part] > 1: # check whether partition contain multiple cell types
            g.write('  <Grid Name="Partition{}" GridType="Collection">\n'.format(part))
        else:
            g.write('  <Grid Name="Partition{}" GridType="Uniform">\n'.format(part))
    
        for cellType in range(numCellTypes):
            if part in partitions[cellType]: # check whether these cells exist in this partition
                xfname = os.path.join('con_' + firstOrderCellType[cellType] + '_p' + str(part) + '.xml')
                dname  = os.path.join('spt_' + firstOrderCellType[cellType] + '_p' + str(part))
                if numTypes[part] > 1:
                    g.write('  <Grid Name="cellType{}" GridType="Uniform">\n'.format(cellType))
                writeTopology(g, xdmfTopologyType[firstOrderCellType[cellType]],
                              partitions[cellType][part], xfname)
                writeGeometry(g, ndims[firstOrderCellType[cellType]], 
                               partitions[cellType][part], 
                               nodeIDs[cellType], 
                               args.mesh, dname)
                if numTypes[part] > 1:
                    g.write('  </Grid>\n')
                # connectivities file
                writeConnectivities(xfname, partitions[cellType][part], 
                                    len(nodeIDs[cellType]))
    
        writeAttribute(g, part)
        g.write('  </Grid>\n')
    
    writeFooter(g)
    g.close()