class Physical_Quantity(object):
    def __init__(self):
        pass
        # Scalars NT11,RFL11
        # Vectors displacement,HFL
        # Tensors Stress,spstress

    def pq_update(self, var_id):
        if var_id == 'S':
            self.pq_class = Stress()
        if var_id == 'SPS':
            self.pq_class = Spstress()
        if var_id == 'DISP':
            self.pq_class = Displacement()
        if var_id == 'NT11':
            self.pq_class = Temperature()
        self.var_type = self.pq_class.var_type
        self.NumberOfComponents = self.pq_class.NumberOfComponents

    def insert(self, var_data):
        return self.pq_class.insertData(var_data)

    def write(self, *args):
        # print(args)
        self.pq_class.writeData(self.nids, self.eids, args)


class Stress:
    def __init__(self):
        self.var_type = 'Tensors'
        self.NumberOfComponents = 9

    def insertData(self, var_data):
        this_pq = {}
        for var_line in var_data:
            eid = var_line.elementLabel
            sxx = var_line.data[0]
            syy = var_line.data[1]
            szz = var_line.data[2]
            sxy = var_line.data[3]
            sxz = var_line.data[4]
            syz = var_line.data[5]
            this_pq[int(eid)] = [float(sxx), float(syy), float(
                szz), float(sxy), float(sxz), float(syz)]
        return this_pq

    def writeData(self, stress, ofs):
        for key in self.eids:
            s11, s22, s33, s12, s13, s23 = stress[key][0:]
            ofs.write("%f %f %f %f %f %f %f %f %f " %
                      (s11, s12, s13, s12, s22, s23, s13, s23, s33))


class Spstress:
    def __init__(self):
        self.var_type = 'Vectors'
        self.NumberOfComponents = 3

    def insertData(self, var_data):
        this_pq = {}
        for var_line in var_data:
            eid = var_line.elementLabel
            smin = var_line.maxPrincipal
            smin = var_line.midPrincipal
            smin = var_line.minPrincipal
            this_pq[int(eid)] = [float(smin), float(smin), float(smin)]
        return this_pq
    
    def writeData(self, nids, eids, args):
        spstress, ofs = args
        for key in eids:
            sp1, sp2, sp3 = spstress[key][0:]
            ofs.write("%f %f %f " % (sp1, sp2, sp3))


class Displacement:
    def __init__(self):
        self.var_type = 'Vectors'
        self.NumberOfComponents = 3

    def insertData(self, var_data):
        this_pq = {}
        for var_line in var_data:
            nid = var_line.nodeLabel
            ux = var_line.data[0]
            uy = var_line.data[1]
            uz = var_line.data[2]
            this_pq[int(nid)] = [float(ux),  float(uy),  float(uz)]
        return this_pq

    def writeData(self, nids, eids, args):
        displacement, ofs = args
        for key in nids:
            ux, uy, uz = displacement[key]
            ofs.write("%f %f %f " % (ux, uy, uz))


class Temperature(object):
    def __init__(self):
        self.var_type = 'Scalars'
        self.NumberOfComponents = 1

    def insertData(self, var_data):
        this_pq = {}
        for var_line in var_data:
            nid = var_line.nodeLabel
            temp = var_line.data
            this_pq[int(nid)] = [float(temp)]
        return this_pq

    def writeData(self, nids, eids, args):
        temperature, ofs = args
        for key in nids:
            temp = temperature[key][0]
            # print(temp)
            ofs.write("%f " % (temp,))


# ### define PointsData element ###
#             ofs.write('<PointData Vectors="displacement">')
#             ofs.write('<DataArray type="Float32" NumberOfComponents="3" format="ascii" Name="displacement">')
#             for key in self.nids:
#                 ux, uy, uz = self.displacement[key]
#                 ofs.write("%f %f %f " % (ux, uy, uz))
#             ofs.write('</DataArray>')
#             ofs.write('</PointData>')

# ### define CellData element ###
#             ofs.write('<CellData Vectors="PrincipalStress" Tensors="Stress">')
#             ofs.write('<DataArray type="Float32" Name="PrincipalStress" NumberOfComponents="3" format="ascii">')
#             for key in self.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 self.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()