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':
        ke_s=element_data[3][5]
        K=np.array([[ke_s,-ke_s],
                    [-ke_s,ke_s]])
    elif element_data[0]=='2dModElasticBeam':
        E=element_data[4]
        A=element_data[5]
        Ie=element_data[10]
        L=np.linalg.norm(element_data[3]-element_data[2])
        S22mod=element_data[6]
        S23mod=element_data[7]
        S32mod=element_data[8]
        S33mod=element_data[9]
        K=computeInitialLocalMaterialElemStiffnessMatrix_2dModElasticBeam(E,A,S22mod,S23mod,S32mod,S33mod,Ie,L)
    elif element_data[0]=='2dDispBeamColumn':
        E=element_data[6][0]
        L=np.linalg.norm(element_data[3]-element_data[2])
        coordinates_fibers=element_data[7]
        area_fibers=element_data[8]
        A=0
        I=0
        for iFiber in range(len(coordinates_fibers)):
            A+=area_fibers[iFiber]
            I+=area_fibers[iFiber]*(coordinates_fibers[iFiber][1])**2
        K=computeInitialLocalMaterialElemStiffnessMatrix_2dBeam(E,A,I,L)
    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


def computeInitialLocalMaterialElemStiffnessMatrix_2dModElasticBeam(E,A,S22mod,S23mod,S32mod,S33mod,Ie,l): 
    # Computes the local element stiffness matrix for a 2D elastic beam element with stiffness modifiers

    #Tangent stiffness matrix in the basic reference frame
    KBar=np.array([
    [E * A / l, 0, 0],
    [0, S22mod * E * Ie / l, S23mod * E * Ie / l],
    [0, S32mod * E * Ie / l, S33mod * E * Ie / l]])

    #Linear transformation matrix
    LTransfo = np.array([
        [-1, 0, 0, 1, 0, 0],
        [0, 1/l, 1, 0, -1/l, 0],
        [0, 1/l, 0, 0, -1/l, 1]])

    # Tangent stiffness matrix in the local reference frame
    Kloc=LTransfo.T@KBar@LTransfo

    return Kloc