import numpy as np

from computeRotationMatrix import computeRotationMatrix
from basicLocalGeoTransfo import computeUbar, computeTransfoBasicToLocal

def elemStateDetermin(elemIdx,elem,numEquations, v):
    dofs_elem=numEquations[elemIdx]

    T=computeRotationMatrix(elem)
    u=np.dot(T,v[dofs_elem])

    if elem[0]=='2dElasticBeam':
        qBar, KBar=elemStateDetermin_2dElasticBeam(elem,u)

        l=np.linalg.norm(elem[3]-elem[2])
        LTransfo, Kgeom=computeTransfoBasicToLocal(elem[7],l,u,qBar)

        QLocal_elem = LTransfo.T @ qBar
        QGlobal_elem = T.T @ QLocal_elem
        KMaterialGlobal_elem = T.T  @ LTransfo.T @ KBar @ LTransfo @ T
        KGeomGlobal_elem = T.T @ Kgeom @ T

    return QGlobal_elem, KMaterialGlobal_elem, KGeomGlobal_elem, QLocal_elem   



def elemStateDetermin_2dElasticBeam(elem,u):
    l=np.linalg.norm(elem[3]-elem[2])
    E=elem[4]
    A=elem[5]
    I=elem[6]
    KBar=np.array([
    [E * A / l, 0, 0],
    [0, 4 * E * I / l, 2 * E * I / l],
    [0, 2 * E * I / l, 4 * E * I / l]])

    uBar=computeUbar(elem[7],l,u)

    qBar=np.dot(KBar,uBar)

    return qBar, KBar