from odbAccess import *
from textRepr import *
from string import *

import os
import time


class VTKFile(object):
    def __init__(self, outfile, this_model, N_Frame):
        # compute point quantity
        self.this_model = this_model
        self.N_Frame = N_Frame

        self.outfile = open(outfile, 'w')
        self.before_Results_data()
        self.Results_data()
        self.after_Results_data()
        self.outfile.close()

    def before_Results_data(self):
        # compute the number of element of each block
        p_elements = len(self.this_model.element)/1 + 1
        lp_elements = len(self.this_model.element) - \
            (p_elements*(1-1))  # last block

        # <VTKFile>, including the type of mesh, version, and byte_order
        self.outfile.write(
            '<VTKFile type="UnstructuredGrid" version="0.1" byte_order="LittleEndian">'+'\n')
        # <UnstructuredGrid>
        self.outfile.write('<UnstructuredGrid>'+'\n')
        # <Piece>, including the number of points and cells
        self.outfile.write('<Piece NumberOfPoints="'+str(len(self.this_model.stg_n)) +
                           '"'+' '+'NumberOfCells="'+str(lp_elements)+'">'+'\n')

        print("Writing Nodes ......")
        # <Points> Write nodes into vtk files
        self.outfile.write('<Points>'+'\n')
        self.outfile.write(
            '<DataArray type="Float64" NumberOfComponents="3" format="ascii">'+'\n')
        for i in range(0, len(self.this_model.stg_n)):
            nt = self.this_model.stg_n[i]
            # base shape
            X, Y, Z = self.this_model.node[nt].coordinates[0], self.this_model.node[
                nt].coordinates[1], self.this_model.node[nt].coordinates[2]
            # modify shape
            # X,Y,Z = node[nt].coordinates[0]+ux,node[nt].coordinates[1]+uy,node[nt].coordinates[2]+uz
            self.outfile.write(' '+'%11.8e' % X+'  '+'%11.8e' %
                               Y+'  '+'%11.8e' % Z+'\n')
        self.outfile.write('</DataArray>'+'\n')
        self.outfile.write('</Points>'+'\n')
        # </Points>

    def Results_data(self):
        print("Writing Results data ......")
        # <PointData> Write results data into vtk files

        # 'U','A', 'V', 'RF','PEEQ','S'
        col = {"Scalars": [], "Vevtors": [], "Tensors": []}
        for var_id in ['NT11']:
            try:
                col["Scalars"].append("Temperature")
            except:
                print('jump', var_id)
        con = "Scalars="+'"'+','.join(col["Scalars"])+'"'
        # con = "Tensors="+'"'+','.join(col["Tensors"])+'"' + " "+ \
        #       "Vevtors="+'"'+','.join(col["Vevtors"])+'"' + " "+ \
        #       "Scalars="+'"'+','.join(col["Scalars"])+'"'
        self.outfile.write("<"+"PointData"+" "+con+">"+'\n')

        for var_id in ['NT11']:
            fieldOutputs = self.N_Frame.fieldOutputs[var_id]
            from utils_odb import NT11 as pq_class
            pq_class(fieldOutputs, self.outfile, self.this_model.stg_n)

        self.outfile.write("</PointData>"+'\n')
        # </PointData>

    def after_Results_data(self):
        print("Writing Cells ......")
        # <Cells> Write cells into vtk files
        self.outfile.write('<Cells>'+'\n')

        # Connectivity 8 node
        self.outfile.write(
            '<DataArray type="Int32" Name="connectivity" format="ascii">'+'\n')
        if (self.this_model.mesh_type == 12):
            for i in range(len(self.this_model.stg_e)/8):
                self.outfile.write(
                    ' '.join([str(self.this_model.stg_e[i*8+idx]) for idx in range(8)])+'\n')
        if (self.this_model.mesh_type == 10):
            for i in range(len(self.this_model.stg_e)/4):

                self.outfile.write(
                    ' '.join([str(self.this_model.stg_e[i*4+idx]) for idx in range(4)])+'\n')
        self.outfile.write('</DataArray>'+'\n')

        # Offsets
        self.outfile.write(
            '<DataArray type="Int32" Name="offsets" format="ascii">'+'\n')
        for i in range(len(self.this_model.stg_e)/self.this_model.mesh_conner):
            self.outfile.write(
                str(i*self.this_model.mesh_conner+self.this_model.mesh_conner)+'\n')
        self.outfile.write('</DataArray>'+'\n')

        # Mesh Type
        self.outfile.write(
            '<DataArray type="UInt8" Name="types" format="ascii">'+'\n')
        for i in range(len(self.this_model.stg_e)/self.this_model.mesh_conner):
            self.outfile.write(str(self.this_model.mesh_type)+'\n')
        self.outfile.write('</DataArray>'+'\n')
        self.outfile.write('</Cells>'+'\n')
        # </Cells>

        # </Piece>
        self.outfile.write('</Piece>'+'\n')
        # </UnstructuredGrid>
        self.outfile.write('</UnstructuredGrid>'+'\n')
        # </VTKFile>
        self.outfile.write('</VTKFile>'+'\n')


class MODEL(object):
    def __init__(self, this_instance, mesh_type):
        self.mesh_type = mesh_type
        self.mesh()

        # access nodes & elements
        self.this_instance = this_instance

        self.count()

    def mesh(self):
        if (self.mesh_type == 12):
            mesh_conner = 8
            mesh_name = "Hexahedron"
        elif (self.mesh_type == 10):
            mesh_conner = 4
            mesh_name = "Tetra"
        else:
            print("Mesh type error or unidentified")
            os._exit(0)

        print("Mesh type:", mesh_name)
        self.mesh_conner = mesh_conner

    def count(self):
        node = self.this_instance.nodes
        element = self.this_instance.elements

        # match nodes' label and its order in sequence (for empty nodes in tetra mesh)
        MLN = node[len(node)-1].label
        TOTAL = []
        # read node in sequence, and get the largest label of node(non-empty)
        # MLN is the max label of nodeset
        for i in node:
            TOTAL.append(i.label)
            if (i.label > MLN):
                MLN = i.label
        # match (the key)
        L = []
        n = 0
        for i in range(MLN):
            L.append(0)
        for i in TOTAL:
            L[i-1] = n
            n += 1

        # estimate whether the node has already existed
        stg_p = []
        # store the reorganized node for element
        stg_e = []
        # store the reorganized node for node
        stg_n = []
        for i in range(MLN):
            stg_p.append(-1)
        nodecount = 0
        # reorganize the node and element (reconstruct the mesh)

        for i in range(0, len(element)):
            for j in range(self.mesh_conner):
                k = element[i].connectivity[j] - 1
                if (stg_p[k] < 0):
                    stg_p[k] = nodecount
                    stg_n.append(int(L[k]))
                    stg_e.append(nodecount)
                    nodecount += 1
                else:
                    stg_e.append(stg_p[k])

        self.stg_p = stg_p
        self.stg_e = stg_e
        self.stg_n = stg_n

        self.element = element
        self.node = node


# Modify the default value of filename here to specify the default configuration file
def ConvertOdb2Vtk(odb_path):
    # open an ODB ( Abaqus output database )
    odb = openOdb(odb_path, readOnly=True)
    print("ODB opened")

    # instance cycle # access geometry and topology information ( odb->rootAssembly->instances->(nodes, elements) )
    instance = odb.rootAssembly.instances
    for instancename in instance.keys():
        print("Instance: ", instancename)
        this_model = MODEL(
            mesh_type=mesh_type, this_instance=instance[instancename])
        # step cycle
        step = odb.steps
        for stepname in step.keys():
            print("Step: ", stepname)
            # frame cycle # access attribute(fieldOutputs) information
            for i_frame, N_Frame in enumerate(step[stepname].frames):
                # if i_frame<len(step[stepname].frames)-1:
                # if i_frame > 20:
                #     continue
                # Access a frame
                print("Frame:", i_frame)

                time_temp = time.time()

                # create and open a VTK(.vtu) files
                VTKFile(outfile=os.path.join(vtk_path)+stepname +
                        '_'+instancename+'_f%04d' % int(i_frame)+'.vtu',
                        this_model=this_model,
                        N_Frame=N_Frame)

                print("Time elapsed: ", time.time() - time_temp, "s")
    odb.close()


# get odb file's path
odb_path = 'C:/Users/OPTX/Downloads/Digital-Twin/odb2vtk/Thermal.odb'
# get the output files' path
vtk_path = 'C:/Users/OPTX/Downloads/Digital-Twin/odb2vtk/output/'  # not.vtu,.vtk
mesh_type = 12

ConvertOdb2Vtk(odb_path)