import numpy
#
#
#  Gdyr_Principal_Strains
#    Rev 0
#    Aug 27, 2015
#  Dennis Conklin - Engineering Mechanics
#      Modularized using Utility routines
#  
#  Paraview 4.3.1 Progammable Filter
#  Adds Cell Variables:
#      Principal Strains:         P1_strain, P2_strain, P3_strain
#      Principal Strain Cosines:  P1_vector, P2_vector, P3_vector
#      max shear strain:          tauMax
#
#   
def check_USTRTOT_loaded(block):
   if (check_elem_variable_exists(block,'USTRTOTXX') and
       check_elem_variable_exists(block,'USTRTOTYY') and
       check_elem_variable_exists(block,'USTRTOTZZ') and
       check_elem_variable_exists(block,'USTRTOTXY') and
       check_elem_variable_exists(block,'USTRTOTZX') and
       check_elem_variable_exists(block,'USTRTOTYZ')     ):
       return 1 # true - all loaded
   else:
       return 0 # false - not all loaded
#
#
def append_elem_variable(block,variable,variable_name):
   # append only if variable doesn't already exist
   # save time if calc in multiple filters
   # running filter multiple times will not change things
   #
   if not check_elem_variable_exists(block,variable_name):
      block.CellData.append(variable,variable_name)
#   
def check_elem_variable_exists(block,variable_name):
   keys=block.CellData.keys()
   # true if exists
   return (variable_name in keys)
#
#
#
def process_block(block):
   # 
   #
   # First check if USTRTOT loaded in Paraview
   if not check_USTRTOT_loaded(block):
      print 'Element tensor USTRTOT read failed, please load all 6 USTRTOT'
      return
   xx=block.CellData['USTRTOTXX']
   yy=block.CellData['USTRTOTYY']
   zz=block.CellData['USTRTOTZZ']
   xy=block.CellData['USTRTOTXY']
   xz=block.CellData['USTRTOTZX']
   yz=block.CellData['USTRTOTYZ']

   numElems = block.GetNumberOfCells()
     
   # declare arrays before using them
   # create output arrays with correct dimensions
   P1=zeros( (numElems,1) )
   P2=zeros( (numElems,1) )
   P3=zeros( (numElems,1) )
   #create vector arrays with correct dimensions
   P1_vector=zeros( (numElems,3) )
   P2_vector=zeros( (numElems,3) )
   P3_vector=zeros( (numElems,3) )
   
   max_shear=zeros( (numElems,1) )
   iso=numpy.zeros( (3,3) )
   dev=numpy.zeros( (3,3) )
    
   for elem in range(numElems):
      sigma = numpy.array([ [xx[elem],xy[elem],xz[elem]],
                            [xy[elem],yy[elem],yz[elem]],
                            [xz[elem],yz[elem],zz[elem]]
                         ])
      
      # isotropic strain matrix(3,3)
      iso = 1.0/3.0*numpy.trace(sigma)*numpy.eye(3)
      # deviatoric strain matrix(3,3)
      dev = sigma - iso
      
      #principal strains
      eigvals, eigvectors=numpy.linalg.eigh(sigma)
      #eigvals=list(eigvals)
      #print eigvals
      #print eigvectors
      for index in range(3):
         if eigvals[index]==max(eigvals):
            P1_index=index
         elif eigvals[index]==min(eigvals):
            P3_index=index
         else:
            P2_index=index
      
      P1[elem]=eigvals[P1_index]
      P2[elem]=eigvals[P2_index]
      P3[elem]=eigvals[P3_index]
      # max shear
      max_shear[elem]=(max(eigvals)-min(eigvals))/2.0 
      P1_vector[elem]=eigvectors[:,P1_index]
      P2_vector[elem]=eigvectors[:,P2_index]
      P3_vector[elem]=eigvectors[:,P3_index]

      
   # Add  variables to Elements  
   append_elem_variable(block,P1,"P1_strain")
   append_elem_variable(block,P2,"P2_strain")
   append_elem_variable(block,P3,"P3_strain")
   append_elem_variable(block,max_shear,"tauMax")  
   append_elem_variable(block,P1_vector,"P1_vector")
   append_elem_variable(block,P2_vector,"P2_vector")
   append_elem_variable(block,P3_vector,"P3_vector")

# Loop over blocks in composite (Exodus) data set
for block in output:

   # only process HEX elements
   # VTK element types:
   #     12 = VTK_HEXAHEDRON (Gdyr Hex)
   #      9 = VTK_QUAD  (Gdyr Membrane)
   #      3 = VTK_LINE  (Gdyr 3D Truss)
   # since EXODUS requires single element type per block, test 1st element only
   if block.GetCellType(0)==12:
      # process each hex block
      process_block(block)