# Customized FEM solver for FRAP theory
from dolfin import *
import fem_utils
import matplotlib.pyplot as plt
import numpy as np
import time

set_log_level(40)

def point_expr(p_c, p_list, phi_tot_int, phi_tot_ext, phi_init,
               rad_drop, tol):
    '''
    Create dolfin expression for initial condition of Phi or 
    overall profile of Phi_tot.
    p_c ... center point
    p_list ... list of points around center point
    '''
    s = ('sqrt(pow((x[0]-'+str(p_c[0])+'),2)+pow((x[1]-'+str(p_c[1])+'),2)+'+
         'pow((x[2]-'+str(p_c[2])+'),2))')
    s = s + '<= '+str(rad_drop)+' ? '+str(phi_init)+' : '
    for p in p_list:
        y = ('sqrt(pow((x[0]-'+str(p[0])+'),2)+pow((x[1]-'+str(p[1])+'),2)+'+
             'pow((x[2]-'+str(p[2])+'),2))')
        s = s + y + '<= '+str(rad_drop)+' ? '+str(phi_tot_int)+' : '
    s = s + str(phi_tot_ext)
    return Expression(s, degree=2, tol=tol)


class frap_solver:

    def __init__(self, cent_poin, mesh, tol=1E-14, dt=0.1, G_in=1,
                 G_out=1, name='FRAP', phi_tot_int=0.99, phi_tot_ext=0.01,
                 phi_tot_initial=0, point_list=[], rad_drop=0.25, T=100):
        '''
        Initialise solver
        '''
        self.cent_poin = cent_poin  # Center point
        self.mesh = Mesh(mesh)  # Create Mesh object from file
        self.name = name  # name of output file
        self.tol = tol  # Tolerance for expressions
        self.dt = dt  # Time step
        self.G_in = G_in  # Gamma inside, mobility
        self.G_out = G_out  # Gamma outside, mobility
        self.phi_tot_int = phi_tot_int  # Value inside droplet
        self.phi_tot_ext = phi_tot_ext  # Value outside droplet
        self.phi_tot_initial = phi_tot_initial # Initial value inside droplet
        self.point_list = point_list  # List of points around center droplet
        self.rad_drop = rad_drop  # Radius of all droplets
        self.T = T  # Number of time steps

        self.cs = []  # Intensity profiles over time


    def solve_frap(self):
        '''
        Solver routine 
        '''
        V = FunctionSpace(self.mesh, 'CG', 1)
        Phi_u = Function(V)
        v = TestFunction(V)

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

        Phi_u0 = point_expr(self.cent_poin, self.point_list, self.phi_tot_int,
                            self.phi_tot_ext, self.phi_tot_initial,
                            self.rad_drop, self.tol)
        Phi_u0 = interpolate(Phi_u0, V)

        Phi_tot = point_expr(self.cent_poin, self.point_list, self.phi_tot_int,
                             self.phi_tot_ext, self.phi_tot_int, self.rad_drop,
                             self.tol)
        Phi_tot = interpolate(Phi_tot, V)
        Gamma_0 = interpolate(Expression('(x[0]-xc)*(x[0]-xc)+'
                                         '(x[1]-yc)*(x[1]-yc)+'+
                                         '(x[2]-zc)*(x[2]-zc) <= rad_drop*rad_drop ? '+
                                         'G_in : G_out' ,
                              degree=2, rad_drop=self.rad_drop, G_in=self.G_in,
                              G_out=self.G_out, xc=self.cent_poin[0],
                              yc=self.cent_poin[1], zc=self.cent_poin[2]), V)

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

        t = 0
        cFile = XDMFFile(self.name+'.xdmf')
        cFile.write(Phi_u0, t)
        
        ti = time.time()
        for i in range(self.T):
            fem_utils.save_time_point(Phi_u0,self.name+'t_p_'+str(i)+'.h5')
            self.cs.append([Phi_u0([x, self.cent_poin[1], self.cent_poin[2]])
                            for x in np.linspace(4, 4.5, 100)])
            solve(Form == 0, Phi_u)
            assign(Phi_u0, Phi_u)
            t += self.dt
            cFile.write(Phi_u0, t)
            print(time.time() - ti)

        cFile.close()