from fipy import CellVariable, Grid2D, Viewer, TransientTerm, DiffusionTerm, ImplicitSourceTerm, FaceVariable
from fipy.tools import numerix
import numpy as np
from matplotlib import pyplot as plt

def plot_internal_boundary_2d(x, y, internal_boundary_mask, L):

    mask_np = np.array(internal_boundary_mask.value)

    # Plot
    plt.figure(figsize=(6, 6))
    plt.scatter(x[mask_np], y[mask_np], c='red')
    plt.title("Internal boundary")
    plt.xlabel("x")
    plt.ylabel("y")
    plt.axis('equal')
    plt.grid(True)
    plt.xlim(0, L)
    plt.ylim(0, L)
    plt.legend()
    plt.show()

voxel_size = 1.

img = np.zeros((20, 20))
img[5:15, 5:15] = 1

solid_mask = numerix.array(img > 0, dtype=bool)
solid_mask_flat = solid_mask.flatten()

nx, ny = solid_mask.shape[1], solid_mask.shape[0]

dx = dy = voxel_size


dx = voxel_size
dy = dx
L = dx * nx
mesh = Grid2D(dx=dx, dy=dy, nx=nx, ny=ny)


# ---------determine internal boundary faces ---------
face_ids = mesh.facesLeft | mesh.facesRight | mesh.facesTop | mesh.facesBottom

adjacent_cells = mesh.faceCellIDs # the ids of the 2 cells adjacent to each face
cell1 = solid_mask_flat[adjacent_cells[0]]
cell2 = solid_mask_flat[adjacent_cells[1]]

internal_boundary_mask_np = (cell1 != cell2) & ~adjacent_cells[1].mask # a mask of the faces at the interface between cells labelled 0 and 1
internal_boundary_mask = FaceVariable(mesh=mesh, value=internal_boundary_mask_np)
# ----------------------------------------------------

x, y = mesh.faceCenters[0], mesh.faceCenters[1]
plot_internal_boundary_2d(x, y, internal_boundary_mask, L)


s = CellVariable(name = "solution variable",
                   mesh = mesh,
                   value = 0.)

s.constrain(0.0, where=~solid_mask_flat) # only solve in solid

# with the constant flux BC the solution is unqiue up to a constant. Randomly set the value of one cell in the solid:
random_id_mask = numerix.zeros_like(solid_mask_flat, dtype=bool)
random_id_mask[np.random.choice(np.where(solid_mask_flat)[0])] = True 
s.constrain(0.0, where=random_id_mask)


D0 = 1
D = FaceVariable(mesh=mesh, value=D0)


gradient = - 1 / D0 

largeValue = 1e10


eq = (TransientTerm() == DiffusionTerm(coeff=D)
       + DiffusionTerm(coeff=largeValue * internal_boundary_mask)
       - ImplicitSourceTerm((internal_boundary_mask * largeValue * gradient
                            * mesh.faceNormals).divergence))



timeStepDuration = 10 * 0.9 * dx**2 / (2 * D0)
steps = 10

s_values = []

from builtins import range
for step in range(steps):
    eq.solve(var=s,
             dt=timeStepDuration)
    
    s_values.append(s.value.copy().reshape(nx, ny))