from . import function, _, numeric
import numpy

def _tryall( obj, prefix, kwargs ):
  for name in dir( obj ):
    if name.startswith( prefix ):
      try:
        return getattr( obj, name )( **kwargs )
      except TypeError:
        pass
  raise Exception, 'not supported: ' + ', '.join( kwargs.keys() )

class Hooke:

  def __init__( self, **kwargs ):
    
    verify = kwargs.pop( 'verify', True )

    if len(kwargs)!=2:
      raise ValueError( 'exactly two arguments expected, found %d' % len(kwargs) )

    _tryall( self, '_set_from_', kwargs )

    if verify:
      for key, value in kwargs.items():
        numpy.testing.assert_almost_equal( value, getattr(self,key) )

  def _set_from_lame( self, lmbda, mu ):
    self.lmbda = float(lmbda)
    self.mu = float(mu)

  def _set_from_poisson_young( self, nu, E ):
    self.lmbda = (nu*E)/((1.+nu)*(1.-2.*nu))
    self.mu = E/(2.*(1.+nu))

  def __call__ ( self, disp, coords ): 
    return self.lmbda * disp.div(coords)[...,_,_] * function.eye( coords.shape[0] ) \
      + (2*self.mu) * disp.symgrad(coords)
    
  def __str__( self ):
    return 'Hooke(mu=%s,lmbda=%s)' % ( self.mu, self.lmbda )

  @property
  def E( self ):
    return self.mu * (3.*self.lmbda+2.*self.mu) / (self.lmbda+self.mu)

  @property
  def nu( self ):
    return self.lmbda / (2.*(self.lmbda+self.mu))

class Orthotropic(object):
  
  def __init__( self, E, G, nu ):

    vbar = numpy.product( nu )
    vv = nu[:,_]*nu[_,:]
    NE = -numpy.sqrt( E[:,_]*E[_,:] ) * ( vbar/vv + vv ) 
    NEdiag = numeric.takediag( NE ) 
    NEdiag += E*(2.*nu**2-1.+vbar/nu**2)
    NE /= (nu**2).sum() + 2*vbar - 1
    
    self.NE     = NE
    self.Gbar   = numpy.product( G )
    self.Ginv   = 1./G


  def __call__( self, disp, coords ):
    gradu = disp.grad( coords )
    diag_gradu = function.takediag( gradu )
    GG = self.Ginv[:,_]*self.Ginv[_,:]*self.Gbar
    return function.diagonalize( (self.NE * diag_gradu[...,_,:]).sum() - 2.*self.Gbar*self.Ginv**2*diag_gradu ) \
        + GG*( gradu+gradu.swapaxes(-2,-1) ) #transpose
      

from . import function, _, numeric
import numpy

def _tryall( obj, prefix, kwargs ):
  for name in dir( obj ):
    if name.startswith( prefix ):
      try:
        return getattr( obj, name )( **kwargs )
      except TypeError:
        pass
  raise Exception, 'not supported: ' + ', '.join( kwargs.keys() )

class Hooke:

  def __init__( self, **kwargs ):
    
    verify = kwargs.pop( 'verify', True )

    if len(kwargs)!=2:
      raise ValueError( 'exactly two arguments expected, found %d' % len(kwargs) )

    _tryall( self, '_set_from_', kwargs )

    if verify:
      for key, value in kwargs.items():
        numpy.testing.assert_almost_equal( value, getattr(self,key) )

  def _set_from_lame( self, lmbda, mu ):
    self.lmbda = float(lmbda)
    self.mu = float(mu)

  def _set_from_poisson_young( self, nu, E ):
    self.lmbda = (nu*E)/((1.+nu)*(1.-2.*nu))
    self.mu = E/(2.*(1.+nu))

  def __call__ ( self, disp, coords ): 
    return self.lmbda * disp.div(coords)[...,_,_] * function.eye( coords.shape[0] ) \
      + (2*self.mu) * disp.symgrad(coords)
    
  def __str__( self ):
    return 'Hooke(mu=%s,lmbda=%s)' % ( self.mu, self.lmbda )

  @property
  def E( self ):
    return self.mu * (3.*self.lmbda+2.*self.mu) / (self.lmbda+self.mu)

  @property
  def nu( self ):
    return self.lmbda / (2.*(self.lmbda+self.mu))

class Orthotropic(object):
  
  def __init__( self, E, G, nu ):

    vbar = function.product( nu )
    vv = nu[:,_]*nu[_,:]
    NE = -function.sqrt( E[:,_]*E[_,:] ) * ( vbar/vv + vv )  + function.diagonalize(E*(2.*nu**2-1.+vbar/nu**2))
    NE /= (nu**2).sum() + 2*vbar - 1
    
    self.NE     = NE
    self.Gbar   = function.product( G )
    self.Ginv   = 1./G


  def __call__( self, disp, coords ):
    gradu = disp.grad( coords )
    diag_gradu = function.takediag( gradu )
    GG = self.Ginv[:,_]*self.Ginv[_,:]*self.Gbar
    return function.diagonalize( (self.NE * diag_gradu[...,_,:]).sum() - 2.*self.Gbar*self.Ginv**2*diag_gradu ) \
        + GG*( gradu+gradu.swapaxes(-2,-1) )
      

use function instead of numpy