import numpy as np 


def computeRotationMatrix(element_data):
    if element_data[0]=='2dTruss':
        coordinates_element=np.array([element_data[2],element_data[3]])
        R=computeRotationMatrix_2dTruss(coordinates_element)
    elif element_data[0]=='2dElasticBeam' or element_data[0]=='2dModElasticBeam' or element_data[0]=='2dDispBeamColumn': 
        coordinates_element=np.array([element_data[2],element_data[3]])
        R=computeRotationMatrix_2dBeam(coordinates_element)
    elif element_data[0]=='spring':
        R=np.identity(2)
    return R


def computeRotationMatrix_2dTruss(coordinates_element):
    #Computes the rotation matrix for a 2D truss element    

    theta=np.arctan2(coordinates_element[1,1]-coordinates_element[0,1],coordinates_element[1,0]-coordinates_element[0,0])

    C=np.cos(theta)
    S=np.sin(theta)

    R = np.array([[C, S, 0, 0],[-S,C,0,0],[0,0,C,S],[0,0,-S,C]])

    return R


def computeRotationMatrix_2dBeam(coordinates_element):
    #Computes the rotation matrix for a 2D beam element    

    theta=np.arctan2(coordinates_element[1,1]-coordinates_element[0,1],coordinates_element[1,0]-coordinates_element[0,0])

    C=np.cos(theta)
    S=np.sin(theta)

    R =  np.array([
    [ C,  S, 0, 0, 0, 0],
    [-S,  C, 0, 0, 0, 0],
    [ 0,  0, 1, 0, 0, 0],
    [ 0,  0, 0, C,  S, 0],
    [ 0,  0, 0, -S, C, 0],
    [ 0,  0, 0, 0, 0, 1]
    ])

    return R