import igl
import numpy as np
from scipy import sparse
from scipy.sparse import csr_matrix
import numpy.linalg as la
def vertex_tri_sum(data,F,V):
    '''
    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     : array of shape (#t,)

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

    if len(data) == len(F):
        data = data[j]

    m = sparse.coo_matrix((data, (i, j)), (len(V), len(F)))
    return np.array(m.sum(axis=1)).flatten()


def compute_mass_matrix(V, F):
    ''' Assemble the Mass matrix by computing quantities per face.

    Parameters:
    - V : np.array (#v, 3)
    - F : np.array (#f, 3)

    Returns:
    - M : scipy sparse diagonal matrix (#v, #v)
        Mass matrix
    '''
    A = V[F[:,0],:]
    B = V[F[:,1],:]
    C = V[F[:,2],:]
    v = la.norm(np.cross((B-A),(C-A),axis=1),axis=1)/6
    M = sparse.diags(vertex_tri_sum(v,F,V) ,format='csr')
    return M

def compute_cotangent(a, b, c, A):
    ''' Compute the cotangent of an angle in a triangle by using the triangle edge lengths and area only. The input parameters are defined in the handout figure. The purpose of this function is to check that your formula is correct. You should not directly use this in the `compute_laplacian_matrix` function.

    Parameters:
    - a : float
    - b : float
    - c : float
    - A : float

    '''
    s_c = (2*A/(a*b))
    c_c = np.sqrt(1-s_c**2)
    return c_c/s_c


vi = [(0, 1, 2), (1, 2, 0), (2, 0, 1)]
def compute_laplacian_matrix(V, F):
    ''' Assemble the Laplacian matrix by computing quantities per face.

    Parameters:
    - V : np.array (#v, 3)
    - F : np.array (#f, 3)

    Returns:
    - L : scipy sparse matrix (#v, #v)
        Laplacian matrix
    '''
    sz = (3*2,F.shape[0])
    W, I, J = np.zeros(sz), np.zeros(sz, np.int32),np.zeros(sz, np.int32)   

    Wc, Ic, Jc = 0,0,0
    for v1, v2, v3 in vi: 
        vi1 = F[:,v1]
        vi2 = F[:,v2]
        vi3 = F[:,v3]

        a, b = V[vi2] - V[vi1], V[vi3] - V[vi1]
        cotan = np.sum(a*b,axis=1) / np.sqrt(np.sum(np.cross(a, b)**2, axis=-1))

        W[Wc] = W[Wc+1] =  cotan/2
        I[Ic] = J[Jc+1] = vi2
        J[Jc] = I[Ic+1] = vi3
        Ic,Jc,Wc=Ic+2,Jc+2,Wc+2

    I, J, W = I.flatten(), J.flatten(), W.flatten()
    
    L = sparse.csr_matrix((W, (I, J)))
    D = sparse.spdiags(L * np.ones(V.shape[0]), 0,  V.shape[0],  V.shape[0])

    L -= D
    return L
