import numpy as np

def computeInitialLocalMaterialElemStiffnessMatrix(element_data):
    if element_data[0]=='2dTruss':
        E=element_data[4]
        A=element_data[5]
        L=np.linalg.norm(element_data[3]-element_data[2])
        K=computeInitialLocalMaterialElemStiffnessMatrix_2dTruss(E,A,L)
    elif element_data[0]=='2dElasticBeam': 
        E=element_data[4]
        A=element_data[5]
        I=element_data[6]
        L=np.linalg.norm(element_data[3]-element_data[2])
        K=computeInitialLocalMaterialElemStiffnessMatrix_2dBeam(E,A,I,L)
    elif element_data[0]=='spring':
        K=np.array(element_data[3])
    return K


def computeInitialLocalMaterialElemStiffnessMatrix_2dTruss(E,A,L):
    #Computes the local element stiffness matrix for a 2D truss element
    
    K = E*A/L*np.array([[1,0,-1,0],
                         [0,0,0,0],
                         [-1,0,1,0],
                         [0,0,0,0]])
    
    return K


def computeInitialLocalMaterialElemStiffnessMatrix_2dBeam(E,A,I,L):
    # Computes the local element stiffness matrix for a 2D beam element

    K = np.array([
    [ E*A/L,         0,          0,        -E*A/L,         0,           0],
    [    0,   12*E*I/L**3,   6*E*I/L**2,         0, -12*E*I/L**3,   6*E*I/L**2],
    [    0,    6*E*I/L**2,    4*E*I/L,         0,  -6*E*I/L**2,    2*E*I/L],
    [-E*A/L,         0,          0,         E*A/L,         0,           0],
    [    0,  -12*E*I/L**3,  -6*E*I/L**2,         0,  12*E*I/L**3,  -6*E*I/L**2],
    [    0,    6*E*I/L**2,    2*E*I/L,         0,  -6*E*I/L**2,    4*E*I/L]
    ])

    return K