from fenics import *
from fenics_adjoint import *
from dolfin_adjoint import *
from energy_density import energy_density, energy_density_initial, get_rotation
from mesh_setup import get_mesh, on_polygon
from deformation import get_deformation




#                                                 the programm                                        
def do_shape_opt(L_start,Lr,Ll,resolution,delta,theta,a1,a2,a3,a4,b_0,b_1,b_2,b_3,rho0,rho1,rho2,rho3):
	mesh, edges, edges_number, chi_a, chi_b, chi_test, verts = get_mesh(L_start,Lr,Ll,resolution,theta)
	x = SpatialCoordinate(mesh)

	class PeriodicBoundary (SubDomain):
		# bottom boundary is target domain
		def inside (self, x, on_boundary): return bool ( (on_polygon(x,edges[9]) and on_boundary) or	(on_polygon(x,edges[10]) and on_boundary)  or (on_polygon(x,edges[11]) and on_boundary) )
		# Map top boundary to bottom boundary
		def map (self, x, y): y[0] = x[0]; y[1] = x[1]-1.0

	# create the function spaces
	U = FunctionSpace (mesh, "CG", 1)
	U2 = FunctionSpace (mesh, "DG", 0)
	V = VectorFunctionSpace (mesh, "CG", 1, constrained_domain=PeriodicBoundary())
	W = VectorFunctionSpace(mesh, 'CG', 1)

	u = Function(V)
	v = TestFunction(V)


	L_opt = Constant(L_start)
	Delta = Constant(theta*0.5)
	ab = Constant(0)
	at = Constant(0)

	psi, dpsi = get_deformation(L_start,Lr,Ll,resolution,delta,theta,L_opt,Delta,ab,at,edges,edges_number,W,x)

	# define the energy density
	GA = Constant (((1,delta), (0,1)))
	GB = Constant (((1,-delta), (0,1)))
	G = chi_a*GA + chi_b*GB

	e1 = Constant((1,0))
	e2 = Constant((0,1))
	f1 = Constant((1/sqrt(2),1/sqrt(2)))
	f2 = Constant((-1/sqrt(2),1/sqrt(2)))

	S2 = Constant(((0,(2*delta*theta-delta)),(0,0)))
	

	class Gamma(SubDomain):
		def inside(self, x, on_boundary):
			return near(x[0], 0., DOLFIN_EPS) and near(x[1], 0, DOLFIN_EPS)


	zero = Constant ((0,0))
	dbcopt = DirichletBC (V, zero, Gamma(), method ="pointwise")
	bcsopt = [dbcopt]


    # Initialisation of u as the solution to the elastic problem with a different energy density to ensure convergence of the newton solver in the actual problem
	E_init =  energy_density_initial (u, psi, G, a1, a2, a3, a4,S2)*dx
	E = energy_density (u, psi, G,b_0,b_1,b_2,b_3,rho0,rho1,rho2,rho3,e1,e2,f1,f2,S2)*dx

	F_init = derivative (E_init, u, v)
	print ("******* compute the initial elastic deformation:")
	solve (F_init == 0, u, bcsopt, solver_parameters={"newton_solver":{"linear_solver":'lu',"relative_tolerance":1e-10}})

	F = derivative (E, u, v)
	print ("******* compute the elastic deformation:")
	solve (F == 0, u, bcsopt, solver_parameters={"newton_solver":{"linear_solver":'lu',"relative_tolerance":1e-10}})
	startE = assemble(E)

	print ("********** E_start = %f" % startE, flush=True)

	

	
	controls = [Control(L_opt),Control(Delta),Control(ab),Control(at)]
	Ehat = ReducedFunctional(assemble(E), controls)

	#         Minimization            
	rL_opt, rDelta, rab, rat = minimize (Ehat, method = 'L-BFGS-B', options = {'disp': True,'maxiter':200,'ftol':1e-16,'gtol':1e-12})
	print (float(rL_opt),float(rDelta),float(rab),float(rat))
	


	# calculate the final deformations and energies
	psi, dpsi = get_deformation(L_start,Lr,Ll,resolution,delta,theta,float(rL_opt),float(rDelta),float(rab),float(rat),edges,edges_number,W,x)
	u_end = Function (V, name='displacement')



	energy_final_init =  energy_density_initial (u_end, psi, G, a1, a2, a3, a4,S2)*dx
	energy_final = energy_density (u_end, psi, G,b_0,b_1,b_2,b_3,rho0,rho1,rho2,rho3,e1,e2,f1,f2,S2) *dx

	F_final_init = derivative (energy_final_init, u_end, v)
	F_final = derivative (energy_final, u_end, v)

	solve (F_final_init == 0, u_end, bcsopt, solver_parameters={"newton_solver":{"linear_solver":'lu',"relative_tolerance":1e-10}})
	solve (F_final == 0, u_end, bcsopt, solver_parameters={"newton_solver":{"linear_solver":'lu',"relative_tolerance":1e-10}})


	
	energy_dens = project(energy_density(u_end, psi, G,b_0,b_1,b_2,b_3,rho0,rho1,rho2,rho3,e1,e2,f1,f2,S2),U2)

	rotation = project(get_rotation(u_end, psi, S2),U2)

	E_end = assemble(energy_density(u_end, psi, G,b_0,b_1,b_2,b_3,rho0,rho1,rho2,rho3,e1,e2,f1,f2,S2)*dx)


	return E_end,float(rat),float(rab),float(rDelta),float(rL_opt), chi_test,dpsi,u,verts,energy_dens,rotation