# Extended Standard Cahn-Hilliard Example to Binary Flory Huggins 
# as discussed on 28/11/2019
# Example can be found at 
# https://bitbucket.org/fenics-project/dolfin/src/master/python/demo/
# documented/cahn-hilliard/demo_cahn-hilliard.py.rst#rst-header-id1
# Runs with fenics 2019.01
# The resulting .pvd file can be opened using default settings in ParaView

import matplotlib.pyplot as plt
#import mshr as ms
import numpy as np
from scipy.optimize import curve_fit
import time
import random
from dolfin import *
def create_expr(p_list):
    string = ''
    for p in p_list:
        string = (string+'(x[0]-'+str(p[0])+')*(x[0]-'+str(p[0])+')+(x[1]-'+str(p[1])+
                  ')*(x[1]-'+str(p[1])+')+(x[2]-'+str(p[2])+')*(x[2]-'+str(p[2])+
                  ')<=.05*.05 ? .76 :' )
    string = string + '.268'
    return string

# test = np.loadtxt('/Users/hubatsch/ownCloud/Dropbox_Lars_Christoph/frap_theory/points_clean.txt',  delimiter=',')
test = [[0.5, 0.5, 0.5]]
f = Expression((create_expr(test), 1), degree=1)

# Class for interfacing with the Newton solver
class CahnHilliardEquation(NonlinearProblem):
    def __init__(self, a, L):
        NonlinearProblem.__init__(self)
        self.L = L
        self.a = a
    def F(self, b, x):
        assemble(self.L, tensor=b)
    def J(self, A, x):
        assemble(self.a, tensor=A)

# Model parameters
kappa  = 0.25*10e-05  # surface parameter
dt     = 0.1e-02  # time step
X = 2.2 # Chi Flory Huggins parameter
# time stepping family, e.g.: 
# theta=1 -> backward Euler, theta=0.5 -> Crank-Nicolson
theta  = 0.5      
# Form compiler optionsmai
parameters["form_compiler"]["optimize"]     = True
parameters["form_compiler"]["cpp_optimize"] = True
# mesh = Mesh('Meshes/te.xdmf')
mesh = Mesh()
with XDMFFile("/home/hubatsch/frap_theory/Meshes/box_with_sphere.xdmf") as infile:
    infile.read(mesh)
# domain = ms.Sphere(Point(0, 0, 0), 1.0)
# mesh = ms.generate_mesh(domain, 100)
P1 = FiniteElement("Lagrange", mesh.ufl_cell(), 1)
ME = FunctionSpace(mesh, P1*P1)
# Define trial and test functions
du    = TrialFunction(ME)
q, v  = TestFunctions(ME)
# Define functions
u   = Function(ME)  # current solution
u0  = Function(ME)  # solution from previous converged step
tis = time.time()
# Split mixed functions
dc, dmu = split(du)
c,  mu  = split(u)
c0, mu0 = split(u0)
# Create intial conditions and interpolate
# u_init = InitialConditions(degree=1)
# u.interpolate(u_init)
# u0.interpolate(u_init)
u.interpolate(f)
u0.interpolate(f)
print(time.time()-tis)
# Compute the chemical potential df/dc
c = variable(c)
# mu_(n+theta)
mu_mid = (1.0-theta)*mu0 + theta*mu
# Weak statement of the equations
L0 = c*q*dx - c0*q*dx + dt*c*(1-c)*dot(grad(mu_mid), grad(q))*dx
L1 = mu*v*dx - (ln(c/(1-c))+X*(1-2*c))*v*dx - kappa*dot(grad(c), grad(v))*dx 
L = L0 + L1
# Compute directional derivative about u in the direction of du (Jacobian)
a = derivative(L, u, du)
# Create nonlinear problem and Newton solver
problem = CahnHilliardEquation(a, L)
solver = NewtonSolver()
# solver.parameters["linear_solver"] = "lu"
solver.parameters["linear_solver"] = 'gmres'
#solver.parameters["preconditioner"] = 'ilu'
#solver.parameters["convergence_criterion"] = "incremental"
#solver.parameters["relative_tolerance"] = 1e-6
# Output file
# file = File("output.pvd", "compressed")
file_c = XDMFFile('c_long.xdmf')
print(mesh.num_cells())
# Step in time
t = 0.0
T = 3*dt
ti = time.time()
while (t < T):
#     file << (u.split()[0], t)
    file_c.write(u.split()[0], t)
    t += dt
    u0.vector()[:] = u.vector()
    solver.solve(problem, u.vector())
    print(time.time() - ti)
file_c.close()