from dolfin import *

# Create mesh and define function space
mesh = Mesh("dolfin-channel.xml")

# Define function space
V = FunctionSpace(mesh, "Lagrange", 1)

# Define Dirichlet boundary (x = 0 or x = 1)
def dirichlet_boundary(x):
    return x[0] < DOLFIN_EPS or x[0] > 1.0 - DOLFIN_EPS

# Define boundary condition
g_D = Constant(0.0)
bc = DirichletBC(V, g_D, dirichlet_boundary)

#Define conductivity
k_1 = Expression("50 + exp(50*pow(0.5 - x[1],2))")
k_2 = 10

# Define source terms
f = Constant(1.0)
# Define g for the Robin boundary condition
g = Expression("sin(DOLFIN_PI*x[0])*sin(DOLFIN_PI*x[1])")

# Define a sub domain
class CircleDomain(SubDomain):
    def inside(self, x, on_boundary):
        dist = (x[0]-0.5)*(x[0]-0.5) + (x[1]-0.5)*(x[1] - 0.5)
        max_dist = 0.35*0.35
        return  dist < max_dist + DOLFIN_EPS 

#Create a CircleDomain object
circle_domain = CircleDomain()

# Create cell/element markers
domain_markers = CellFunction("uint",mesh)
domain_markers.set_all(2)
circle_domain.mark(domain_markers, 1)

# Write out domain markers in VTK format.
# You can do something similar for the boundary markers
domain_marker_file = File("domain_marker_5.pvd")
domain_marker_file << domain_markers

# Create facet markers
# Note: It marks *all* facets, not only boundary ones.
# But *ds integrate over over boundary facets, so it does
# not matter
boundary_markers = FacetFunction("uint",mesh)
boundary_markers.set_all(2)
circle_domain.mark(boundary_markers, 1)

#Redefine dx measure to make dx aware of domains
dx = Measure("dx")[domain_markers]

#Redefine ds measure to make dx aware of domains
ds = Measure("ds")[boundary_markers]

# Define trial an test functions
u = TrialFunction(V)
v = TestFunction(V)

# Define variational forms
a = (k_1*inner(grad(u), grad(v))  + u*v)*dx(1) \
    + k_2*inner(grad(u), grad(v))*dx(2) + k_2*u*v*ds(2)
L = f*v*dx(1) + f*v*dx(2) + k_2*g*v*ds(2)

# Compute solution
u = Function(V)
solve(a == L, u, bc)

# Save solution in VTK format
u_file = File("poisson_5.pvd")
u_file << u

# Plot solution
plot(u, interactive=True)