using JFVM, PyPlot

function diffusion_newton()
L= 1.0 # domain length
Nx= 100 # number of cells
m= createMesh1D(Nx, L) # create a 1D mesh
D0= 1.0 # diffusion coefficient constant
# Define the diffusion coefficientand its derivative
D(phi)=D0*(1.0+phi.^2)
dD(phi)=2.0*D0*phi
# create boundary condition
BC = createBC(m)
BC.left.a[:]=0.0
BC.left.b[:]=1.0
BC.left.c[:]=5.0
BC.right.a[:]=0.0
BC.right.b[:]=1.0
BC.right.c[:]=0.0
(Mbc, RHSbc)= boundaryConditionTerm(BC)
# define the initial condition
phi0= 0.0 # initial value
phi_old= createCellVariable(m, phi0, BC) # create initial cells
phi_val= copyCell(phi_old)
# define the time steps
dt= 0.001*L*L/D0 # a proper time step for diffusion process
for i=1:10
    err=100
    (Mt, RHSt) = transientTerm(phi_old, dt)
    while err>1e-10
        # calculate the diffusion coefficient
        Dface= harmonicMean(cellEval(D,phi_val))
        # calculate the face value of phi_0
        phi_face= linearMean(phi_val)
        # calculate the velocity for convection term
        u= faceEval(dD, phi_face).*gradientTerm(phi_val)
        # diffusion term
        Mdif= diffusionTerm(Dface)
        # convection term
        Mconv= convectionTerm(u)
        # divergence term on the RHS
        RHSlin= divergenceTerm(u.*phi_face)
        # matrix of coefficients
        M= Mt-Mdif-Mconv+Mbc
        # RHS vector
        RHS= RHSbc+RHSt-RHSlin
        # call the linear solver
        phi_new= solveLinearPDE(m, M, RHS)
        # calculate the error
        err= maximum(abs(phi_new.value-phi_val.value))
        # assign the new phi to the old phi
        phi_val=copyCell(phi_new)
    end
    phi_old= copyCell(phi_val)
    visualizeCells(phi_val)
end
    return phi_val
end

@time phi_n=diffusion_newton()

function diffusion_direct()
L= 1.0 # domain length
Nx= 100 # number of cells
m= createMesh1D(Nx, L) # create a 1D mesh
D0= 1.0 # diffusion coefficient constant
# Define the diffusion coefficientand its derivative
D(phi)=D0*(1.0+phi.^2)
dD(phi)=2.0*D0*phi
# create boundary condition
BC = createBC(m)
BC.left.a[:]=0.0
BC.left.b[:]=1.0
BC.left.c[:]=5.0
BC.right.a[:]=0.0
BC.right.b[:]=1.0
BC.right.c[:]=0.0
(Mbc, RHSbc)= boundaryConditionTerm(BC)
# define the initial condition
phi0= 0.0 # initial value
phi_old= createCellVariable(m, phi0, BC) # create initial cells
phi_val= copyCell(phi_old)
# define the time steps
dt= 0.001*L*L/D0 # a proper time step for diffusion process
for i=1:10
    err=100
    (Mt, RHSt) = transientTerm(phi_old, dt)
    while err>1e-10
        # calculate the diffusion coefficient
        Dface= harmonicMean(cellEval(D,phi_val))
        # diffusion term
        Mdif= diffusionTerm(Dface)
        # matrix of coefficients
        M= Mt-Mdif+Mbc
        # RHS vector
        RHS= RHSbc+RHSt
        # call the linear solver
        phi_new= solveLinearPDE(m, M, RHS)
        # calculate the error
        err= maximum(abs(phi_new.value-phi_val.value))
        # assign the new phi to the old phi
        phi_val=copyCell(phi_new)
    end
    phi_old= copyCell(phi_val)
    visualizeCells(phi_val)
end
    return phi_val
end

@time phi_ds=diffusion_direct()

plot(phi_ds.domain.cellcenters.x, phi_ds.value[2:end-1], "y",
phi_n.domain.cellcenters.x, phi_n.value[2:end-1], "--r") 
legend(["direct subs", "Newton"])