INPPATH = 'C:/abq2020/test1204MC.inp'
ODBPATH = 'C:/abq2020/test1204MC.odb'
OUTPATH = 'C:/abq2020/test1204MC.vtk'

from odbAccess import openOdb
from abaqus import *
from abaqusConstants import *

class DataSet:
    def __init__(self):
        self.nodes = {}
        self.elements = {}
        self.stress = {}
        self.spstress = {}
        self.displacement = {}

    def insertNode(self, nid, x, y, z):
        self.nodes[int(nid)] = [float(x), float(y), float(z)]

    def insertHexa(self, eid, n0, n1, n2, n3, n4, n5, n6, n7):
        self.elements[int(eid)] = [int(n0), int(n1), int(n2), int(n3), int(n4), int(n5), int(n6), int(n7)]
        
    def insertStress(self, eid, s11, s22, s33, s12, s13, s23):
        self.stress[int(eid)] = [float(s11), float(s22), float(s33), float(s12), float(s13), float(s23)]
              
    def insertSpstress(self, eid, sp1, sp2, sp3):
        self.spstress[int(eid)] = [float(sp1), float(sp2), float(sp3)]

    def insertDisplacement(self, nid, ux, uy, uz):
        self.displacement[int(nid)] = [float(ux),  float(uy),  float(uz)]

    def save(self, file_name):
        nids = sorted(self.nodes.keys())
        eids = sorted(self.elements.keys())
        with open(file_name, 'w') as ofs:
            ofs.write('<?xml version="1.0"?>')
            ofs.write('<VTKFile type="StructuredGrid" version="0.1" byte_order="LittleEndian" header_type="UInt32">')
            ofs.write('<StructuredGrid> WholeExtent="x1 x2 x3 y1 y2 y3 z1 z2 z3">')
            ofs.write('<Piece Extent="x1 x2 x3 y1 y2 y3 z1 z2 z3">')
            
### define Points element ###
            ofs.write('<Points>')
            ofs.write('<DataArray type="Float32" Name="Points" NumberOfComponents="3" format="ascii">')
            for key in nids:
                x, y, z = self.nodes[key]
                ofs.write("%f %f %f " % (x, y, z))                        
            ofs.write('</DataArray>')
            ofs.write('</Points>')  
                      
### define PointsData element ###            
            ofs.write('<PointData Vectors="displacement">')
            ofs.write('<DataArray type="Float32" NumberOfComponents="3" format="ascii" Name="displacement">')
            for key in nids:
                ux, uy, uz = self.displacement[key]
                ofs.write("%f %f %f " % (ux, uy, uz))
            ofs.write('</DataArray>')
            ofs.write('</PointData>')
            
### define Cells element ###             
            ofs.write('<Cells>')
            ofs.write('<DataArray type="Int64" Name="connectivity" format="ascii">')
            for key in eids:
                ns = self.elements[key]
                nns = map(lambda x: str(x - 1), ns)
                cons = " ".join(nns)
                ofs.write(cons + ' ')
            ofs.write('</DataArray>')
            ofs.write('<DataArray type="Int64" Name="offsets" format="ascii">')
            for key in eids:
                ofs.write("%d " % (key * 8))
            ofs.write('</DataArray>')
            ofs.write('<DataArray type="UInt8" Name="types" format="ascii">')
            for key in eids:
                ofs.write("10 ")
            ofs.write('</DataArray>')
            ofs.write('</Cells>')
            
### define CellData element ###              
            ofs.write('<CellData Vectors="PrincipalStress" Tensors="Stress">')                
            ofs.write('<DataArray type="Float32" Name="PrincipalStress" NumberOfComponents="3" format="ascii">')
            for key in eids:
                sp1, sp2, sp3 = self.spstress[key][0:]
                ofs.write("%f %f %f " % (sp1, sp2, sp3))      
            ofs.write('</DataArray>')
            ofs.write('<DataArray type="Float32" Name="Stress" NumberOfComponents="9" format="ascii">')
            for key in eids:
                s11, s22, s33, s12, s13, s23 = self.stress[key][0:]
                ofs.write("%f %f %f %f %f %f %f %f %f " % (s11, s12, s13, s12, s22, s23, s13, s23, s33))
            ofs.write('</DataArray>')
            ofs.write('</CellData>')
            ofs.write('</Piece>')
            ofs.write('</StructuredGrid>')
            ofs.write('</VTKFile>')
            ofs.close()


data = DataSet()
# 0 - nothing
# 1 - node
# 2 - hexahedron
state = 0
# read input file
with open(INPPATH, 'r') as input_file:
    for line in input_file:
        if line.strip() == '*Node':
            state = 1
        elif line.startswith('*Element, type=C3D8R'):
            state = 2
        else:
            if line.startswith("*") and (not line.startswith("**")):
                state = 0
            elif state == 1:
                uid, x, y, z = line.strip().split(',')
                data.insertNode(uid, x, y, z)
            elif state == 2:
                el, n0, n1, n2, n3, n4, n5, n6, n7 = line.strip().split(',')
                data.insertHexa(el, n0, n1, n2, n3, n4, n5, n6, n7)
            else:
                pass

# Open the odb
myodb = openOdb(ODBPATH)

# Get the frame repository for the step, find number of frames (starts at frame 0)
stepName = myodb.steps.keys()[0]
frames = myodb.steps[stepName].frames
numFrames = len(frames)
print("num Frame %d" % (numFrames))

# Isolate the instance, get the number of nodes and elements
instanceName = myodb.rootAssembly.instances.keys()[0]
myInstance = myodb.rootAssembly.instances[instanceName]
numNodes = len(myInstance.nodes)
numElements = len(myInstance.elements)
print("num Element %d, num Node %d" % (numElements, numNodes))

S=myodb.steps[stepName].frames[-1].fieldOutputs['S'].getSubset(position=INTEGRATION_POINT)
SPS=myodb.steps[stepName].frames[-1].fieldOutputs['S'].getSubset(position=INTEGRATION_POINT)
DISP=myodb.steps[stepName].frames[-1].fieldOutputs['U']
print("S: %d DISP: %d" % (len(S.values), len(DISP.values)))

for sts in S.values:
    eid = sts.elementLabel
    sxx = sts.data[0]
    syy = sts.data[1]
    szz = sts.data[2]
    sxy = sts.data[3]
    sxz = sts.data[4]
    syz = sts.data[5]
    data.insertStress(eid, float(sxx), float(syy), float(szz), float(sxy), float(sxz), float(syz))

for spsts in SPS.values:
    eid = spsts.elementLabel  
    smin = spsts.maxPrincipal
    smid = spsts.midPrincipal
    smax = spsts.minPrincipal
    data.insertSpstress(eid, float(smin), float(smid), float(smax))
    
for disp in DISP.values:
    nid = disp.nodeLabel
    ux = disp.data[0]
    uy = disp.data[1]
    uz = disp.data[2]
    data.insertDisplacement(nid, ux, uy, uz)

data.save(OUTPATH)