import numpy as np
from numpy.core.einsumfunc import einsum
from numpy.core.fromnumeric import swapaxes

class ElasticEnergy:
    def __init__(self, young, poisson):
        '''
        Input:
        - young   : Young's modulus [Pa]
        - poisson : Poisson ratio
        '''
        self.young   = young
        self.poisson = poisson
        self.lbda = young * poisson / ((1 + poisson) * (1 - 2 * poisson))
        self.mu   = young / (2 * (1 + poisson))

        self.psi = None
        self.E   = None
        self.P   = None

        self.dE = None
        self.dP = None

    def make_energy_density(self, jac):
        '''
        This method computes the energy density at each tetrahedron (#t,),
        and stores the result in self.psi

        Input:
        - jac  : jacobian of the deformation (#t, 3, 3)
        
        Updated attributes:
        - psi : energy density per tet (#t,)
        '''

        print("Please specify the kind of elasticity model.")
        raise NotImplementedError
    
    def make_strain_tensor(self, jac):
        '''
        This method computes the strain tensor (#t, 3, 3), and stores it in self.E

        Input:
        - jac : jacobian of the deformation (#t, 3, 3)
        
        Updated attributes:
        - E : strain induced by the deformation (#t, 3, 3)
        '''

        print("Please specify the kind of elasticity model.")
        raise NotImplementedError

    def make_piola_kirchhoff_stress_tensor(self, jac):
        '''
        This method computes the stress tensor (#t, 3, 3), and stores it in self.P

        Input:
        - jac : jacobian of the deformation (#t, 3, 3)
        
        Updated attributes:
        - P : stress tensor induced by the deformation (#t, 3, 3)
        '''

        print("Please specify the kind of elasticity model.")
        raise NotImplementedError

    def make_differential_strain_tensor(self, jac, dJac):
        '''
        This method computes the differential of strain tensor (#t, 3, 3), 
        and stores it in self.dE

        Input:
        - jac  : jacobian of the deformation (#t, 3, 3)
        - dJac : differential of the jacobian of the deformation (#t, 3, 3)
        
        Updated attributes:
        - dE : differential of the strain tensor (#t, 3, 3)
        '''

        print("Please specify the kind of elasticity model.")
        raise NotImplementedError

    def make_differential_piola_kirchhoff_stress_tensor(self, jac, dJac):
        '''
        This method computes the differential of the stress tensor (#t, 3, 3), and stores it in self.dP

        Input:
        - jac  : jacobian of the deformation (#t, 3, 3)
        - dJac : differential of the jacobian of the deformation (#t, 3, 3)
        
        Updated attributes:
        - dP : differential of the stress tensor (#t, 3, 3)
        '''

        print("Please specify the kind of elasticity model.")
        raise NotImplementedError

class LinearElasticEnergy(ElasticEnergy):
    def __init__(self, young, poisson):
        super().__init__(young, poisson)

    def make_energy_density(self, jac):
        fa = self.mu * np.einsum("ijk,ijk->i",self.E, self.E)
        fb = self.lbda/2.0 * np.einsum('ijj->i' ,self.E)**2
        self.psi = fa + fb 

    def make_strain_tensor(self, jac):
        self.E = (jac+np.transpose(jac,axes=(0,2,1)))/2.0 - np.identity(3)

    def make_piola_kirchhoff_stress_tensor(self, jac):
        fa = 2.0*self.mu*self.E 
        fb = self.lbda* np.einsum('i,jk->ijk',np.einsum('ijj->i',self.E), np.identity(3))
        self.P = fa + fb

    def make_differential_strain_tensor(self, jac, dJac):
        self.dE = (dJac+np.transpose(dJac,axes=(0,2,1)))/2.0

    def make_differential_piola_kirchhoff_stress_tensor(self, jac, dJac):
        fa = 2.0*self.mu*self.dE 
        fb = self.lbda* np.einsum("i,jk->ijk",np.trace(self.dE, axis1=1,axis2=2), np.identity(3))
        self.dP = fa + fb

class NeoHookeanElasticEnergy(ElasticEnergy):
    def __init__(self, young, poisson):
        super().__init__(young, poisson)
        self.logJ = None
        self.Finv = None

    def make_energy_density(self, jac):
        i1 = np.einsum('ijk,ijk->i',jac,jac)
        j = np.log(np.linalg.det(jac))
        fa = self.mu/2.0 * (i1-3 - 2*j)
        fb = self.lbda/2.0 * j**2
        self.psi = fa + fb

    def make_strain_tensor(self, jac):
         pass

    def make_piola_kirchhoff_stress_tensor(self, jac):
        '''
        Additional updated attributes:
        - logJ ; log of the determinant of the jacobians (#t,)
        - Finv : inverse of the jacobians (#t, 3, 3)
        '''
        self.logJ = np.log(np.linalg.det(jac))
        self.Finv = np.linalg.inv(jac)
        FinvT = np.transpose(self.Finv,axes=(0,2,1))
        fa = self.mu * (jac - FinvT)
        fb = self.lbda * np.einsum("i,ijk->ijk",self.logJ, FinvT)
        self.P = fa +fb

    def make_differential_strain_tensor(self, jac, dJac):
        pass

    def make_differential_piola_kirchhoff_stress_tensor(self, jac, dJac):
        self.logJ = np.log(np.linalg.det(jac))
        self.Finv = np.linalg.inv(jac)
        FinvT = np.transpose(self.Finv,axes=(0,2,1))
        dFT = np.transpose(dJac,axes=(0,2,1))
        fa = self.mu * dJac
        fb = np.einsum("i,ijk->ijk",(self.mu - self.lbda * self.logJ),FinvT @ dFT @ FinvT)
        fc = self.lbda * np.einsum("i,ijk->ijk",np.trace(self.Finv @ dJac, axis1=1,axis2=2),FinvT)
        self.dP = fa + fb + fc