from fenics import *
from fenics_adjoint import *
from dolfin_adjoint import *
from mshr import *
from pyadjoint.overloaded_type import create_overloaded_object
import numpy


tol= 1E-14


class Edge:
	def __init__(self, data):
		self.pointnumber = len(data)
		self.data = data

class Subdomain:
	def __init__(self, edges, orientations):
		self.edgenumber = len(edges)
		self.edges = edges
		self.orientations = orientations

	def get_polygon(self):
		polygon = []
		for i in range (0,self.edgenumber):
			if (bool(self.orientations[i])):
				for j in range (0, (self.edges[i]).pointnumber-1):
					polygon.append((self.edges[i]).data[j])
			else:
				for j in range (0, (self.edges[i]).pointnumber-1):
					polygon.append((self.edges[i]).data[(self.edges[i]).pointnumber-1-j])
		return polygon


# check if a point x is on the edge
def on_polygon(x,edge,closed=True):
	for i in range(0,edge.pointnumber-1):
		x0 = ((edge.data)[i]).x()
		y0 = ((edge.data)[i]).y()
		x1 = ((edge.data)[i+1]).x()
		y1 = ((edge.data)[i+1]).y()
		if (near(x[0],x0,tol) and near(x[1],y0,tol) or near(x[0],x1,tol) and near(x[1],y1,tol)):
			if closed:
				return True
			else: 
				return False
		if near(x1,x0,tol):
			r = (x1 - x0) * (x[1] - y0) / (y1 - y0) + x0 - x[0] 
			if (near(r, 0.0, tol) and (((x[1] >= y0) and (x[1] <= y1)) or ((x[1] >= y1) and (x[1] <=y0)))):
				return True
		else:
			r = (y1 - y0) * (x[0] - x0) / (x1 - x0) + y0 - x[1]
			if (near(r, 0.0, tol) and (((x[0] >= x0) and (x[0] <= x1)) or ((x[0] >= x1) and (x[0] <= x0)))):
				return True
		#print ("( ", x0, " , ", y0, " ) ; ( ", x[0], " , ", x[1], " ) ; ( ", x1, " , ", y1, " )           return_val: ", result)
	return False


#                                                setting up the mesh                                         
def get_mesh(L_start,Lr,Ll,resolution,theta):

	# input edges
	top_left = Edge([Point(0.,1.0),Point(-Ll, 1.0)])
	top_mid = Edge([Point(L_start,1-0.5*theta),Point(0., 1.)])
	top_right = Edge([Point(L_start+Lr,1-0.5*theta),Point(L_start, 1-0.5*theta)])
	right_top = Edge([Point(L_start+Lr,0.5*theta),Point(L_start+Lr, 1.-0.5*theta)])
	right_bottom = Edge([Point(L_start+Lr,-0.5*theta),Point(L_start+Lr, 0.5*theta)])
	mid_right = Edge([Point(L_start,0.5*theta ),Point(L_start+Lr,0.5*theta )])
	mid_left = Edge([Point(0.,0.),Point(0., 1.)])
	mid_bottom = Edge([Point(0.,0.),Point(L_start, 0.5*theta)])
	left = Edge([Point(-Ll,1.),Point(-Ll, 0.)])
	bottom_left = Edge([Point(-Ll,0.),Point(0., 0.)])
	bottom_mid = Edge([Point(0.,0.),Point(L_start,-0.5*theta )])
	bottom_right = Edge([Point(L_start,-0.5*theta ),Point(L_start+Lr,-0.5*theta )])

	# define a vector with the edges
	edges = [top_left,top_mid,top_right,right_top,right_bottom,mid_right,mid_left,mid_bottom,left,bottom_left,bottom_mid,bottom_right]
	edges_number = len(edges)

	# input domain
	domain_complete = Subdomain([top_left,left,bottom_left, bottom_mid, bottom_right, right_bottom, right_top, top_right, top_mid],[0,0,0,0,0,0,0,0,0])

	# input subdomains
	domain_left = Subdomain([top_left,left,bottom_left,mid_left],[0,0,0,0])
	domain_top = Subdomain([mid_bottom,mid_right,right_top,top_right,top_mid,mid_left],[0,0,0,0,0,1])
	domain_bottom = Subdomain([bottom_mid,bottom_right,right_bottom,mid_right,mid_bottom],[0,0,0,1,1])

	# define vector of subdomains
	subdomains = [domain_left,domain_top,domain_bottom]
	subdomain_number = len(subdomains)

	# defining the domain, and the subdomains
	domain = Polygon(domain_complete.get_polygon())
	for i in range(0, subdomain_number):
		domain.set_subdomain (i+1, Polygon((subdomains[i]).get_polygon()))

	# generat the mesh
	mesh = generate_mesh (domain, resolution)
	sudom = MeshFunction ('size_t', mesh, 2, mesh.domains())


    # refining the mesh at the needle tip and the end of the needle part
	for k in range(1):
		cell_markers = MeshFunction("bool", mesh,2)
		cell_markers.set_all(False)
		coords = mesh.coordinates()
		for cell in cells(mesh):
			if abs(cell.midpoint().x()-L_start) < 0.5:
				cell_markers[cell] = True
		mesh = refine(mesh, cell_markers)
		sudom = adapt(sudom, mesh)
		print ("Anzahl Knoten:", mesh.num_vertices())

	for k in range(1):
		cell_markers = MeshFunction("bool", mesh,2)
		cell_markers.set_all(False)
		coords = mesh.coordinates()
		for cell in cells(mesh):
			if abs(cell.midpoint().x()-L_start) < 0.25:
				cell_markers[cell] = True
		mesh = refine(mesh, cell_markers)
		sudom = adapt(sudom, mesh)
		print ("Anzahl Knoten:", mesh.num_vertices())

	for k in range(2):
		cell_markers = MeshFunction("bool", mesh,2)
		cell_markers.set_all(False)
		coords = mesh.coordinates()
		for cell in cells(mesh):
			if cell.midpoint().distance(dolfin.Point((0,0))) < 0.09:
				cell_markers[cell] = True
		mesh = refine(mesh, cell_markers)
		sudom = adapt(sudom, mesh)
		print ("Anzahl Knoten:", mesh.num_vertices())

	for k in range(1):
		cell_markers = MeshFunction("bool", mesh,2)
		cell_markers.set_all(False)
		coords = mesh.coordinates()
		for cell in cells(mesh):
			if cell.midpoint().distance(dolfin.Point((0,0))) < 0.06:
				cell_markers[cell] = True
		mesh = refine(mesh, cell_markers)
		sudom = adapt(sudom, mesh)
		print ("Anzahl Knoten:", mesh.num_vertices())

	for k in range(1):
		cell_markers = MeshFunction("bool", mesh,2)
		cell_markers.set_all(False)
		coords = mesh.coordinates()
		for cell in cells(mesh):
			for v in vertices(cell):
				for i in [3,4,6,8]:
					if on_polygon((coords[v.index()][0],coords[v.index()][1]),edges[i]):
						cell_markers[cell] = True
		mesh = refine(mesh, cell_markers)
		sudom = adapt(sudom, mesh)
		print ("Anzahl Knoten:", mesh.num_vertices())


    # checking if every vertex on the top boundary has a corresponding vertex on the bottom boundary to satisfy the periodic boundary conditions
	for k in range(6):
		cell_markers = MeshFunction("bool", mesh,2)
		cell_markers.set_all(False)
		coords = mesh.coordinates()
		for v in vertices(mesh):
			matching_point_exists = False


			if on_polygon((coords[v.index()][0],coords[v.index()][1]),edges[9]) or on_polygon((coords[v.index()][0],coords[v.index()][1]),edges[10]) or on_polygon((coords[v.index()][0],coords[v.index()][1]),edges[11]):
				for v2 in vertices(mesh):
					if near(coords[v.index()][0],coords[v2.index()][0],1e-10) and near(coords[v.index()][1]+1.,coords[v2.index()][1],1e-10):
						matching_point_exists = True

				if not matching_point_exists:
					min_dist = 1.
					for cell in cells(mesh):
						min_dist = min(min_dist,cell.distance(dolfin.Point((coords[v.index()][0],coords[v.index()][1]+1.))))
						
						if cell.contains(dolfin.Point((coords[v.index()][0],coords[v.index()][1]+1))):
							cell_markers[cell] = True
		mesh = refine(mesh, cell_markers)
		sudom = adapt(sudom, mesh)
		print ("Anzahl Knoten:", mesh.num_vertices())


    # checking if every vertex on the bottom boundary has a corresponding vertex on the top boundary to satisfy the periodic boundary conditions
	for k in range(3):
		cell_markers = MeshFunction("bool", mesh,2)
		cell_markers.set_all(False)
		coords = mesh.coordinates()
		for v in vertices(mesh):
			matching_point_exists = False


			if on_polygon((coords[v.index()][0],coords[v.index()][1]),edges[0]) or on_polygon((coords[v.index()][0],coords[v.index()][1]),edges[1]) or on_polygon((coords[v.index()][0],coords[v.index()][1]),edges[2]):
				for v2 in vertices(mesh):
					if near(coords[v.index()][0],coords[v2.index()][0],1e-10) and near(coords[v.index()][1]-1.,coords[v2.index()][1],1e-10):
						matching_point_exists = True
				if not matching_point_exists:
					min_dist = 1.
					for cell in cells(mesh):
						min_dist = min(min_dist,cell.distance(dolfin.Point((coords[v.index()][0],coords[v.index()][1]-1.))))
						
						if cell.contains(dolfin.Point((coords[v.index()][0],coords[v.index()][1]-1))):
							cell_markers[cell] = True
		mesh = refine(mesh, cell_markers)
		sudom = adapt(sudom, mesh)
		print ("Anzahl Knoten:", mesh.num_vertices())


    # checking if every vertex on the top boundary has a corresponding vertex on the bottom boundary to satisfy the periodic boundary conditions
	for k in range(3):
		cell_markers = MeshFunction("bool", mesh,2)
		cell_markers.set_all(False)
		coords = mesh.coordinates()
		for v in vertices(mesh):
			matching_point_exists = False


			if on_polygon((coords[v.index()][0],coords[v.index()][1]),edges[9]) or on_polygon((coords[v.index()][0],coords[v.index()][1]),edges[10]) or on_polygon((coords[v.index()][0],coords[v.index()][1]),edges[11]):
				for v2 in vertices(mesh):
					if near(coords[v.index()][0],coords[v2.index()][0],1e-10) and near(coords[v.index()][1]+1.,coords[v2.index()][1],1e-10):
						matching_point_exists = True

				if not matching_point_exists:
					min_dist = 1.
					for cell in cells(mesh):
						min_dist = min(min_dist,cell.distance(dolfin.Point((coords[v.index()][0],coords[v.index()][1]+1.))))
						
						if cell.contains(dolfin.Point((coords[v.index()][0],coords[v.index()][1]+1))):
							cell_markers[cell] = True
		mesh = refine(mesh, cell_markers)
		sudom = adapt(sudom, mesh)
		print ("Anzahl Knoten:", mesh.num_vertices())


	mesh = create_overloaded_object(mesh)
	

	# defining coefficients on subdomains
	X = FunctionSpace (mesh, "DG", 0)
	dm = X.dofmap()
	sudom_arr = numpy.asarray (sudom.array(), dtype=numpy.int)
	for cell in cells (mesh): sudom_arr [dm.cell_dofs (cell.index())] = sudom [cell]

	def sudom_fct (sudom_arr, vals, fctspace):
		f = Function (fctspace)
		f.vector()[:] = numpy.choose (sudom_arr, vals)
		return f

	chi_a = sudom_fct (sudom_arr, [0,1,0,1], X)
	chi_b = sudom_fct (sudom_arr, [0,0,1,0], X)
	chi_test = sudom_fct (sudom_arr, [0,1,2,1], X)
	return mesh, edges, edges_number, chi_a, chi_b, chi_test,mesh.num_vertices()