import numpy as np
from numpy import linalg
from scipy import sparse
import igl


# -----------------------------------------------------------------------------
# -----------------------------------------------------------------------------
#                             ELASTIC SOLID CLASS
# -----------------------------------------------------------------------------
# -----------------------------------------------------------------------------


class ElasticSolid(object):

    def __init__(self, v_rest, t, rho=1, pin_idx=[]):
        '''
        Input:
        - v_rest      : position of the vertices of the mesh (#v, 3)
        - t           : indices of the element's vertices (#t, 4)
        - rho         : mass per unit volume [kg.m-3]
        - pin_idx     : list of vertex indices to pin
        '''

        self.v_rest   = v_rest.copy()
        self.v_def    = v_rest.copy()
        self.t        = t
        self.rho      = rho
        self.pin_idx  = pin_idx
        self.free_idx = None
        self.pin_mask = None
        
        self.W0 = None
        self.Dm = None
        self.Bm = None
        self.rest_barycenters = None

        self.W  = None
        self.Ds = None
        self.F  = None
        self.def_barycenters = None

        self.make_free_indices_and_pin_mask()
        self.update_rest_shape(self.v_rest)
        self.update_def_shape(self.v_def)

    ## Precomputation ##

    def make_free_indices_and_pin_mask(self):
        '''
        Should list all the free indices and the pin mask.

        Updated attributes:
        - free_index : np array of shape (#free_vertices,) containing the list of unpinned vertices
        - pin_mask   : np array of shape (#v, 1) containing 1 at free vertex indices and 0 at pinned vertex indices
        '''
        self.pin_mask = np.ones((self.v_rest.shape[0],1))
        self.pin_mask[self.pin_idx] = 0
        self.free_idx = np.where(self.pin_mask==1)[0]

    ## Methods related to rest quantities ##

    def make_rest_barycenters(self):
        '''
        Construct the barycenters of the undeformed configuration

        Updated attributes:
        - rest_barycenters : np array of shape (#t, 3) containing the position of each tet's barycenter
        '''
        self.rest_barycenters = np.einsum("ijk->ik",self.v_rest[self.t[:,:]])/4

    def make_rest_shape_matrices(self):
        '''
        Construct Dm that has shape (#t, 3, 3), and its inverse Bm

        Updated attributes:
        - Dm : np array of shape (#t, 3, 3) containing the shape matrix of each tet
        - Bm : np array of shape (#t, 3, 3) containing the inverse shape matrix of each tet
        '''
        tv = self.v_rest[self.t[:,:]]

        tva = tv[:,:3,:]
        tvb = tv[:,3,:].reshape((self.t.shape[0],1,3))
        tvb= np.repeat(tvb, 3, axis=1)
        
        tvs = tva - tvb
        tvsf = np.transpose(tvs,(0,2,1))
        
        self.Dm = tvsf
        self.Bm = np.linalg.inv(tvsf)

    def update_rest_shape(self, v_rest):
        '''
        Updates the vertex position, the shape matrices Dm and Bm, the volumes W0,
        and the mass matrix at rest

        Input:
        - v_rest : position of the vertices of the mesh at rest state (#v, 3)

        Updated attributes:
        - v_rest : np array of shape (#v, 3) containing the position of each vertex at rest
        - W0     : np array of shape (#t) containing the volume of each tet
        '''
        self.v_rest = v_rest
        self.make_rest_barycenters()
        self.make_rest_shape_matrices()
        self.W0 = -np.linalg.det(self.Dm)/6
        
        self.update_def_shape(self.v_def)

    ## Methods related to deformed quantities ##

    def make_def_barycenters(self):
        '''
        Construct the barycenters of the deformed configuration

        Updated attributes:
        - def_barycenters : np array of shape (#t, 3) containing the position of each tet's barycenter
        '''
        self.def_barycenters = np.einsum("ijk->ik",self.v_def[self.t[:,:]])/4

    def make_def_shape_matrices(self):
        '''
        Construct Ds that has shape (#t, 3, 3)

        Updated attributes:
        - Ds : np array of shape (#t, 3, 3) containing the shape matrix of each tet
        '''
        tv = self.v_def[self.t[:,:]]

        tva = tv[:,:3,:]
        tvb = tv[:,3,:].reshape((self.t.shape[0],1,3))
        tvb= np.repeat(tvb, 3, axis=1)
        
        tvs = tva - tvb
        tvsf = np.transpose(tvs,(0,2,1))
        
        self.Ds = tvsf

    def make_jacobians(self):
        '''
        Compute the current Jacobian of the deformation

        Updated attributes:
        - F : np array of shape (#t, 3, 3) containing Jacobian of the deformation in each tet
        '''
        self.F = self.Ds @ self.Bm

    def update_def_shape(self, v_def):
        '''
        Updates the vertex position, the Jacobian of the deformation, and the 
        resulting elastic forces.

        Input:
        - v_def : position of the vertices of the mesh (#v, 3)

        Updated attributes:
        - v_def : np array of shape (#v, 3) containing the position of each vertex after deforming the solid
        - W     : np array of shape (#t, 3) containing the volume of each tet
        '''
        self.v_def = self.v_rest
        self.v_def[self.free_idx] = v_def[self.free_idx]
        self.make_def_barycenters()
        self.make_def_shape_matrices()
        self.W = -np.linalg.det(self.Ds)/6
        self.make_jacobians()
        
    def displace(self, v_disp):
        '''
        Displace the whole mesh so that v_def += v_disp

        Input:
        - v_disp : displacement of the vertices of the mesh (#v, 3)
        '''
        self.update_def_shape(self.v_def + v_disp)
