import numpy as np 
from computeLocalElemStiffnessMatrix import computeLocalElemStiffnessMatrix 
from computeRotationMatrix import computeRotationMatrix


def assembleGlobalStiffnessMatrix(nDofTot,numEquations,element_data):
    K_global=np.zeros((nDofTot,nDofTot))

    for elemIdx, elem in enumerate(element_data):
        # Compute element stiffness matrices in local reference system
        K_elem_inLocalRef=computeLocalElemStiffnessMatrix(elem)

        # Compute element stiffness matrices in global reference system
        R=computeRotationMatrix(elem)
        K_elem_inGlobalRef=np.dot(np.dot(R.T,K_elem_inLocalRef),R)

        # Assemble global stiffness matrix
        dof_elem=np.array(numEquations[elemIdx]).astype(int)
        K_global[np.ix_(dof_elem,dof_elem)]+=K_elem_inGlobalRef

    return K_global