from fenics import *
from mesh_setup import on_polygon
from dolfin_adjoint import *



class Identity2(UserExpression):
	def __init__(self, **kwargs):
		super().__init__(**kwargs)

	def eval_cell(self,values,x,cell):
		values[0] = x[0]
		values[1] = x[1]

	def value_shape(self):
		return (2,)

# calculating the deformation on every edge                    
def get_deformation(L_start,Lr,Ll,resolution,delta,theta,L_opt,Delta,ab,at,edges,edges_number,W,x):
	sindt = delta*theta

	#top left
	edges[0].deformation = project(as_vector((x[0]-sindt,1.)),W)
	def boundary_top_left(x, on_boundary):
		return on_polygon(x,edges[0])

	#top mid
	edges[1].deformation = project(as_vector(((x[0]*(L_opt/L_start))  -  ((ab*(x[0]*(1/L_start))*(x[0]*(1/L_start))  +  ((Delta-theta)/1  - ab*1)  *  (x[0]*(1/L_start)) +1.) * sindt)  ,   ab*(x[0]*(1/L_start))*(x[0]*(1/L_start))  +  ((Delta-theta)/1  - ab*1)  *  (x[0]*(1/L_start)) +1.  )),W)
	def boundary_top_mid(x, on_boundary):
		return on_polygon(x,edges[1], closed=False)

	#top right
	edges[2].deformation = project(as_vector((((x[0]-L_start)/Lr)* (L_start+Lr-L_opt) + L_opt - (1.-theta+Delta) * sindt, 1.- theta + Delta)),W)
	def boundary_top_right(x, on_boundary):
		return on_polygon(x,edges[2])

	#right_top
	edges[3].deformation = project(as_vector(( L_start+Lr-   (x[1]+ Delta-0.5*theta)*sindt, x[1]+ Delta-0.5*theta)),W)
	def boundary_right_top(x, on_boundary):
		return on_polygon(x,edges[3], closed=False)

	#right_bottom
	edges[4].deformation = project(as_vector(( L_start+Lr-(x[1]+ Delta-0.5*theta)*sindt,x[1]+ Delta-0.5*theta)),W)
	def boundary_right_bottom(x, on_boundary):
		return on_polygon(x,edges[4], closed=False)

	#mid_right
	edges[5].deformation = project(as_vector((((x[0]-L_start)/Lr)* (L_start+Lr-L_opt) + L_opt - Delta * sindt,Delta)),W)
	def boundary_mid_right(x, on_boundary):
		return on_polygon(x,edges[5])

	#mid_left
	edges[6].deformation = project(as_vector((-x[1]*sindt,x[1])),W)
	def boundary_mid_left(x, on_boundary):
		return on_polygon(x,edges[6], closed=False)

	#mid_bottom
	edges[7].deformation = project(as_vector(((x[0]*(L_opt/L_start))  -  ((at*  (x[0]*(1/L_start))*(x[0]*(1/L_start))  +  (Delta/1  - at*1)  *  (x[0]*(1/L_start))) * sindt) ,   at*  (x[0]*(1/L_start))*(x[0]*(1/L_start))  +  (Delta/1  - at*1)  *  (x[0]*(1/L_start)))),W)
	def boundary_mid_bottom(x, on_boundary):
		return on_polygon(x,edges[7], closed=False)

	#left
	edges[8].deformation = project(as_vector((x[0]-x[1]*sindt,x[1])),W)
	def boundary_left(x, on_boundary):
		return on_polygon(x,edges[8], closed=False)

	#bottom_left
	edges[9].deformation = project(as_vector((x[0],x[1])),W)
	def boundary_bottom_left(x, on_boundary):
		return on_polygon(x,edges[9])

	#bottom_mid
	edges[10].deformation = project(as_vector(((x[0]*(L_opt/L_start))  -  ((ab*(x[0]*(1/L_start))*(x[0]*(1/L_start))  +  ((Delta-theta)/1  - ab*1)  *  (x[0]*(1/L_start))) * sindt)  ,   ab*(x[0]*(1/L_start))*(x[0]*(1/L_start))  +  ((Delta-theta)/1  - ab*1)  *  (x[0]*(1/L_start)))),W)
	def boundary_bottom_mid(x, on_boundary):
		return on_polygon(x,edges[10], closed=False)

	#bottom_right
	edges[11].deformation = project(as_vector((((x[0]-L_start)/Lr)* (L_start+Lr-L_opt) + L_opt - (Delta-theta) * sindt,Delta-theta)),W)
	def boundary_bottom_right(x, on_boundary):
		return on_polygon(x,edges[11])

	# define a vector with the boundary conditions
	boundary_edges = [boundary_top_left,boundary_top_mid,boundary_top_right,boundary_right_top,boundary_right_bottom,boundary_mid_right,boundary_mid_left,boundary_mid_bottom,boundary_left,boundary_bottom_left,boundary_bottom_mid,boundary_bottom_right]

	bcs = []
	for i in range(0,edges_number):
		bc = DirichletBC(W,(edges[i]).deformation, boundary_edges[i])
		bcs.append(bc)

	#print (bcs)

	# compute the deformation following eq. (4.4) 
	print ("******* compute the deformation from computational domain to fundamental cell:")
	psit=TrialFunction(W)
	vt=TestFunction(W)
	a = inner(grad(psit),grad(vt)) *dx
	psi=Function(W, name='psi')
	solve(lhs(a)==rhs(a), psi, bcs)

	# compute the displacement
	id = project(Identity2(),W)
	dpsi = Function(W, name='dpsi')
	dpsi = project(psi-id,W)
	return psi,dpsi