#!/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')
    if not os.path.exists(param):
        raise argparse.ArgumentTypeError("{} not found".format(param))
    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 cells are quadrilaterals
                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 cells are triangles
                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 cells are hexahedrons
                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")
            elif re.search('pyr', line.decode()):  # check whether cells are pyramids
                if nnodes == 5:
                    pyr5[partno] = ncells
                elif nnodes == 14:
                    pyr14[partno] = ncells
                elif nnodes == 30:
                    pyr30[partno] = ncells
                elif nnodes == 55:
                    pyr55[partno] = ncells
                elif nnodes == 91:
                    pyr91[partno] = ncells
                elif nnodes == 140:
                    pyr140[partno] = ncells
                else:
                    print("unknown cell order")
            elif re.search('pri', line.decode()): # check whether cells are prisms
                if nnodes == 6:
                    pri6[partno] = ncells
                elif nnodes == 18:
                    pri18[partno] = ncells
                elif nnodes == 40:
                    pri40[partno] = ncells
                elif nnodes == 75:
                    pri75[partno] = ncells
                elif nnodes == 126:
                    pri126[partno] = ncells
                elif nnodes == 196:
                    pri196[partno] = ncells
                else:
                    print("unknown cell order")
            elif re.search('tet', line.decode()): # check whether cells are tetrahedrons
                if nnodes == 4:
                    tet4[partno] = ncells
                elif nnodes == 10:
                    tet10[partno] = ncells
                elif nnodes == 20:
                    tet20[partno] = ncells
                elif nnodes == 35:
                    tet35[partno] = ncells
                elif nnodes == 56:
                    tet56[partno] = ncells
                elif nnodes == 84:
                    tet84[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): reduces high-order cells to first order
# example, for second-order pyramid:
# >>> from pyfr.readers.nodemaps import GmshNodeMaps
# >>> [GmshNodeMaps.to_pyfr['pyr', 14][i] for i in range(5)]
nodeIDs = {}
nodeIDs[29] = [0, 2, 8, 6]                 # quad9
nodeIDs[28] = [0, 1, 3, 2]                 # quad4
nodeIDs[24] = [0, 2, 5]                    # tri6
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
nodeIDs[13] = [0, 2, 8, 6, 13]             # pyr14
nodeIDs[12] = [0, 1, 3, 2, 4]              # pyr5
nodeIDs[7]  = [0, 2, 5, 12, 14, 17]        # pri18
nodeIDs[6]  = [0, 1, 2, 3, 4, 5]           # pri6
nodeIDs[1]  = [0, 2, 5, 9]                 # tet10
nodeIDs[0]  = [0, 1, 2, 3]                 # tet4

# keys are partition numbers
tet4Keys  = list(ntet4.keys())    
tet10Keys = list(ntet10.keys())
pri6Keys  = list(npri6.keys())    
pri18Keys = list(npri18.keys())
pyr5Keys  = list(npyr5.keys())    
pyr14Keys = list(npyr14.keys())
hex8Keys  = list(nhex8.keys())    
hex27Keys = list(nhex27.keys())
tri3Keys  = list(ntri3.keys())
tri6Keys  = list(ntri6.keys())
quad4Keys = list(nquad4.keys())
quad9Keys = list(nquad9.keys())
# concatenate keys
allKeys   = (tet4Keys + tet10Keys +
             pri6Keys + pri18Keys +
             pyr5Keys + pyr14Keys +
             hex8Keys + hex27Keys +
             tri3Keys + tri6Keys +
             quad4Keys + quad9Keys)
# number of types present in each partition
numTypes  = Counter(allKeys)        
# list of partition keys
partKeys  = list(numTypes.keys())   
# list of cell dictionaries
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]   
numCellTypes = len(partitions)

if partKeys:
    # write files
    g = open(os.path.join(base + '.xdmf'), 'w')
    writeHeader(g)
    
    for part in partKeys:
        g.write('  <Grid Name="Partition{}" GridType="Collection">\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))
                g.write('   <Grid Name="{}s[{}]" GridType="Uniform">\n'.format(xdmfTopologyType[firstOrderCellType[cellType]], part))
                writeTopology(g, xdmfTopologyType[firstOrderCellType[cellType]],
                              partitions[cellType][part], xfname)
                writeGeometry(g, ndims[firstOrderCellType[cellType]], 
                               partitions[cellType][part], 
                               nodeIDs[cellType], 
                               args.mesh, dname)
                writeAttribute(g, part)
                g.write('   </Grid>\n')
                # connectivities file
                writeConnectivities(xfname, partitions[cellType][part], 
                                    len(nodeIDs[cellType]))
        g.write('  </Grid>\n')
    
    writeFooter(g)
    g.close()
else:
    print("mesh does not contain any supported cell types")