from odbAccess import openOdb
from abaqus import *
from abaqusConstants import *
from utils_odb2vtk import Physical_Quantity


class CAE_Model(Physical_Quantity, object):
    def __init__(self):
        super(CAE_Model, self).__init__()
        # print('self.var_type', self.var_type)
        self.nodes = {}
        self.elements = {}
        self.physical_quantity = {}

    def read_inp(self, INPPATH):
        # 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:
                        nid, x, y, z = line.strip().split(',')
                        # insertNode(uid, x, y, z)
                        self.nodes[int(nid)] = [float(x), float(y), float(z)]
                    elif state == 2:
                        eid, n0, n1, n2, n3, n4, n5, n6, n7 = line.strip().split(',')
                        # insertHexa(el, 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)]
                    else:
                        pass

    def read_odb(self, ODBPATH):
        # 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))

        for var_id in ['U', 'A', 'V', 'RF', 'S', 'LE', 'PE', 'PEEQ', 'NT11', 'HFL', 'RFL11']:
            # try:
            if var_id == 'NT11':
                # S=myodb.steps[stepName].frames[-1].fieldOutputs['S'].getSubset(position=INTEGRATION_POINT)
                var_data = myodb.steps[stepName].frames[-1].fieldOutputs[var_id].values
                print(var_id+": %d" % (len(var_data)))

                self.pq_update(var_id)
                self.physical_quantity[var_id] = self.insert(var_data)
            # except:
            #     print('jump ', var_id)

    def write_vtk(self, OUTPATH, var_ids):
        self.nids = sorted(self.nodes.keys())
        self.eids = sorted(self.elements.keys())
        with open(OUTPATH, '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 self.nids:
                x, y, z = self.nodes[key]
                ofs.write("%f %f %f " % (x, y, z))
            ofs.write('</DataArray>')
            ofs.write('</Points>')

            ### define Cells element ###
            ofs.write('<Cells>')
            ofs.write(
                '<DataArray type="Int64" Name="connectivity" format="ascii">')
            for key in self.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 self.eids:
                ofs.write("%d " % (key * 8))
            ofs.write('</DataArray>')
            ofs.write('<DataArray type="UInt8" Name="types" format="ascii">')
            for key in self.eids:
                ofs.write("10 ")
            ofs.write('</DataArray>')
            ofs.write('</Cells>')

            # node data/element data
            for var_id in var_ids:
                self.pq_update(var_id)

                # var_name = var_stress.__name__
                var_name = var_id
                var_type = self.var_type
                NumberOfComponents = str(self.NumberOfComponents)

                ofs.write('<'+var_type+'Data '+var_type+'="' + var_name + '">')
                ofs.write('<DataArray type="Float32" format="ascii" Name="' +
                          var_name+'" NumberOfComponents="'+NumberOfComponents+'" >')
                self.write(self.physical_quantity[var_id], ofs)
                ofs.write('</DataArray>')
                ofs.write('</'+var_type+'Data>')

            # finshed
            ofs.write('</Piece>')
            ofs.write('</StructuredGrid>')
            ofs.write('</VTKFile>')
            ofs.close()


INPPATH = 'C:/Users/Harvo/Downloads/Digital-Twin/thermal/Thermal.inp'
ODBPATH = 'C:/Users/Harvo/Downloads/Digital-Twin/thermal/Thermal.odb'
OUTPATH = 'C:/Users/Harvo/Downloads/Digital-Twin/thermal/Thermal.vtk'

my_model = CAE_Model()
my_model.read_inp(INPPATH)
print('finish read_inp')
my_model.read_odb(ODBPATH)
print('finish read_odb')
my_model.write_vtk(OUTPATH, var_ids=['NT11'])
print('finish write_vtk')