from fenics import *

def g(x,y,rho):
    return conditional(le(x,(1-rho)*y),(1-(1-rho)**2)*y**2,conditional(le(y,(1-rho)*x),(1-(1-rho)**2)*x**2,x**2+y**2-2*(1-rho)*x*y))

def h(x,rho):
    return conditional(le(x,1-rho),0,(x-(1-rho))**2)

def phi0(F,rho):
    x = det(F)
    return conditional(le(x,0),1e10,conditional(le(x,rho),(ln(rho/x))**2,0))

def phi1(F,rho,f1,f2):
    x = sqrt(dot(F*f1,F*f1))
    y = sqrt(dot(F*f2,F*f2))
    return 0.5*g(x,y,rho) - rho*det(F)

def phi2(F,rho,e1,e2):
    x = sqrt(dot(F*e1,F*e1))
    y = sqrt(dot(F*e2,F*e2))
    return 2*rho*(1-rho)+h(x,rho)+h(y,rho)-2*rho*det(F)

def phi3(F,rho):
    x = sqrt(tr(F*F.T))/sqrt(2)
    return 2*rho*(1-rho)+2*h(x,rho)-2*rho*det(F)


def energy_density(u, psi, G,b0, b1, b2, b3,rho0, rho1,rho2,rho3,e1,e2,f1,f2,Sr):
    F = (( Identity(2)+ Sr+  grad(u)* inv(grad(psi)) ) * inv(G))
    return (b0*phi0(F,rho0)+b1*phi1(F,rho1,f1,f2)+b2*phi2(F,rho2,e1,e2)+b3*phi3(F,rho3))*abs(det(grad(psi)))



def energy_density_initial(u, psi, G, a1, a2, a3, a4,Sr):
	F = ( Identity(2)+ Sr +  grad(u)* inv(grad(psi)) ) * inv(G)
	C = F.T*F
	return (a1*(tr(C))**2 + a2*det(C) - a3*ln(det(C)) + a4*(C[0,0]**2+C[1,1]**2) - (4*a1+a2+2*a4))*abs(det(grad(psi)))



def get_rotation (u, psi,Sr):
	A = Identity(2)+ Sr + grad(u)* inv(grad(psi))
	U = (1/sqrt(abs(det(A+abs(det(A))*inv(A.T)))))*(A+abs(det(A))*inv(A.T))
	return atan(U[1,0]/U[0,0])