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


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


class ElasticSolid(object):

    def __init__(self, v_rest, t, ee, rho=1, pin_idx=[], f_mass=None):
        '''
        Input:
        - v_rest      : position of the vertices of the mesh (#v, 3)
        - t           : indices of the element's vertices (#t, 4)
        - ee          : elastic energy object that can be found in elasticenergy.py
        - rho         : mass per unit volume [kg.m-3]
        - pin_idx     : list or np array of vertex indices to pin
        - f_mass      : external force per unit mass (3,) [N.kg-1]
        '''

        self.v_rest   = v_rest.copy()
        self.v_def    = v_rest.copy()
        self.t        = t
        self.ee       = ee
        self.rho      = rho
        self.pin_idx  = np.array(pin_idx)
        self.f_mass   = f_mass.copy()
        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.energy_el  = None
        self.energy_ext = None

        self.f     = None
        self.f_vol = None
        self.f_ext = None

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

    ## Utils ##

    def vertex_tet_sum(self, data):
        '''
        Distributes data specified at each tetrahedron to the neighboring vertices.
        All neighboring vertices will receive the value indicated at the corresponding tet position in data.

        Input:
        - data : np array of shape (#t,) or (4*#t,)

        Output:
        - data_sum : np array of shape (#v,), containing the summed data
        '''
        i = self.t.flatten('F')        # (4*#t,)
        j = np.arange(len(self.t))     # (#t,)
        j = np.tile(j, 4)              # (4*#t,)

        if len(data) == len(self.t):
            data = data[j]

        # Has shape (#v, #t)
        m = sparse.coo_matrix((data, (i, j)), (len(self.v_rest), len(self.t)))
        return np.array(m.sum(axis=1)).flatten()

    ## 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.free_idx = np.arange(0,self.v_rest.shape[0])
        if(self.pin_idx.shape[0]<=0):
            return
        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 Ds 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 signed 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)
        self.make_volumetric_and_external_forces()

    ## 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 # <=> np.einsum("lij,ljk->lik",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,) containing the signed 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()
        if(self.f_vol is None):
            self.make_volumetric_and_external_forces()
        self.make_elastic_forces()
        self.make_elastic_energy()
        self.make_external_energy()
        
    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)

    ## Energies ##

    def make_elastic_energy(self):
        '''
        This updates the elastic energy

        Updated attributes:
        - energy_el  : elastic energy of the system [J]
        '''
        self.ee.make_strain_tensor(self.F)
        self.ee.make_energy_density(self.F)
        self.energy_el = np.sum(self.W0 * self.ee.psi)

    def make_external_energy(self):
        '''
        This computes the external energy potential

        Updated attributes:
        - energy_ext : postential energy due to external forces [J]
        '''
        dx = self.def_barycenters - self.rest_barycenters
        ifv = np.einsum('i,ij -> ij', self.W0, self.f_vol)
        self.energy_ext =  - np.einsum("i,ij,ij->",self.W0, self.f_vol, dx)

    ## Forces ##

    def make_elastic_forces(self):
        '''
        This method updates the elastic forces stored in self.f (#v, 3)

        Updated attributes:
        - f  : elastic forces per vertex (#v, 3)
        - ee : elastic energy, some attributes should be updated
        '''
        self.ee.make_strain_tensor(self.F)
        self.ee.make_piola_kirchhoff_stress_tensor(self.F)
        tdm = np.transpose(self.Bm,(0,2,1))
        ptdm = self.ee.P @ tdm
        f123 = -np.einsum("i,ijk->ikj",self.W0,ptdm )
        f4 = -f123[:,0] - f123[:,1] - f123[:,2] 
        H = np.hstack((f123,f4[:,None]))

        x = self.vertex_tet_sum(H[:,:,0].flatten('F'))
        y = self.vertex_tet_sum(H[:,:,1].flatten('F'))
        z = self.vertex_tet_sum(H[:,:,2].flatten('F'))

        self.f =  np.hstack((x[:,None],y[:,None],z[:,None]))

    def make_volumetric_and_external_forces(self):
        '''
        Convert force per unit mass to volumetric forces, then distribute
        the forces to the vertices of the mesh.

        Updated attributes:
        - f_vol : np array of shape (#t, 3) external force per unit volume acting on the tets
        - f_ext : np array of shape (#v, 3) external force acting on the vertices
        '''
        mpv = self.f_mass*self.rho
        self.f_vol = np.ones((self.W0.shape[0],1)) * mpv

        ifv = np.einsum('i,ij -> ij', self.W0, self.f_vol)

        f_ext_x = self.vertex_tet_sum(ifv[:,0])
        f_ext_y = self.vertex_tet_sum(ifv[:,1])
        f_ext_z = self.vertex_tet_sum(ifv[:,2])

        self.f_ext = np.stack((f_ext_x,f_ext_y,f_ext_z),axis=1) / 4.0

    ## Force Differentials

    def compute_force_differentials(self, v_disp):
        '''
        This computes the differential of the force given a displacement dx,
        where df = df/dx|x . dx = - K(x).dx. Where K(x) is the stiffness matrix (or Hessian)
        of the solid. Note that the implementation doesn't need to construct the stiffness matrix explicitly.

        Input:
        - v_disp : displacement of the vertices of the mesh (#v, 3)

        Output:
        - df : force differentials at the vertices of the mesh (#v, 3)

        Updated attributes:
        - ee : elastic energy, some attributes should be updated
        '''

        # First compute the differential of the Jacobian
        tv = v_disp[self.t]
        tva = tv[:,:3,:]
        tvb = tv[:,3,:].reshape((-1,1,3))

        tvb= np.repeat(tvb, 3, axis=1)
        
        tvs = tva - tvb
        tvsf = np.transpose(tvs,(0,2,1))
        
        dDs = tvsf
        dF = dDs @ self.Bm
        
        # Then update differential quantities in self.ee
        self.ee.make_differential_strain_tensor(self.F,dF)
        self.ee.make_differential_piola_kirchhoff_stress_tensor(self.F,dF)

        # Compute the differential of the forces
        tdm = np.transpose(self.Bm,(0,2,1))
        ptdm = self.ee.dP @ tdm
        f123 = -np.einsum("i,ijk->ikj",self.W0,ptdm)
        f4 = -f123[:,0] - f123[:,1] - f123[:,2] 
        H = np.hstack((f123,f4[:,None]))

        x = self.vertex_tet_sum(H[:,:,0].flatten('F'))
        y = self.vertex_tet_sum(H[:,:,1].flatten('F'))
        z = self.vertex_tet_sum(H[:,:,2].flatten('F'))

        return np.hstack((x[:,None],y[:,None],z[:,None]))
    
    def equilibrium_step(self, step_size_init = 2, max_l_iter = 20, c1 = 1e-4):
        '''
        This function displaces the whole solid to the next deformed configuration
        using a Newton-CG step.

        Updated attributes:
        - LHS : The hessian vector product
        - RHS : Right hand side for the conjugate gradient linear solve
        Other than them, only attributes updated by displace(self, v_disp) should be changed
        '''

        # Define LHS
        def LHS(dx):
            '''
            Should implement the Hessian-Vector Product L(dx), and take care of pinning constraints
            as described in the handout.
            '''
            if len(dx.shape)==1:
                dx = dx.reshape((-1,3))

            if (dx.shape[0] < self.v_rest.shape[0]):
                dxf = np.zeros((self.v_rest.shape[0],3))
                dxf[self.free_idx] = dx
                dx = dxf
            
            return -self.compute_force_differentials(dx)[self.free_idx,:].flatten()
        
        self.LHS = LHS # Save to class for testing
        
        # Define RHS
        RHS = (self.f + self.f_ext)[self.free_idx].flatten()
        
        self.RHS = RHS # Save to class for testing
        
        # Use conjugate gradient to find a descent direction
        # (see conjugate_gradient in Utils.py)
        dx_CG = conjugate_gradient(LHS,RHS).reshape(-1,3) # shape (#v, 3)
        if (dx_CG.shape[0] < self.v_rest.shape[0]):
                dxf = np.zeros((self.v_rest.shape[0],3))
                dxf[self.free_idx] = dx_CG
                dx_CG = dxf
        
        # Run line search on the direction
        step_size = step_size_init
        ft = self.f + self.f_ext
        energy_tot_prev = (self.energy_el + self.energy_ext)
        for l_iter in range(max_l_iter):
            step_size *= 0.5 # DO NOT CHANGE
            dx_search = dx_CG * step_size
            self.displace(dx_search)
            
            # Current force norm for unpinned vertices
            f_tmp = (self.f + self.f_ext)[self.free_idx]
            g = np.linalg.norm(f_tmp)

            energy_tot_tmp = self.energy_el + self.energy_ext
            
            #Make sure you only use the gradient of the unpinned vertices
            # USE parameter c1 = 1e-4 (from the function arguments) when you submit
            armijo = energy_tot_tmp < energy_tot_prev + c1 * step_size*(-dx_CG*dx_CG).sum()
            
            if armijo or l_iter == max_l_iter-1:
                print("Energy: " + str(energy_tot_tmp) + " Force norm: " + str(g) + " Line search Iters: " + str(l_iter))
                break
            else:
                self.displace(-dx_search)
