from __future__ import print_function
from fenics import *
# from mshr import *
from dolfin import *
import matplotlib.pyplot as plt
import numpy as np
import time
from math import *

set_log_level(40)

def four_points(point, dist, rad):
    e = Expression('sqrt((x[1]-4)*(x[1]-4)+(x[0]-4)*(x[0]-4)+(x[2]-0.5)*(x[2]-0.5))'+
                    '<= Radius_Droplet ? 0 : '+
                    'sqrt((x[1]-5)*(x[1]-5)+(x[0]-4)*(x[0]-4)+(x[2]-0.5)*(x[2]-0.5))'+
                    '<= Radius_Droplet ? Phi_tot_int : '+
                    'sqrt((x[1]-3)*(x[1]-3)+(x[0]-4)*(x[0]-4)+(x[2]-0.5)*(x[2]-0.5))'+
                    '<= Radius_Droplet ? Phi_tot_int : '+
                    'sqrt((x[1]-4)*(x[1]-4)+(x[0]-5)*(x[0]-5)+(x[2]-0.5)*(x[2]-0.5))'+
                    '<= Radius_Droplet ? Phi_tot_int : '+
                    'sqrt((x[1]-4)*(x[1]-4)+(x[0]-3)*(x[0]-3)+(x[2]-0.5)*(x[2]-0.5))'+
                    '<= Radius_Droplet ? Phi_tot_int : Phi_tot_ext',
                    Phi_tot_int=Phi_tot_int, degree=2, tol=tol, Phi_tot_ext=Phi_tot_ext,
                    Radius_Droplet=Radius_Droplet)
    return e

dt = 0.1
Radius_Droplet = 0.25
mesh = Mesh('Meshes/multi_drop_gauss_med.xml')
V = FunctionSpace(mesh, 'CG', 1)
Phi_u = Function(V)
v = TestFunction(V)

Phi_u0 = Function(V)
Phi_tot = Function(V)

tol = 1E-14

Phi_tot_int = 0.99
Phi_tot_ext = 0.001

Phi_u0 = Expression('sqrt((x[1]-4)*(x[1]-4)+(x[0]-4)*(x[0]-4)+(x[2]-0.5)*(x[2]-0.5))'+
                    '<= Radius_Droplet ? 0 : '+
                    'sqrt((x[1]-5)*(x[1]-5)+(x[0]-4)*(x[0]-4)+(x[2]-0.5)*(x[2]-0.5))'+
                    '<= Radius_Droplet ? Phi_tot_int : '+
                    'sqrt((x[1]-3)*(x[1]-3)+(x[0]-4)*(x[0]-4)+(x[2]-0.5)*(x[2]-0.5))'+
                    '<= Radius_Droplet ? Phi_tot_int : '+
                    'sqrt((x[1]-4)*(x[1]-4)+(x[0]-5)*(x[0]-5)+(x[2]-0.5)*(x[2]-0.5))'+
                    '<= Radius_Droplet ? Phi_tot_int : '+
                    'sqrt((x[1]-4)*(x[1]-4)+(x[0]-3)*(x[0]-3)+(x[2]-0.5)*(x[2]-0.5))'+
                    '<= Radius_Droplet ? Phi_tot_int : Phi_tot_ext',
                    Phi_tot_int=Phi_tot_int, degree=2, tol=tol, Phi_tot_ext=Phi_tot_ext,
                    Radius_Droplet=Radius_Droplet)
Phi_u0 = interpolate(Phi_u0, V)

Phi_tot = Expression('sqrt((x[1]-4)*(x[1]-4)+(x[0]-4)*(x[0]-4)+(x[2]-0.5)*(x[2]-0.5))'+
                     ' <= Radius_Droplet ? Phi_tot_int : '+
                     'sqrt((x[1]-5)*(x[1]-5)+(x[0]-4)*(x[0]-4)+(x[2]-0.5)*(x[2]-0.5))'+
                     '<= Radius_Droplet ? Phi_tot_int: '+
                     'sqrt((x[1]-3)*(x[1]-3)+(x[0]-4)*(x[0]-4)+(x[2]-0.5)*(x[2]-0.5))'+
                     '<= Radius_Droplet ? Phi_tot_int: '+
                     'sqrt((x[1]-4)*(x[1]-4)+(x[0]-5)*(x[0]-5)+(x[2]-0.5)*(x[2]-0.5))'+
                     '<= Radius_Droplet ? Phi_tot_int: '+
                     'sqrt((x[1]-4)*(x[1]-4)+(x[0]-3)*(x[0]-3)+(x[2]-0.5)*(x[2]-0.5))'+
                     '<= Radius_Droplet ? Phi_tot_int: Phi_tot_ext', 
                     degree=2, tol=tol, Phi_tot_int=Phi_tot_int, Phi_tot_ext=Phi_tot_ext,
                     Radius_Droplet=Radius_Droplet)
Phi_tot = interpolate(Phi_tot,V)

# Gamma_int = 1/2.5/10;
# Gamma_ext = 1/10;
# Gamma0 = Expression('sqrt((x[1]-2)*(x[1]-2)+(x[0]-2)*(x[0]-2)+(x[2]-0.5)*(x[2]-0.5))'+
#                      ' <= Radius_Droplet ? Gamma_int : '+
#                      'sqrt((x[1]-2.5)*(x[1]-2.5)+(x[0]-2.5)*(x[0]-2.5)+(x[2]-0.5)*(x[2]-0.5))'+
#                      '<= Radius_Droplet ? Gamma_int: Gamma_ext', 
#                      degree=2, tol=tol, Gamma_int=Gamma_int, Gamma_ext=Gamma_ext,
#                      Radius_Droplet=Radius_Droplet)
# Gamma0 = interpolate(Gamma0,V)

Form = (inner((Phi_u-Phi_u0)/dt, v) - inner(((1-Phi_tot)/Phi_tot)*Phi_u*grad(Phi_tot)- 
                                            (1-Phi_tot)*grad(Phi_u), grad(v)))*dx

t = 0
cFile = XDMFFile('multi_drop_6_neighbors_med.xdmf')
cFile.write(Phi_u0, t)

# c_in_b = []
# c_in_b1 = []
c_in_6n_med = []
ti = time.time()
for i in range(3):
    c_in_6n_med.append([Phi_u([x, 4, 0.5]) for x in np.linspace(4, 4.49, 100)])
    solve(Form == 0, Phi_u)
    assign(Phi_u0, Phi_u)
    t += dt
    cFile.write(Phi_u0, t)
    print(time.time() - ti)

cFile.close()