from dolfin import *
import random

"""
    OBS!!!
    Tested for fenics version 2019
    The fenics library can be loaded on the terminal computers at uio using
    the terminal command "module load FEniCS". Then run the script from command line
    with "python3 cahn-hilliardFEM.py".
    To install software on own computer you can find installation information for your
    operating system here: https://fenicsproject.org/download/archive/

    Solver is implemented with homogenous NeumanBC at all edges.
    To include Dirichlet conditions add BCs to solve (see line 132/146)

    Writes to results/c.pvd and results/mu.pvd which can be read in paraview
    For other options, access c/mu.vector().array() and
    write to file.
"""



##### DEFINE VARIABLES/PARAMETERS #####

Cahn_number  = 0.05
Pechlet_number = 2
dt     = 0.01
t      = 0.0     # initial time
T      = 100*dt  # end time
theta  = 0.5     # coefficient for discretization scheme:
                 # 0: forward Euler
                 # 0.5: midpoint,
                 # 1: backwards euler,
N      = 100     # number of elements in each dimension
IC     = 2       # Initial condition, 1=i, 2=ii


##### DEFINE MESH AND FUNCTION SPACE #####

mesh    = UnitSquareMesh(N, N)
elms    = FiniteElement('CG', mesh.ufl_cell(), 1) # Linear elements
melms   = MixedElement([elms, elms])
V       = FunctionSpace(mesh, elms)  # Linear functionspace
ME      = FunctionSpace(mesh, melms)   # Mixed function space


##### DEFINE BOUNDARIES #####

class Left(SubDomain):
    def inside(self, x, on_boundary):
        return (x[0] < DOLFIN_EPS) and on_boundary

class Right(SubDomain):
    def inside(self, x, on_boundary):
        return (1 - x[0] < DOLFIN_EPS) and on_boundary


boundaries = MeshFunction("size_t", mesh, mesh.topology().dim() - 1)
boundaries.set_all(0)
left = Left()
left.mark(boundaries, 1)
right = Right()
right.mark(boundaries, 2)

##### FUNCTIONS #####

u       = Function(ME)
u0      = Function(ME)

du      = TrialFunction(ME)
q, v    = TestFunction(ME)

dc, dmu = split(du)
c, mu   = split(u)
c0, mu0 = split(u0)



##### DEFINE INITIAL CONDITIONS #####

class half_and_half(UserExpression):
    """
        Assignment initial condition
    """
    def eval(self, values, x):
        if x[0] < 0.5:
            values[0] = -1
        elif x[0] == 0.5:
            values[0] = 0
        else:
            values[0] = 1
    def value_shape(self):
       return (2, )


class RandomIC(UserExpression):
    """
        Random noise
    """
    def eval(self, values, x):
        values[0] = 0.02*(0.5 - random.random())

    def value_shape(self):
       return (2, )

if IC == 1:
    IC  = half_and_half()
elif IC == 2:
    IC  = RandomIC()

u.interpolate(IC)
u0.interpolate(IC)


##### CHOOSE F AS A FOURTH ORDER POLYNOMIAL #####

c       = variable(c)
f       = 0.25*(c-1)**2*(c+1)**2
dfdc    = diff(f, c)


##### DEFINE AND COMBINE VARIATIONAL FORMS #####

mu_mid = (1.0-theta)*mu0 + theta*mu

L0 = inner(c, q)*dx - c0*q*dx + Pechlet_number*dt*dot(grad(mu_mid), grad(q))*dx
L1 = inner(mu, v)*dx - inner(dfdc, v)*dx - Cahn_number**2*dot(grad(c), grad(v))*dx
L = L0 + L1


##### DEFINE BOUNDARY CONDITIONS #####
##### To include bcs to solver; add to line 146: solve(L == 0, u, BCs) #####

bc_left = DirichletBC(ME.sub(0), Constant(-1), left)
bc_right = DirichletBC(ME.sub(0), Constant(1), right)

BCs = [bc_left, bc_right]



_c = File('results/c.pvd')
_mu = File('results/mu.pvd')

while t < T + dt:
    t += dt
    solve(L == 0, u)
    u0.assign(u)

    _c << u0.split(deepcopy=True)[0]
    _mu << u0.split(deepcopy=True)[1]