import dolfin as df
import matplotlib.pyplot as plt

import mshr

# Geometry
isolation_depth = .6  # m
isolation_thickness = .06  # m
house_half_size = 4  # m
domain_size = 14  # m

domain = mshr.Rectangle(df.Point(-domain_size, -1.5*domain_size),
                        df.Point(house_half_size, 0.))
isolation = mshr.Rectangle(df.Point(0.0, -isolation_depth),
                           df.Point(isolation_thickness, .0))

domain.set_subdomain(1, isolation)

# Mesh
mesh = mshr.generate_mesh(domain, 120)

df.plot(mesh);
plt.xlim([-2, 2]); plt.ylim([-2, 0]);  # zoom

class SurfaceExt(df.SubDomain):
    def inside(self, x, on_boundary):
        surface = on_boundary and\
                  df.near(x[1], 0, df.DOLFIN_EPS) and\
                  x[0]<0
        left_side = df.near(x[0], -domain_size, df.DOLFIN_EPS)
        bottom_side = df.near(x[1], -1.5*domain_size, df.DOLFIN_EPS)
        return surface or left_side or bottom_side
    
class SurfaceInt(df.SubDomain):
    def inside(self, x, on_boundary):
        return on_boundary and\
               df.near(x[1], 0, df.DOLFIN_EPS) and\
               x[0] > isolation_thickness-df.DOLFIN_EPS

# Marquage des frontières:
surface_ext = SurfaceExt()
surface_int = SurfaceInt()

boundary_markers = df.FacetFunction("size_t", mesh)

boundary_markers.set_all(0)
surface_ext.mark(boundary_markers, 1)
surface_int.mark(boundary_markers, 2)

# boundary_markers.array()

# Marquage des cellules, à partir de la géométrie:
cell_markers = df.CellFunction("size_t", mesh)
markers = df.MeshFunction('size_t', mesh, 2, mesh.domains())

df.plot( markers )
plt.xlim([-2, 2]); plt.ylim([-2, 0]);  # zoom

# Définition des élements infinitésimaux pour les intégrations
dx = df.Measure('dx', domain=mesh, subdomain_data=markers)
ds = df.Measure('ds', domain=mesh, subdomain_data=boundary_markers)

V = df.FunctionSpace(mesh, "Lagrange", 1)

# Define boundary condition
u_ext = df.Constant(0.0)
bc_ext = df.DirichletBC(V, u_ext, surface_ext)

u_int = df.Constant(20.0)
bc_int = df.DirichletBC(V, u_int, surface_int)

bcs = [bc_ext, bc_int]

# Define input data
k_soil = df.Constant(1.28)  # Soil, W/m/K
k_isolant = df.Constant(0.04) # Isolation

f = df.Constant(.0)  # Source

u = df.TrialFunction(V)
v = df.TestFunction(V)

# Define variational form
a = df.inner(k_soil*df.grad(u), df.grad(v))*dx(0)\
    + df.inner(k_isolant*df.grad(u), df.grad(v))*dx(1)

L = f*v*dx(0) + f*v*dx(1)

# Solve problem
u = df.Function(V)
df.solve(a == L, u, bcs)

# Graph
plt.figure(figsize=(12, 8))
ax = df.plot(u)
plt.xlabel('x [m]');
plt.ylabel('depth [m]');
#plt.xlim([-2, 2]); plt.ylim([-2, 0]);
plt.colorbar(ax); plt.title('temperature stationnaire');

# Zoom
plt.figure(figsize=(12, 8))
ax = df.plot(u)
plt.xlabel('x [m]');
plt.ylabel('depth [m]');
plt.xlim([-2, 2]); plt.ylim([-2, 0]);
plt.colorbar(ax); plt.title('temperature stationnaire');

# Length of the interior boundary:
L_int = df.assemble( df.Constant(1.)*ds(2) )
print(L_int)

# Flux chaleur vers le sol (linéïque):
F = df.assemble( k_soil * df.Dx(u, 1)*ds(2) )
print(F, ' W/m')

# Flux moyen surfacique:
print(F/L_int, ' W/m2')

# Flux sortant de `domain`
n = df.FacetNormal(mesh)
flux_1 = df.assemble( k_soil*df.dot(df.grad(u), n)*ds(1))
flux_1