import numpy as np
import openpnm as op
%config InlineBackend.figure_formats = ['svg']
import matplotlib.pyplot as plt
np.random.seed(10)
ws = op.Workspace()

"""
Explanation: Fickian Diffusion with Variable Diffusivity
In this example, we will implement a Fickian diffusion algorithm on a Cubic network assumming a variable diffusivity, such that diffusivity is significantly higher at high concentrations.
End of explanation
"""


net = op.network.Cubic([1, 15, 15], spacing=1e-5)
geom = op.geometry.SpheresAndCylinders(network=net, pores=net.Ps, throats=net.Ts)
air = op.phases.Water(network=net)

"""
Explanation: Creating network, phase and geometry
We create a Cubic network and add the geometry and phase objects. Note that the same procedure can be applied on extracted networks and 3D networks. Here for simplicity and visualization we chose a 2D cubic network.
End of explanation
"""


def variable_diffusivity(target, c_ref, pore_concentration="pore.concentration"):
    X = target[pore_concentration]
    val = 1e-9 * (1 + 5 * (X / c_ref) ** 2)
    return val

air.add_model(propname="pore.diffusivity", model=variable_diffusivity, c_ref=10.0,
              regen_mode="deferred")
phys = op.physics.Standard(network=net, geometry=geom, phase=air)

"""
Explanation: Adding the diffusivity model
We assume that the diffusivity of the phase air is not constant and is a function of the concentration of the species. Therefore, we define a model to calculate the pore.diffusivity at each concentration.
End of explanation
"""


fd = op.algorithms.FickianDiffusion(network=net, phase=air)
fd.set_value_BC(pores=net.pores('top'), values=10.0)
fd.set_value_BC(pores=net.pores('bottom'), values=0.0)
fd.run()

"""
Explanation: Performing Fickian diffusion
We then add the the FickianDiffusion algorithm to our simulation and define "Dirichlet" boundary conditions assigned to the  pores that are on the top and the bottom face of the network. Note that values=10.0 for the first boundary condition means that the pore concentration of each pore (that is on the top boundary face) is equal to 10.0.
End of explanation
"""


air.update(fd.results())
c=air['pore.concentration']
fig, ax = plt.subplots()
c_xz_avg = c.reshape(op.topotools.get_shape(net)).mean(axis=(0, 1))
ax.plot(c_xz_avg, "ko-")
ax.set_xlabel(r"z")
ax.set_ylabel(r"c")
plt.show()

"""
Explanation: Visualizing
We can visualize the average concentration distribution along the x direction (from the left boundary face towards the right boundary face)
End of explanation
"""


c2d = c.reshape(op.topotools.get_shape(net)).squeeze()
plt.imshow(np.rot90(c2d))
plt.colorbar()

"""
Explanation: The concentration of species throughout the network is:
End of explanation
"""


