import numpy as np

def computeUbar(theGeoTransfo,l, u):
    if theGeoTransfo=='linear':
        uBar=computeUbar_Linear(l,u)
    elif theGeoTransfo=='corotational':
        uBar=computeUbar_Corotational(l,u)

    return uBar
       
        
def computeUbar_Linear(l,u):
    L = np.array([
        [-1, 0, 0, 1, 0, 0],
        [0, 1/l, 1, 0, -1/l, 0],
        [0, 1/l, 0, 0, -1/l, 1]])

    uBar =L@u

    return uBar


def computeUbar_Corotational(l,u):
    DeltaUx = u[3] - u[0]
    DeltaUy = u[4] - u[1]

    ln = np.sqrt((l + DeltaUx)**2 + DeltaUy**2)
    beta = np.arctan(DeltaUy / (l + DeltaUx))

    u1Bar = ln - l
    u2Bar = u[2] - beta
    u3Bar = u[5] - beta
    uBar = np.array([u1Bar, u2Bar, u3Bar])

    return uBar



def computeTransfoBasicToLocal(theGeoTransfo,l,u,qBar):
    if theGeoTransfo=='linear':
        LTransfo, Kgeom=computeTransfoBasicToLocal_Linear(l)
    elif theGeoTransfo=='corotational':
        LTransfo, Kgeom=computeTransfoBasicToLocal_Corotational(l,u,qBar)

    return LTransfo, Kgeom


def computeTransfoBasicToLocal_Linear(l):
    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]])

    Kgeom = np.zeros((6,6))

    return LTransfo, Kgeom


def computeTransfoBasicToLocal_Corotational(l,u,qBar):
    DeltaUx = u[3] - u[0]
    DeltaUy = u[4] - u[1]

    ln = np.sqrt((l + DeltaUx)**2 + DeltaUy**2)
    beta = np.arctan(DeltaUy / (l + DeltaUx))

    c = np.cos(beta)
    s = np.sin(beta)

    LTransfo = np.array([
        [-c, -s, 0, c, s, 0],
        [-s / ln, c / ln, 1, s / ln, -c / ln, 0],
        [-s / ln, c / ln, 0, s / ln, -c / ln, 1]])

    term1 = (qBar[0] / ln) * np.array([
        [s**2, -c * s, 0, -s**2, c * s, 0],
        [-c * s, c**2, 0, c * s, -c**2, 0],
        [0, 0, 0, 0, 0, 0],
        [-s**2, c * s, 0, s**2, -c * s, 0],
        [c * s, -c**2, 0, -c * s, c**2, 0],
        [0, 0, 0, 0, 0, 0]])

    term2 = ((qBar[1] + qBar[2]) / ln**2) * np.array([
        [-2 * s * c, c**2 - s**2, 0, 2 * s * c, -c**2 + s**2, 0],
        [c**2 - s**2, 2 * c * s, 0, -c**2 + s**2, -2 * c * s, 0],
        [0, 0, 0, 0, 0, 0],
        [2 * s * c, -c**2 + s**2, 0, -2 * s * c, c**2 - s**2, 0],
        [-c**2 + s**2, -2 * c * s, 0, c**2 - s**2, 2 * c * s, 0],
        [0, 0, 0, 0, 0, 0]])

    Kg = term1 + term2

    return LTransfo, Kg