text:
using Polynomials, Roots, JLD, PyPlot
using JFVM
include("peteng-master/functions/rel_perms_real.jl")

struct RelativePermeabilityDME
    kro0::Function
    krw0::Real
    n_o::Real
    n_w::Real
    sor::Function
    dsor_dc::Function #used in DME_solver_gravity_pc_compressible
    swc::Real
end

struct CapillaryPressureDME
    pc::Function
    dpc_dS::Function
    d2pc_d2S::Function
    d2pc_dcdS::Function
    theta::Real #Contact angle
    lambda::Real
    b::Real
    pc0::Real #maximum allowed pressure at low water saturations
end

struct Geology
    m::MeshStructure
    perm_val::Real
    poros_val::Real
    k::CellValue
    ϕ::CellValue
end

struct SimulationTime
    t_final::Real # [s] final time for simulation
    dt::Real # [s] time step
end

struct BoundaryConditionDME
    BCc_oil::BoundaryCondition # concentration (DME in oil) boundary condition
    BCc_water::BoundaryCondition # concentration (DME in water) boundary condition
    BCp::BoundaryCondition # pressure boundary condition
    BCs::BoundaryCondition # saturation boundary condition
    BCt::BoundaryCondition # tracer concentration boundary condition
end

struct InitialConditionDME
    c_oil_init::CellValue
    c_water_init::CellValue
    p_init::CellValue
    sw_init::CellValue
    c_t_init::CellValue
end

struct FluidParametersDME
    ρ_oil::Function
    dρ_dc_oil::Function
    dρ_dp_oil::Function #used in DME_solver_gravity_pc_compressible
    ρ_water::Function
    dρ_dp_water::Function #used in DME_solver_gravity_pc_compressible
    μ_oil::Function
    dμ_dc_oil::Function
    μ_water::Function
    dμ_dc_water::Function
    K::Function
    dK_dc::Function #used in DME_solver_gravity_pc_compressible
    a::Real
    D_oil::Real
    D_water::Real
end

function DME_solver_no_gravity_no_pc_incompressible(G::Geology, K::RelativePermeabilityDME, MT::FluidParametersDME, T::Array{SimulationTime,1}, BC::Array{BoundaryConditionDME,1}, IC::InitialConditionDME, k_o::Real, k_w::Real, plotFigures::Bool)

    # =============================================================================
    # Unpacking the relative permeabillity parameters
    # =============================================================================
    krw0 = K.krw0
    kro0_c = K.kro0 #function
    n_w  = K.n_w
    n_o  = K.n_o
    swc  = K.swc
    sor_c = K.sor #function

    KRW  = (sw, sor) -> krw.(sw, krw0, sor, swc, n_w)
    dKRW = (sw, sor) -> dkrwdsw.(sw, krw0, sor, swc, n_w)
    KRO  = (sw, kro0, sor) -> kro.(sw, kro0, sor, swc, n_o)
    dKRO = (sw, kro0, sor) -> dkrodsw.(sw, kro0, sor, swc, n_o)

    # =============================================================================
    # Unpacking the boundary and initial conditions
    # =============================================================================

    BCc_water   = BC[1].BCc_water # concentration (DME in water) boundary condition
    BCc_oil     = BC[1].BCc_oil # concentration (DME in oil) boundary condition
    BCp         = BC[1].BCp # pressure boundary condition
    BCs         = BC[1].BCs # saturation boundary condition
    BCt         = BC[1].BCt # tracer concentration boundary condition

    # discretize boundary conditions
    M_bc_c_water, RHS_bc_c_water    = boundaryConditionTerm(BCc_water)
    M_bc_c_oil, RHS_bc_c_oil        = boundaryConditionTerm(BCc_oil)
    M_bc_p, RHS_bc_p                = boundaryConditionTerm(BCp)
    M_bc_s, RHS_bc_s                = boundaryConditionTerm(BCs)
    M_bc_t, RHS_bc_t                = boundaryConditionTerm(BCt)

    # Iinitial conditions
    c_water_init    = IC.c_water_init
    c_oil_init      = IC.c_oil_init
    p_init          = IC.p_init
    sw_init         = IC.sw_init
    c_t_init        = IC.c_t_init

    # new values of each variable
    c_water_val     = copyCell(c_water_init)
    c_oil_val       = copyCell(c_oil_init)
    p_val           = copyCell(p_init)
    c_t_val         = copyCell(c_t_init)
    sw_val          = copyCell(sw_init)
    c_water_face    = arithmeticMean(c_water_val)
    c_oil_face      = arithmeticMean(c_oil_val)
    sor_face        = faceEval(sor_c, c_oil_face)

    # =============================================================================
    # Unpacking the fluid parameters
    # =============================================================================
    D_water = MT.D_water# m^2/s
    D_oil   = MT.D_oil # m^2/s

    ρ_water         = MT.ρ_water
    ρ_oil           = MT.ρ_oil
    dρ_dc_oil       = MT.dρ_dc_oil
    μ_water         = MT.μ_water
    dμ_dc_water     = MT.dμ_dc_water
    μ_oil           = MT.μ_oil
    dμ_dc_oil       = MT.dμ_dc_oil

    Kval = MT.K #Equilibrium constant

    #surface area only active if DME flooding is executed, otherwise 0
    DME_injected = false
    DME_injection_ended = false
    if maximum(BCc_oil.left.c[:]) > 0
        a = MT.a # m^2/m^3
        DME_injected = true
    else
        a = 0
    end

    K_w = (Kval(0)*k_o*k_w)/(Kval(0)*k_o+k_w)
    K_o = (k_o*k_w)/(Kval(0)*k_o+k_w)

    # =============================================================================
    # Unpacking the geology/mesh
    # =============================================================================
    perm_val    = G.perm_val      # [m^2] permeability
    poros_val   = G.poros_val     # [-] porosity
    k           = G.k
    k_face      = harmonicMean(k)     # permeability on the cell faces
    poros_val   = G.poros_val
    ϕ           = G.ϕ

    m   = G.m
    Nx  = m.dims[1]  # number of grids in the x direction
    Lx  = sum(m.cellsize.x[2:end-1]) # [m]
    N   = (Nx+2)

    if m.dimension>=2
        Ny  = m.dims[2]  # number of grids in the y direction
        Ly  = sum(m.cellsize.y[2:end-1]) # [m]
        N   = (Nx+2)*(Ny+2)
    end

    if m.dimension>=3
        Nz  = m.dims[3]  # number of grids in the y direction
        Lz  = sum(m.cellsize.z[2:end-1]) # [m]
        N   = (Nx+2)*(Ny+2)*(Ny+2)
    end

    # =============================================================================
    # Unpacking the time related parameters
    # =============================================================================

    t_final1  = T[1].t_final
    dt0       = T[1].dt
    dt        = dt0

    t_final_end  = T[end].t_final
    no_of_switch = length(T)

    # =============================================================================
    # Preparing for time-stepping
    # =============================================================================

    # outside the loop: initialization
    uw       = gradientTerm(p_val) # only for initialization of the water velocity vector
    uo       = gradientTerm(p_val) # only for initialization of the oil velocity vector
    ut       = 0 # initialize total velocity

    # Diffusion terms in the DME equations will be constant throughout the simulation
    M_c_water_diff  = diffusionTerm(createFaceVariable(m, poros_val*D_water))
    M_c_oil_diff    = diffusionTerm(createFaceVariable(m, poros_val*D_oil))

    # this whole thing goes inside two loops (Newton and time)
    tol_s = 1e-7
    tol_c = 1e-7
    max_change_c = 0.1 # 10 % relative change
    max_change_s = 0.1 # 10 % relative change
    max_int_loop = 40
    t = 0.0
    oil_init=domainInt(cellEval(ρ_oil, c_oil_val).*(1-sw_init)) # initial oil volume in the core
    rec_fact=zeros(1)
    dp_hist = zeros(1)
    t_s=zeros(1)
    switch_count = 2
    FL = fluxLimiter("SUPERBEE") # flux limiter

    while t < t_final_end
        # Switching injected fluid
        if switch_count <= no_of_switch && t > T[switch_count-1].t_final
            dt0 = T[switch_count].dt

            BCc_water   = BC[switch_count].BCc_water
            BCc_oil     = BC[switch_count].BCc_oil
            BCp         = BC[switch_count].BCp
            BCt         = BC[switch_count].BCt

            c_water_init    = createCellVariable(m, internalCells(c_water_val), BCc_water)
            c_oil_init      = createCellVariable(m, internalCells(c_oil_val), BCc_oil)
            p_val           = createCellVariable(m, internalCells(p_val), BCp)
            c_t_init        = createCellVariable(m, internalCells(c_t_val), BCt)

            M_bc_c_water, RHS_bc_c_water    = boundaryConditionTerm(BCc_water)
            M_bc_c_oil, RHS_bc_c_oil        = boundaryConditionTerm(BCc_oil)
            M_bc_p, RHS_bc_p                = boundaryConditionTerm(BCp)
            M_bc_t, RHS_bc_t                = boundaryConditionTerm(BCt)

            #changing surface area if DME flooding is beginning
            if maximum(BCc_oil.left.c[:]) > 0
                a = MT.a # m^2/m^3
                DME_injected = true
            else
                if DME_injected
                    DME_injection_ended = true
                end
            end

            switch_count += 1
        end

        error_s = 2*tol_s
        error_c = 2*tol_c
        loop_countor = 0

        #Newton-loop
        while error_s>tol_s || error_c_oil>tol_c || error_c_water>tol_c
            loop_countor += 1
            if loop_countor > max_int_loop
              c_water_val   = copyCell(c_water_init)
              c_oil_val     = copyCell(c_oil_init)
              p_val         = copyCell(p_init)
              sw_val        = copyCell(sw_init)
              c_t_val       = copyCell(c_t_init)
              dt = dt/3.0
              break
            end

            ∇p0 = gradientTerm(p_val)

            c_water_face  = upwindMean(c_water_val, uw)
            c_oil_face    = upwindMean(c_oil_val, uo)
            sw_face       = upwindMean(sw_val, uw)
            sor_face      = faceEval(sor_c, c_oil_face)
            sor_cell      = cellEval(sor_c, c_oil_val)
            kro0_face     = faceEval(kro0_c, sor_face)
            kro0_cell     = cellEval(kro0_c, sor_cell)

            ρ_water_cell        = cellEval(ρ_water, c_water_val)
            ρ_water_face        = faceEval(ρ_water, c_water_face)
            ρ_oil_cell          = cellEval(ρ_oil, c_oil_val)
            ρ_oil_face          = faceEval(ρ_oil, c_oil_face)
            dρ_dc_oil_cell      = cellEval(dρ_dc_oil, c_oil_val)
            dρ_dc_oil_face      = faceEval(dρ_dc_oil, c_oil_face)
            μ_water_face        = faceEval(μ_water, c_water_face)
            dμ_dc_water_face    = faceEval(dμ_dc_water, c_water_face)
            μ_oil_face          = faceEval(μ_oil, c_oil_face)
            dμ_dc_oil_face      = faceEval(dμ_dc_oil,c_oil_face)

            dc_μ_dc_water   = (μ_water_face-c_water_face.*dμ_dc_water_face)./(μ_water_face.*μ_water_face)
            d1_μ_dc_water   = (-dμ_dc_water_face)./(μ_water_face.*μ_water_face)
            dc_μ_dc_oil     = (μ_oil_face-c_oil_face.*dμ_dc_oil_face)./((μ_oil_face).*(μ_oil_face))
            dρ_μ_dc_oil    = (dρ_dc_oil_face.*μ_oil_face-dμ_dc_oil_face.*ρ_oil_face)./(μ_oil_face.*μ_oil_face)

            krw_face        = faceEval(KRW, sw_face, sor_face)
            kro_face        = faceEval(KRO, sw_face, kro0_face, sor_face)
            dkrw_dc_face    = faceEval(dKRW, sw_face, sor_face)
            dkro_dc_face    = faceEval(dKRO, sw_face, kro0_face, sor_face)
            uw              = -k_face.*krw_face./μ_water_face.*∇p0
            uo              = -k_face.*kro_face./μ_oil_face.*∇p0
            ut              = uw + uo

            #System matrices
            M_c_water_1     = M_bc_c_water  - M_c_water_diff  + linearSourceTerm(createCellVariable(m, K_w.*a.*poros_val))
            M_c_oil_1       =                                   linearSourceTerm(createCellVariable(m, -K_o.*a.*poros_val))
            M_c_water_2     =                                   linearSourceTerm(createCellVariable(m, -K_w.*a.*poros_val))
            M_c_oil_2       = M_bc_c_oil    - M_c_oil_diff    + linearSourceTerm(createCellVariable(m, K_o.*a.*poros_val))
            M_p_3           = M_bc_p
            M_S_4           = M_bc_s

            #Jacobians related to coupled reactive transport of water (row 1)
            J_t_c_water_1, RHS_t_c_water_1  = transientTerm(c_water_init, dt, ϕ.*sw_val)
            J_adv_c_water_1                 = convectionUpwindTerm(-k_face.*dc_μ_dc_water.*krw_face.*∇p0, ut)
            J_c_water_1                     = J_t_c_water_1 + J_adv_c_water_1

            J_p_1 = diffusionTerm(-k_face.*(1.0./μ_water_face).*krw_face.*c_water_face)

            J_t_S_1, RHS_t_S_1  = transientTerm(sw_init, dt, ϕ.*c_water_val)
            J_adv_S_1           = convectionUpwindTerm(-k_face.*(c_water_face./μ_water_face).*dkrw_dc_face.*∇p0, ut)
            J_S_1               = J_t_S_1 + J_adv_S_1

            #Jacobians related to coupled reactive transport of oil (row 2)
            J_t_c_oil_2, RHS_t_c_oil_2  = transientTerm(c_oil_init, dt, ϕ.*(1.0-sw_val))
            J_adv_c_oil_2               = convectionUpwindTerm(-k_face.*dc_μ_dc_oil.*kro_face.*∇p0, ut)
            J_c_oil_2                   = J_t_c_oil_2 + J_adv_c_oil_2

            J_p_2 = diffusionTerm(-k_face.*(1.0./μ_oil_face).*kro_face.*c_oil_face)

            J_t_S_2, RHS_t_S_2  = transientTerm(sw_init, dt, -ϕ.*c_oil_val)
            J_adv_S_2           = convectionUpwindTerm(-k_face.*(c_oil_face./μ_oil_face).*dkro_dc_face.*∇p0, ut)
            J_S_2               = J_t_S_2 + J_adv_S_2

            #Jacobians related to two-phase flow of water (row 3)
            J_adv_c_water_3                 = convectionUpwindTerm(-k_face.*d1_μ_dc_water.*krw_face.*ρ_water_face.*∇p0, ut)
            J_c_water_3                     = J_adv_c_water_3

            J_p_3 = diffusionTerm(-k_face.*(1.0./μ_water_face).*krw_face.*ρ_water_face)

            J_t_S_3, RHS_t_S_3  = transientTerm(sw_init, dt, ϕ.*ρ_water_cell)
            J_adv_S_3           = convectionUpwindTerm(-k_face.*(1.0./μ_water_face).*dkrw_dc_face.*ρ_water_face.*∇p0, ut)
            J_S_3               = J_t_S_3 + J_adv_S_3

            #Jacobians related to two-phase flow of oil (row 4)
            J_t_c_oil_4, RHS_t_c_oil_4  = transientTerm(c_oil_init, dt, ϕ.*(1.0-sw_val).*dρ_dc_oil_cell)
            J_adv_c_oil_4               = convectionUpwindTerm(-k_face.*dρ_μ_dc_oil.*kro_face.*∇p0, ut)
            J_c_oil_4                   = J_t_c_oil_4 + J_adv_c_oil_4

            J_p_4 = diffusionTerm(-k_face.*(ρ_oil_face./μ_oil_face).*kro_face)

            J_t_S_4, RHS_t_S_4  = transientTerm(sw_init, dt, -ϕ.*ρ_oil_cell)
            J_adv_S_4           = convectionUpwindTerm(-k_face.*(ρ_oil_face./μ_oil_face).*dkro_dc_face.*∇p0, ut)
            J_S_4               = J_t_S_4 + J_adv_S_4

            # create the PDE system M [c_water,c_oil,p,S]=RHS
            # x_val = [c_water_val.value[:], c_oil_val.value[:], p_val.value[:]; sw_val.value[:]]
            M = [M_c_water_1+J_c_water_1 M_c_oil_1 J_p_1 J_S_1;
                 M_c_water_2 M_c_oil_2+J_c_oil_2 J_p_2 J_S_2;
                 J_c_water_3 zeros(N,N) M_p_3+J_p_3 J_S_3;
                 zeros(N,N) J_c_oil_4 J_p_4 M_S_4+J_S_4]

            RHS_1 = RHS_bc_c_water  + RHS_t_c_water_1   + RHS_t_S_1     + J_adv_c_water_1*c_water_val.value[:]  + J_adv_S_1*sw_val.value[:]
            RHS_2 = RHS_bc_c_oil    + RHS_t_c_oil_2     + RHS_t_S_2     + J_adv_c_oil_2*c_oil_val.value[:]      + J_adv_S_2*sw_val.value[:]
            RHS_3 = RHS_bc_p                            + RHS_t_S_3     + J_adv_c_water_3*c_water_val.value[:]  + J_adv_S_3*sw_val.value[:]
            RHS_4 = RHS_bc_s        + RHS_t_c_oil_4     + RHS_t_S_4     + J_adv_c_oil_4*c_oil_val.value[:]      + J_adv_S_4*sw_val.value[:]

            RHS = [RHS_1[:];
                   RHS_2[:];
                   RHS_3[:];
                   RHS_4[:]]

            x_sol = M\RHS

            c_water_new = x_sol[1:N]
            c_oil_new   = x_sol[(1:N)+N]
            p_new       = x_sol[(1:N)+(2*N)]
            s_new       = x_sol[(1:N)+(3*N)]

            error_c_water = sumabs(c_water_new[2:end-1]-c_water_val.value[2:end-1])
            error_c_oil = sumabs(c_oil_new[2:end-1]-c_oil_val.value[2:end-1])
            error_s = sumabs(s_new[2:end-1]-sw_val.value[2:end-1])

            c_water_val.value[:]    = c_water_new[:]
            c_oil_val.value[:]      = c_oil_new[:]
            p_val.value[:]          = p_new[:]

            #Appleyard approach to updating saturations
            if !DME_injection_ended
                dsw_apple = 0.2
                eps_apple = sqrt(eps()) # 1e-5
                for i in eachindex(s_new)
                    if s_new[i]>=(1-sor_cell.value[i])
                        if sw_val.value[i]<(1-sor_cell.value[i]-eps_apple)
                            sw_val.value[i] = 1-sor_cell.value[i]-eps_apple
                        else
                            sw_val.value[i] = 1-sor_cell.value[i]
                        end
                    elseif s_new[i]<=swc
                        if sw_val.value[i]> swc+eps_apple
                            sw_val.value[i] = swc+eps_apple
                        else
                            sw_val.value[i] = swc
                        end
                    elseif abs(s_new[i]-sw_val.value[i])>dsw_apple
                        sw_val.value[i] += dsw_apple*sign(s_new[i]-sw_val.value[i])
                    else
                        sw_val.value[i] = s_new[i]
                    end
                end
            else
                w_c_w = 0.6
                w_c_o = 0.6
                w_p = 0.9
                w_sw = 0.6
                c_water_val.value[:] = w_c_w*c_water_new[:]+(1-w_c_w)*c_water_val.value[:]
                c_oil_val.value[:] = w_c_o*c_oil_new[:]+(1-w_c_o)*c_oil_val.value[:]
                p_val.value[:] = w_p*p_new[:]+(1-w_p)*p_val.value[:]
                sw_val.value[:] = w_sw*s_new[:]+(1-w_sw)*sw_val.value[:]
            end

            sw_val = createCellVariable(m, internalCells(sw_val), BCs)
        end

        if loop_countor < max_int_loop
            for tracer_count = 1:4
                # tracer calculation
                M_t_t, RHS_t_t = transientTerm(c_t_init, dt, sw_val.*ϕ)
                M_s_t          = linearSourceTerm(ϕ.*(sw_val-sw_init)/dt)
                M_a_t, RHS_a_t = convectionTvdTerm(uw, c_t_val, FL, ut)
                c_t_val = solveLinearPDE(m, M_t_t+M_s_t+M_a_t+M_bc_t, RHS_a_t+RHS_bc_t+RHS_t_t)
            end
            c_water_init    = copyCell(c_water_val)
            c_oil_init      = copyCell(c_oil_val)
            p_init          = copyCell(p_val)
            sw_init         = copyCell(sw_val)
            c_t_init        = copyCell(c_t_val)

            sor_face        = faceEval(sor_c, c_oil_face)
            kro0_face       = faceEval(kro0_c, sor_face)
            t += dt
            dt = dt0
            #println("progress: $(t/t_final_end*100) [%]")
            rec_fact=push!(rec_fact, (oil_init-domainInt(cellEval(ρ_oil,c_oil_val).*(1-sw_val)))/oil_init)
            t_s=push!(t_s, t)
        end
    end

    #Rounding off solutions, to avoid Julia plotting issues
    c_water_init.value[:] = round(c_water_init.value[:],2)
    c_oil_init.value[:] = round(c_oil_init.value[:],2)
    c_t_init.value[:] = round(c_t_init.value[:],2)

    if m.dimension >= 1 && m.dimension < 2 && plotFigures
        plot_DME_1D(k, sor_c, kro0_c, KRO, KRW, BC_DME, c_water_init, c_oil_init, p_init, sw_init, c_t_init, t_s, rec_fact)
    elseif m.dimension >= 2 && m.dimension < 3 && plotFigures
        plot_DME_2D(k, sor_c, kro0_c, KRO, KRW, BC_DME, c_water_init, c_oil_init, p_init, sw_init, c_t_init, t_s, rec_fact)
    elseif plotFigures
        plot_DME_3D(k, KRO, KRW, BC_DME, c_water_init, c_oil_init, p_init, sw_init, c_t_init, t_s, rec_fact)
    end

    println("Finished Simulation")
    return t_s, rec_fact
end

function DME_solver_gravity_pc_compressible(G::Geology, K::RelativePermeabilityDME, PC::CapillaryPressureDME, MT::FluidParametersDME, T::Array{SimulationTime,1}, BC::Array{BoundaryConditionDME,1}, IC::InitialConditionDME, k_o::Real, k_w::Real, plotFigures::Bool)
    # =============================================================================
    # Unpacking the relative permeabillity parameters
    # =============================================================================
    krw0    = K.krw0
    kro0    = K.kro0
    n_w     = K.n_w
    n_o     = K.n_o
    swc     = K.swc
    sor_c   = K.sor
    dsor_c  = K.dsor_dc

    KRW             = (sw, sor_c) -> krw.(sw, krw0, sor_c, swc, n_w)
    dKRW_dc_oil     = (sw, sor_c, dsor_c) -> dkrw_dc_oil.(sw, krw0, sor_c, dsor_c, swc, n_w)
    dKRW_dS         = (sw, sor_c) -> dkrwdsw.(sw, krw0, sor_c, swc, n_w)
    KRO             = (sw, sor_c) -> kro.(sw, kro0, sor_c, swc, n_o)
    dKRO_dc_oil     = (sw, sor_c, dsor_c) -> dkro_dc_oil.(sw, kro0, sor_c, dsor_c, swc, n_o)
    dKRO_dS         = (sw, sor_c) -> dkrodsw.(sw, kro0, sor_c, swc, n_o)

    # =============================================================================
    # Unpacking the capillary pressure functions
    # =============================================================================

    pc = PC.pc
    dpc_dS = PC.dpc_dS
    d2pc_d2S = PC.d2pc_d2S
    d2pc_dcdS = PC.d2pc_dcdS

    # =============================================================================
    # Unpacking the boundary and initial conditions
    # =============================================================================

    BCc_water   = BC[1].BCc_water # concentration (DME in water) boundary condition
    BCc_oil     = BC[1].BCc_oil # concentration (DME in oil) boundary condition
    BCp         = BC[1].BCp # pressure boundary condition
    BCs         = BC[1].BCs # saturation boundary condition
    BCt         = BC[1].BCt # tracer concentration boundary condition

    # discretize boundary conditions
    M_bc_c_water, RHS_bc_c_water    = boundaryConditionTerm(BCc_water)
    M_bc_c_oil, RHS_bc_c_oil        = boundaryConditionTerm(BCc_oil)
    M_bc_p, RHS_bc_p                = boundaryConditionTerm(BCp)
    M_bc_s, RHS_bc_s                = boundaryConditionTerm(BCs)
    M_bc_t, RHS_bc_t                = boundaryConditionTerm(BCt)

    # Iinitial conditions
    c_water_init    = IC.c_water_init
    c_oil_init      = IC.c_oil_init
    p_init          = IC.p_init
    sw_init         = IC.sw_init
    c_t_init        = IC.c_t_init

    # new values of each variable
    c_water_val     = copyCell(c_water_init)
    c_oil_val       = copyCell(c_oil_init)
    p_val           = copyCell(p_init)
    c_t_val         = copyCell(c_t_init)
    sw_val          = copyCell(sw_init)
    c_water_face    = arithmeticMean(c_water_val)
    c_oil_face      = arithmeticMean(c_oil_val)
    sor_face        = faceEval(sor_c, c_oil_face)

    # =============================================================================
    # Unpacking the fluid parameters
    # =============================================================================
    D_water = MT.D_water# m^2/s
    D_oil   = MT.D_oil # m^2/s

    ρ_water         = MT.ρ_water
    dρ_dp_water     = MT.dρ_dp_water
    ρ_oil           = MT.ρ_oil
    dρ_dc_oil       = MT.dρ_dc_oil
    dρ_dp_oil       = MT.dρ_dp_oil
    μ_water         = MT.μ_water
    dμ_dc_water     = MT.dμ_dc_water
    μ_oil           = MT.μ_oil
    dμ_dc_oil       = MT.dμ_dc_oil

    Kval = MT.K #Equilibrium constant
    dKval_dc = MT.dK_dc

    #surface area only active if DME flooding is executed, otherwise 0
    DME_injected = false
    DME_injection_ended = false
    if maximum(BCc_oil.left.c[:]) > 0
        a = MT.a # m^2/m^3
        DME_injected = true
    else
        a = 0
    end

    K_w = c -> (Kval(c).*k_o.*k_w)./(Kval(c).*k_o+k_w)
    K_o = c -> (k_o.*k_w)./(Kval(c).*k_o+k_w)

    # =============================================================================
    # Unpacking the geology/mesh
    # =============================================================================
    perm_val    = G.perm_val      # [m^2] permeability
    poros_val   = G.poros_val     # [-] porosity
    k           = G.k
    k_face      = harmonicMean(k)     # permeability on the cell faces
    poros_val   = G.poros_val
    ϕ           = G.ϕ

    m       = G.m
    Nx      = m.dims[1]  # number of grids in the x direction
    Lx      = sum(m.cellsize.x[2:end-1]) # [m]
    N       = (Nx+2)
    g       = createFaceVariable(m, -9.82) #gravity in x-direction

    if m.dimension>=2
        Ny      = m.dims[2]  # number of grids in the y direction
        Ly      = sum(m.cellsize.y[2:end-1]) # [m]
        N       = (Nx+2)*(Ny+2)
        g       = createFaceVariable(m, [0.0, -9.82]) #gravity in y-direction
    end

    if m.dimension>=3
        Nz      = m.dims[3]  # number of grids in the y direction
        Lz      = sum(m.cellsize.z[2:end-1]) # [m]
        N       = (Nx+2)*(Ny+2)*(Ny+2)
        g       = createFaceVariable(m, [0.0, 0.0, -9.82]) #gravity in z-direction
    end

    # =============================================================================
    # Unpacking the time related parameters
    # =============================================================================

    t_final1  = T[1].t_final
    dt0       = T[1].dt
    dt        = dt0

    t_final_end  = T[end].t_final
    no_of_switch = length(T)

    # =============================================================================
    # Preparing for time-stepping
    # =============================================================================

    # outside the loop: initialization
    uw       = gradientTerm(p_val) # only for initialization of the water velocity vector
    uo       = gradientTerm(p_val) # only for initialization of the oil velocity vector
    ut       = 0 # initialize total velocity

    # Diffusion terms in the DME equations will be constant throughout the simulation
    M_c_water_diff  = diffusionTerm(createFaceVariable(m, poros_val*D_water))
    M_c_oil_diff    = diffusionTerm(createFaceVariable(m, poros_val*D_oil))

    # this whole thing goes inside two loops (Newton and time)
    tol_s = 1e-7
    tol_c = 1e-7
    max_change_c = 0.1 # 10 % relative change
    max_change_s = 0.1 # 10 % relative change
    max_int_loop = 15
    t = 0.0
    oil_init=domainInt(cellEval(ρ_oil, c_oil_val, p_val).*(1-sw_init)) # initial oil volume in the core
    rec_fact=zeros(1)
    dp_hist = zeros(1)
    t_s=zeros(1)
    switch_count = 2
    FL = fluxLimiter("SUPERBEE") # flux limiter

    while t < t_final_end
        # Switching injected fluid
        if switch_count <= no_of_switch && t > T[switch_count-1].t_final
            dt0 = T[switch_count].dt

            BCc_water   = BC[switch_count].BCc_water
            BCc_oil     = BC[switch_count].BCc_oil
            BCp         = BC[switch_count].BCp
            BCt         = BC[switch_count].BCt

            c_water_init    = createCellVariable(m, internalCells(c_water_val), BCc_water)
            c_oil_init      = createCellVariable(m, internalCells(c_oil_val), BCc_oil)
            p_val           = createCellVariable(m, internalCells(p_val), BCp)
            c_t_init        = createCellVariable(m, internalCells(c_t_val), BCt)

            M_bc_c_water, RHS_bc_c_water    = boundaryConditionTerm(BCc_water)
            M_bc_c_oil, RHS_bc_c_oil        = boundaryConditionTerm(BCc_oil)
            M_bc_p, RHS_bc_p                = boundaryConditionTerm(BCp)
            M_bc_t, RHS_bc_t                = boundaryConditionTerm(BCt)

            #changing surface area if DME flooding is executed, otherwise 0
            if maximum(BCc_oil.left.c[:]) > 0
                a = MT.a # m^2/m^3
                DME_injected = true
            else
                if DME_injected
                    DME_injection_ended = true
                end
            end

            switch_count += 1
        end

        error_s = 2*tol_s
        error_c = 2*tol_c
        loop_countor = 0

        #System matrices
        M_c_water_1     = M_bc_c_water  - M_c_water_diff
        M_c_oil_2       = M_bc_c_oil    - M_c_oil_diff
        M_p_3           = M_bc_p
        M_S_4           = M_bc_s

        #Newton-loop
        while error_s>tol_s || error_c_oil>tol_c || error_c_water>tol_c
            loop_countor += 1
            if loop_countor > max_int_loop
              c_water_val   = copyCell(c_water_init)
              c_oil_val     = copyCell(c_oil_init)
              p_val         = copyCell(p_init)
              sw_val        = copyCell(sw_init)
              c_t_val       = copyCell(c_t_init)
              dt = dt/3.0
              break
            end

            ∇p0 = gradientTerm(p_val)
            ∇S0 = gradientTerm(sw_val)

            c_water_face  = upwindMean(c_water_val, uw)
            c_oil_face    = upwindMean(c_oil_val, uo)
            p_face        = upwindMean(p_val, uw) #CORRECT TO USE uw?
            sw_face       = upwindMean(sw_val, uw) #CORRECT TO USE uw?
            sor_cell      = cellEval(sor_c, c_oil_val)
            sor_face      = faceEval(sor_c, c_oil_face)
            dsor_cell     = cellEval(dsor_c, c_oil_val)
            dsor_face     = faceEval(dsor_c, c_oil_face)

            ρ_water_cell        = cellEval(ρ_water, p_val)
            ρ_water_face        = faceEval(ρ_water, p_face)
            dρ_dp_water_cell    = cellEval(dρ_dp_water, p_val)
            dρ_dp_water_face    = faceEval(dρ_dp_water, p_face)
            ρ_oil_cell          = cellEval(ρ_oil, c_oil_val, p_val)
            ρ_oil_face          = faceEval(ρ_oil, c_oil_face, p_face)
            dρ_dc_oil_cell      = cellEval(dρ_dc_oil, c_oil_val, p_val)
            dρ_dc_oil_face      = faceEval(dρ_dc_oil, c_oil_face, p_face)
            dρ_dp_oil_cell      = cellEval(dρ_dp_oil, c_oil_val, p_val)
            dρ_dp_oil_face      = faceEval(dρ_dp_oil, c_oil_face, p_face)
            μ_water_cell        = cellEval(μ_water, c_water_val)
            μ_water_face        = faceEval(μ_water, c_water_face)
            dμ_dc_water_face    = faceEval(dμ_dc_water, c_water_face)
            μ_oil_cell          = cellEval(μ_oil, c_oil_val)
            μ_oil_face          = faceEval(μ_oil, c_oil_face)
            dμ_dc_oil_face      = faceEval(dμ_dc_oil,c_oil_face)

            krw_cell            = cellEval(KRW, sw_val, sor_cell)
            krw_face            = faceEval(KRW, sw_face, sor_face)
            kro_cell            = cellEval(KRO, sw_val, sor_cell)
            kro_face            = faceEval(KRO, sw_face, sor_face)
            dkrw_dc_face        = faceEval(dKRW_dc_oil, sw_face, sor_face, dsor_face)
            dkro_dc_face        = faceEval(dKRO_dc_oil, sw_face, sor_face, dsor_face)
            dkrw_dS_face        = faceEval(dKRW_dS, sw_face, sor_face)
            dkro_dS_face        = faceEval(dKRO_dS, sw_face, sor_face)

            #dpc_dS_face         = faceEval(dpc_dS, sw_face, sor_face)
            #d2pc_d2S_face       = faceEval(d2pc_d2S, sw_face, sor_face)
            #d2pc_dcdS_face      = faceEval(d2pc_dcdS, sw_face, sor_face, dsor_face)

            Kval_cell     = cellEval(Kval, c_oil_val)
            Kval_face     = faceEval(Kval, c_oil_face)
            dKval_dc_cell = cellEval(dKval_dc, c_oil_val)
            dKval_dc_face = faceEval(dKval_dc, c_oil_face)
            K_w_cell      = cellEval(K_w, c_oil_val)
            K_w_face      = faceEval(K_w, c_oil_face)
            K_o_cell      = cellEval(K_o, c_oil_val)
            K_o_face      = faceEval(K_o, c_oil_face)

            #Derivatives of c_water
            dc_μ_dc_water = (μ_water_face-c_water_face.*dμ_dc_water_face)./(μ_water_face.*μ_water_face)
            d1_μ_dc_water = -dμ_dc_water_face./(μ_water_face.*μ_water_face)

            #Derivatives of c_oil
            dc_μ_kro_dc_oil     = kro_face./μ_oil_face-c_oil_face.*kro_face./(μ_oil_face.*μ_oil_face).*dμ_dc_oil_face+c_oil_face./μ_oil_face.*dkro_dc_face
            dρ_μ_kro_dc_oil     = kro_face./μ_oil_face.*dρ_dc_oil_face-ρ_oil_face.*kro_face./(μ_oil_face.*μ_oil_face).*dμ_dc_oil_face+ρ_oil_face./μ_oil_face.*dkro_dc_face
            dρ2_μ_kro_dc_oil    = 2.0.*ρ_oil_face.*kro_face./μ_oil_face.*dρ_dc_oil_face-(ρ_oil_face.*ρ_oil_face).*kro_face./(μ_oil_face.*μ_oil_face).*dμ_dc_oil_face+(ρ_oil_face.*ρ_oil_face)./μ_oil_face.*dkro_dc_face
            dρ_μ_c_kro_dc_oil   = c_oil_face.*kro_face./μ_oil_face.*dρ_dc_oil_face+ρ_oil_face.*kro_face./μ_oil_face-ρ_oil_face.*c_oil_face.*kro_face./(μ_oil_face.*μ_oil_face).*dμ_dc_oil_face+ρ_oil_face.*c_oil_face./μ_oil_face.*dkro_dc_face
            #dkro_μ_c_pcS_dc_oil = c_oil_face.*dpc_dS_face./μ_oil_face.*dkro_dc_face-kro_face.*c_oil_face.*dpc_dS_face./(μ_oil_face.*μ_oil_face).*dμ_dc_oil_face+kro_face.*dpc_dS_face./μ_oil_face+kro_face.*c_oil_face./μ_oil_face.*d2pc_dcdS_face
            #dkro_μ_ρ_pcS_dc_oil = ρ_oil_face.*dpc_dS_face./μ_oil_face.*dkro_dc_face-kro_face.*ρ_oil_face.*dpc_dS_face./(μ_oil_face.*μ_oil_face).*dμ_dc_oil_face+kro_face.*dpc_dS_face./μ_oil_face.*dρ_dc_oil_face+kro_face.*ρ_oil_face./μ_oil_face.*d2pc_dcdS_face
            dKw_dc              = k_o.*k_w./(Kval_cell.*k_o+k_w).*dKval_dc_cell-Kval_cell.*(k_o.*k_o).*k_w./((Kval_cell.*k_o+k_w).*(Kval_cell.*k_o+k_w)).*dKval_dc_cell
            dKo_c_dc            = k_o.*k_w./(Kval_cell.*k_o+k_w)-(k_o.*k_o).*k_w.*c_oil_val./((Kval_cell.*k_o+k_w).*(Kval_cell.*k_o+k_w)).*dKval_dc_cell

            #Derivatives of p
            dρ2_dp_water    = 2.0.*ρ_water_face.*dρ_dp_water_face
            dρ2_dp_oil      = 2.0.*ρ_oil_face.*dρ_dp_oil_face

            #Derivatives of S
            #dkro_pcS_dS_oil     = dkro_dc_face.*dpc_dS_face+kro_face.*d2pc_d2S_face

            #Phase velocities
            uw = -k_face.*krw_face./μ_water_face.*(∇p0-ρ_water_face.*g)
            uo = -k_face.*kro_face./μ_oil_face.*(∇p0-ρ_oil_face.*g)#+dpc_dS_face.*∇S0
            ut = uw + uo

            #println("uw=", mean(uw.xvalue), " , uo=", mean(uo.xvalue))

            #Jacobians related to coupled reactive transport of water (row 1)
            J_t_c_water_1, RHS_t_c_water_1  = transientTerm(c_water_init, dt, ϕ.*sw_val)
            J_adv_c_water_p_1               = convectionUpwindTerm(-k_face.*dc_μ_dc_water.*krw_face.*∇p0, uw)
            J_adv_c_water_g_1               = convectionUpwindTerm(-k_face.*ρ_water_face.*dc_μ_dc_water.*krw_face.*g, uw)
            J_lin_c_water_cw_1              = linearSourceTerm(createCellVariable(m, internalCells(K_w_cell).*a.*internalCells(ϕ)))
            J_c_water_1                     = J_t_c_water_1 + J_adv_c_water_p_1 - J_adv_c_water_g_1 + J_lin_c_water_cw_1

            J_adv_c_oil_p_1                 = convectionUpwindTerm(-k_face.*c_water_face./μ_water_face.*dkrw_dc_face.*∇p0, uw)
            J_adv_c_oil_g_1                 = convectionUpwindTerm(-k_face.*ρ_water_face.*c_water_face./μ_water_face.*dkrw_dc_face.*g, uw)
            J_lin_c_oil_cw_1                = linearSourceTerm(createCellVariable(m, internalCells(dKw_dc).*a.*internalCells(ϕ).*internalCells(c_water_val)))
            J_lin_c_oil_co_1                = linearSourceTerm(createCellVariable(m, internalCells(dKo_c_dc).*a.*internalCells(ϕ)))
            J_c_oil_1                       = J_adv_c_oil_p_1 - J_adv_c_oil_g_1 + J_lin_c_oil_cw_1 - J_lin_c_oil_co_1

            J_adv_p_g_1                     = convectionUpwindTerm(-k_face.*dρ_dp_water_face.*c_water_face./μ_water_face.*krw_face.*g, uw)
            J_diff_p_1                      = diffusionTerm(-k_face.*c_water_face./μ_water_face.*krw_face)
            J_p_1                           = -J_adv_p_g_1 + J_diff_p_1

            J_t_S_1, RHS_t_S_1              = transientTerm(sw_init, dt, ϕ.*c_water_val)
            J_adv_S_p_1                     = convectionUpwindTerm(-k_face.*(c_water_face./μ_water_face).*dkrw_dS_face.*∇p0, uw)
            J_adv_S_g_1                     = convectionUpwindTerm(-k_face.*ρ_water_face.*c_water_face./μ_water_face.*dkrw_dS_face.*g, uw)
            J_S_1                           = J_t_S_1 + J_adv_S_p_1 - J_adv_S_g_1

            #Jacobians related to coupled reactive transport of oil (row 2)
            J_lin_c_water_cw_2              = linearSourceTerm(createCellVariable(m, internalCells(K_w_cell).*a.*internalCells(ϕ)))
            J_c_water_2                     = -J_lin_c_water_cw_2

            J_t_c_oil_2, RHS_t_c_oil_2      = transientTerm(c_oil_init, dt, ϕ.*(1.0-sw_val))
            J_adv_c_oil_p_2                 = convectionUpwindTerm(-k_face.*dc_μ_kro_dc_oil.*∇p0, uo)
            J_adv_c_oil_pc_2                = zeros(N,N)#convectionUpwindTerm(-k_face.*dkro_μ_c_pcS_dc_oil.*∇S0, uo)
            J_adv_c_oil_g_2                 = convectionUpwindTerm(-k_face.*dρ_μ_c_kro_dc_oil.*g, uo)
            J_lin_c_oil_cw_2                = linearSourceTerm(createCellVariable(m, internalCells(dKw_dc).*a.*internalCells(ϕ).*internalCells(c_water_val)))
            J_lin_c_oil_co_2                = linearSourceTerm(createCellVariable(m, internalCells(dKo_c_dc).*a.*internalCells(ϕ)))
            J_c_oil_2                       = J_t_c_oil_2 + J_adv_c_oil_p_2 + J_adv_c_oil_pc_2 - J_adv_c_oil_g_2 - J_lin_c_oil_cw_2 + J_lin_c_oil_co_2

            J_adv_p_g_2                     = convectionUpwindTerm(-k_face.*dρ_dp_oil_face.*c_oil_face./μ_oil_face.*kro_face.*g, uo)
            J_diff_p_2                      = diffusionTerm(-k_face.*c_oil_face./μ_oil_face.*kro_face)
            J_p_2                           = -J_adv_p_g_2 + J_diff_p_2

            J_t_S_2, RHS_t_S_2              = transientTerm(sw_init, dt, -ϕ.*c_oil_val)
            J_adv_S_p_2                     = convectionUpwindTerm(-k_face.*c_oil_face./μ_oil_face.*dkro_dS_face.*∇p0, uo)
            J_adv_S_pc_2                    = zeros(N,N)#convectionUpwindTerm(-k_face.*c_oil_face./μ_oil_face.*dkro_pcS_dS_oil.*∇S0, uo)
            J_adv_S_g_2                     = convectionUpwindTerm(-k_face.*ρ_oil_face.*c_oil_face./μ_oil_face.*dkro_dS_face.*g, uo)
            J_diff_S_pc_2                   = zeros(N,N)#diffusionTerm(-k_face.*c_oil_face./μ_oil_face.*kro_face.*dpc_dS_face)
            J_S_2                           = J_t_S_2 + J_adv_S_p_2 + J_adv_S_pc_2 - J_adv_S_g_2  + J_diff_S_pc_2

            #Jacobians related to two-phase flow of water (row 3)
            J_adv_c_water_p_3               = convectionUpwindTerm(-k_face.*ρ_water_face.*d1_μ_dc_water.*krw_face.*∇p0, uw)
            J_adv_c_water_g_3               = convectionUpwindTerm(-k_face.*(ρ_water_face.*ρ_water_face).*d1_μ_dc_water.*krw_face.*g, uw)
            J_c_water_3                     = J_adv_c_water_p_3 - J_adv_c_water_g_3

            J_adv_c_oil_p_3                 = convectionUpwindTerm(-k_face.*ρ_water_face.*(1.0./μ_water_face).*dkrw_dc_face.*∇p0, uw)
            J_adv_c_oil_g_3                 = convectionUpwindTerm(-k_face.*(ρ_water_face.*ρ_water_face).*(1.0./μ_water_face).*dkrw_dc_face.*g, uw)
            J_c_oil_3                       = J_adv_c_oil_p_3 - J_adv_c_oil_g_3

            J_t_p_3, RHS_t_p_3              = transientTerm(p_init, dt, ϕ.*sw_val.*dρ_dp_water_cell)
            J_adv_p_p_3                     = convectionUpwindTerm(-k_face.*dρ_dp_water_face.*1.0./μ_water_face.*krw_face.*∇p0, uw)
            J_adv_p_g_3                     = convectionUpwindTerm(-k_face.*dρ2_dp_water.*(1.0./μ_water_face).*krw_face.*g, uw)
            J_diff_p_3                      = diffusionTerm(-k_face.*ρ_water_face.*(1.0./μ_water_face).*krw_face)
            J_p_3                           = J_t_p_3 + J_adv_p_p_3 - J_adv_p_g_3 + J_diff_p_3

            J_t_S_3, RHS_t_S_3              = transientTerm(sw_init, dt, ϕ.*ρ_water_cell)
            J_adv_S_p_3                     = convectionUpwindTerm(-k_face.*ρ_water_face.*(1.0./μ_water_face).*dkrw_dS_face.*∇p0, uw)
            J_adv_S_g_3                     = convectionUpwindTerm(-k_face.*(ρ_water_face.*ρ_water_face).*(1.0./μ_water_face).*dkrw_dS_face.*g, uw)
            J_S_3                           = J_t_S_3 + J_adv_S_p_3 - J_adv_S_g_3

            #Jacobians related to two-phase flow of oil (row 4)
            J_t_c_oil_4, RHS_t_c_oil_4      = transientTerm(c_oil_init, dt, ϕ.*(1.0-sw_val).*dρ_dc_oil_cell)
            J_adv_c_oil_p_4                 = convectionUpwindTerm(-k_face.*dρ_μ_kro_dc_oil.*∇p0, uo)
            J_adv_c_oil_pc_4                = zeros(N,N)#convectionUpwindTerm(-k_face.*dkro_μ_ρ_pcS_dc_oil.*∇S0, uo)
            J_adv_c_oil_g_4                 = convectionUpwindTerm(-k_face.*dρ2_μ_kro_dc_oil.*g, uo)
            J_c_oil_4                       = J_t_c_oil_4 + J_adv_c_oil_p_4 + J_adv_c_oil_pc_4 - J_adv_c_oil_g_4

            J_t_p_4, RHS_t_p_4              = transientTerm(p_init, dt, ϕ.*(1.0-sw_val).*dρ_dp_oil_cell)
            J_adv_p_p_4                     = convectionUpwindTerm(-k_face.*dρ_dp_oil_face./μ_oil_face.*kro_face.*∇p0, uo)
            J_adv_p_pc_4                    = zeros(N,N)#convectionUpwindTerm(-k_face.*dρ_dp_oil_face./μ_oil_face.*kro_face.*dpc_dS_face.*∇S0, uo)
            J_adv_p_g_4                     = convectionUpwindTerm(-k_face.*dρ2_dp_oil.*(1.0./μ_oil_face).*kro_face.*g, uo)
            J_diff_p_4                      = diffusionTerm(-k_face.*ρ_oil_face./μ_oil_face.*kro_face)
            J_p_4                           = J_t_p_4 + J_adv_p_p_4 + J_adv_p_pc_4 - J_adv_p_g_4 + J_diff_p_4

            J_t_S_4, RHS_t_S_4              = transientTerm(sw_init, dt, -ϕ.*ρ_oil_cell)
            J_adv_S_p_4                     = convectionUpwindTerm(-k_face.*(ρ_oil_face./μ_oil_face).*dkro_dS_face.*∇p0, uo)
            J_adv_S_pc_4                    = zeros(N,N)#convectionUpwindTerm(-k_face.*ρ_oil_face./μ_oil_face.*dkro_pcS_dS_oil.*∇S0, uo)
            J_adv_S_g_4                     = convectionUpwindTerm(-k_face.*(ρ_oil_face.*ρ_oil_face)./μ_oil_face.*dkro_dS_face.*g, uo)
            J_diff_S_pc_4                   = zeros(N,N)#diffusionTerm(-k_face.*ρ_oil_face./μ_oil_face.*kro_face.*dpc_dS_face)
            J_S_4                           = J_t_S_4 + J_adv_S_p_4 + J_adv_S_pc_4 - J_adv_S_g_4 + J_diff_S_pc_4

            # Constant terms for all equations
            RHS_const_c_water_cw_1    = constantSourceTerm(createCellVariable(m, internalCells(K_w_cell).*a.*internalCells(ϕ).*internalCells(c_water_val)))
            RHS_const_c_oil_co_1      = constantSourceTerm(createCellVariable(m, internalCells(K_o_cell).*a.*internalCells(ϕ).*internalCells(c_oil_val)))
            RHS_div_g_1             = divergenceTerm(-k_face.*ρ_water_face.*c_water_face./μ_water_face.*krw_face.*g)

            RHS_const_c_water_cw_2    = constantSourceTerm(createCellVariable(m, internalCells(K_w_cell).*a.*internalCells(ϕ).*internalCells(c_water_val)))
            RHS_const_c_oil_co_2      = constantSourceTerm(createCellVariable(m, internalCells(K_o_cell).*a.*internalCells(ϕ).*internalCells(c_oil_val)))
            RHS_div_g_2             = divergenceTerm(-k_face.*ρ_oil_face.*c_oil_face./μ_oil_face.*kro_face.*g)

            RHS_div_g_3             = divergenceTerm(-k_face.*(ρ_water_face.*ρ_water_face)./μ_water_face.*krw_face.*g)

            RHS_div_g_4             = divergenceTerm(-k_face.*(ρ_oil_face.*ρ_oil_face)./μ_oil_face.*kro_face.*g)

            # create the PDE system M [c_water,c_oil,p,S]=RHS
            # x_val = [c_water_val.value[:], c_oil_val.value[:], p_val.value[:]; sw_val.value[:]]
            M = [M_c_water_1+J_c_water_1 J_c_oil_1 J_p_1 J_S_1;
                 J_c_water_2 M_c_oil_2+J_c_oil_2 J_p_2 J_S_2;
                 J_c_water_3 zeros(N,N) M_p_3+J_p_3 J_S_3;
                 zeros(N,N) J_c_oil_4 J_p_4 M_S_4+J_S_4]

            #RHS related to coupled reactive transport of water
            RHS_1 = RHS_bc_c_water + RHS_t_c_water_1 + RHS_t_S_1 + (J_adv_c_water_p_1 - J_adv_c_water_g_1 + J_lin_c_water_cw_1)*c_water_val.value[:] + (J_adv_c_oil_p_1 - J_adv_c_oil_g_1 + J_lin_c_oil_cw_1-J_lin_c_oil_co_1)*c_oil_val.value[:] - J_adv_p_g_1*p_val.value[:] + (J_adv_S_p_1 - J_adv_S_g_1)*sw_val.value[:] - RHS_const_c_water_cw_1 + RHS_const_c_oil_co_1 + RHS_div_g_1

            #RHS related to coupled reactive transport of oil
            RHS_2 = RHS_bc_c_oil + RHS_t_c_oil_2 + RHS_t_S_2 - J_lin_c_water_cw_2*c_water_val.value[:] + (J_adv_c_oil_p_2 + J_adv_c_oil_pc_2 - J_adv_c_oil_g_2 - J_lin_c_oil_cw_2 + J_lin_c_oil_co_2)*c_oil_val.value[:] - J_adv_p_g_2*p_val.value[:] + (J_adv_S_p_2 + J_adv_S_pc_2 - J_adv_S_g_2)*sw_val.value[:] + RHS_const_c_water_cw_2 - RHS_const_c_oil_co_2 + RHS_div_g_2

            #RHS related to two-phase flow of water
            RHS_3 = RHS_bc_p + RHS_t_p_3 + RHS_t_S_3 + (J_adv_c_water_p_3 - J_adv_c_water_g_3)*c_water_val.value[:] + (J_adv_c_oil_p_3 - J_adv_c_oil_g_3)*c_oil_val.value[:] + (J_adv_p_p_3 - J_adv_p_g_3)*p_val.value[:] + (J_adv_S_p_3 - J_adv_S_g_3)*sw_val.value[:] + RHS_div_g_3

            #RHS related to two-phase flow of oil
            RHS_4 = RHS_bc_s + RHS_t_c_oil_4 + RHS_t_p_4 + RHS_t_S_4 + (J_adv_c_oil_p_4 + J_adv_c_oil_pc_4 - J_adv_c_oil_g_4)*c_oil_val.value[:] + (J_adv_p_p_4 + J_adv_p_pc_4 - J_adv_p_g_4)*p_val.value[:] + (J_adv_S_p_4 + J_adv_S_pc_4 - J_adv_S_g_4)*sw_val.value[:] + RHS_div_g_4

            RHS = [RHS_1[:];
                   RHS_2[:];
                   RHS_3[:];
                   RHS_4[:]]

            x_sol = M\RHS

            c_water_new = x_sol[1:N]
            c_oil_new   = x_sol[(1:N)+N]
            p_new       = x_sol[(1:N)+(2*N)]
            s_new       = x_sol[(1:N)+(3*N)]

            error_c_water = sumabs(c_water_new[2:end-1]-c_water_val.value[2:end-1])
            error_c_oil = sumabs(c_oil_new[2:end-1]-c_oil_val.value[2:end-1])
            error_s = sumabs(s_new[2:end-1]-sw_val.value[2:end-1])

            c_water_val.value[:]    = c_water_new[:]
            c_oil_val.value[:]      = c_oil_new[:]
            p_val.value[:]          = p_new[:]

            #Appleyard approach to updating saturations
            if !DME_injection_ended
                dsw_apple = 0.2
                eps_apple = sqrt(eps()) # 1e-5
                for i in eachindex(s_new)
                    if s_new[i]>=(1-sor_cell.value[i])
                        if sw_val.value[i]<(1-sor_cell.value[i]-eps_apple)
                            sw_val.value[i] = 1-sor_cell.value[i]-eps_apple
                        else
                            sw_val.value[i] = 1-sor_cell.value[i]
                        end
                    elseif s_new[i]<=swc
                        if sw_val.value[i]> swc+eps_apple
                            sw_val.value[i] = swc+eps_apple
                        else
                            sw_val.value[i] = swc
                        end
                    elseif abs(s_new[i]-sw_val.value[i])>dsw_apple
                        sw_val.value[i] += dsw_apple*sign(s_new[i]-sw_val.value[i])
                    else
                        sw_val.value[i] = s_new[i]
                    end
                end
            else
                w_c_w = 0.6
                w_c_o = 0.6
                w_p = 0.9
                w_sw = 0.6
                c_water_val.value[:] = w_c_w*c_water_new[:]+(1-w_c_w)*c_water_val.value[:]
                c_oil_val.value[:] = w_c_o*c_oil_new[:]+(1-w_c_o)*c_oil_val.value[:]
                p_val.value[:] = w_p*p_new[:]+(1-w_p)*p_val.value[:]
                sw_val.value[:] = w_sw*s_new[:]+(1-w_sw)*sw_val.value[:]
            end

            sw_val = createCellVariable(m, internalCells(sw_val), BCs)
        end

        if loop_countor < max_int_loop
            for tracer_count = 1:4
                # tracer calculation
                M_t_t, RHS_t_t = transientTerm(c_t_init, dt, sw_val.*ϕ)
                M_s_t          = linearSourceTerm(ϕ.*(sw_val-sw_init)/dt)
                M_a_t, RHS_a_t = convectionTvdTerm(uw, c_t_val, FL, ut)
                c_t_val = solveLinearPDE(m, M_t_t+M_s_t+M_a_t+M_bc_t, RHS_a_t+RHS_bc_t+RHS_t_t)
            end
            c_water_init    = copyCell(c_water_val)
            c_oil_init      = copyCell(c_oil_val)
            p_init          = copyCell(p_val)
            sw_init         = copyCell(sw_val)
            c_t_init        = copyCell(c_t_val)

            sor_face        = faceEval(sor_c, c_oil_face)
            t += dt
            dt = dt0
            println("progress: $(t/t_final_end*100) [%]")
            rec_fact=push!(rec_fact, (oil_init-domainInt(cellEval(ρ_oil,c_oil_val, p_val).*(1-sw_val)))/oil_init)
            t_s=push!(t_s, t)
        end
    end

    if m.dimension >= 1 && m.dimension < 2 && plotFigures
        plot_DME_1D(k, sor_c, kro0_c, KRO, KRW, BC_DME, c_water_init, c_oil_init, p_init, sw_init, c_t_init, t_s, rec_fact)
    elseif m.dimension >= 2 && m.dimension < 3 && plotFigures
        plot_DME_2D(k, KRO, KRW, BC_DME, c_water_init, c_oil_init, p_init, sw_init, c_t_init, t_s, rec_fact)
    elseif plotFigures
        plot_DME_3D(k, KRO, KRW, BC_DME, c_water_init, c_oil_init, p_init, sw_init, c_t_init, t_s, rec_fact)
    end

    return t_s, rec_fact
end

function plot_DME_1D(k::CellValue, sor_c::Function, kro0_c::Function, KRO::Function, KRW::Function, BC_DME::BoundaryConditionDME, c_water_init::CellValue, c_oil_init::CellValue, p_init::CellValue, sw_init::CellValue, c_t_init::CellValue, t_s, rec_fact)
    figure()
    grid("on")
    visualizeCells(k)
    title("Permeability")
    xlabel(L"$L_x$"" [m]")
    ylabel("Permeability, "L"$k$"" "L"[m^2]")
    savefig("1D Results/perm_field.png")

    figure()
    grid("on")
    visualizeCells(cellEval(sor_c, c_oil_init))
    ylim([0, 1])
    title("Residual Oil Saturation")
    xlabel(L"$L_x$"" [m]")
    ylabel("Residual Oil Saturation, "L"$S_{or}$"" "L"[-]")
    savefig("1D Results/residual_oil.png")

    figure()
    grid("on")
    sw_plot = collect(linspace(0,1,1000))
    sor = sor_c(0)
    plot(sw_plot, KRO.(sw_plot, kro0_c(sor), sor), sw_plot, KRW.(sw_plot, sor))
    plot([swc, swc], [0, 1], color=:black, linestyle="--")
    plot([1-sor, 1-sor], [0, 1], color=:black, linestyle="--")
    axis([0,1,0,1])
    title("Relative permeability, "L"$k_{ro},\ k_{rw}$"" [-]")
    xlabel(L"$S_w$"" [-]")
    ylabel(L"$k_r$"" [-]")
    legend([L"$k_{ro}$", L"$k_{rw}$"])
    savefig("1D Results/rel_perm_start.png")

    figure()
    grid("on")
    sw_plot = collect(linspace(0,1,1000))
    sor = sor_c(maximum(BC_DME.BCc_oil.left.c[:]))
    plot(sw_plot, KRO.(sw_plot, kro0_c(sor), sor), sw_plot, KRW.(sw_plot, sor))
    plot([swc, swc], [0, 1], color=:black, linestyle="--")
    plot([1-sor, 1-sor], [0, 1], color=:black, linestyle="--")
    axis([0,1,0,1])
    title("Relative permeability, "L"$k_{ro},\ k_{rw}$"" [-]")
    xlabel(L"$S_w$"" [-]")
    ylabel(L"$k_r$"" [-]")
    legend([L"$k_{ro}$", L"$k_{rw}$"])
    savefig("1D Results/rel_perm_end.png")

    figure()
    grid("on")
    visualizeCells(p_init/1e6) #MPa
    title("Pressure")
    xlabel(L"$L_x$"" [m]")
    ylabel("Pressure, "L"$p$"" [MPa]")
    savefig("1D Results/p_init.png")

    figure()
    grid("on")
    visualizeCells(sw_init)
    title("Water saturation")
    xlabel(L"$L_x$"" [m]")
    ylabel("Water saturation, "L"$S_w$"" [-]")
    savefig("1D Results/sw_init.png")

    figure()
    grid("on")
    c_oil_init.value[:] = round.(c_oil_init.value[:],3)
    visualizeCells(c_oil_init)
    title("Concentration of DME in oil")
    xlabel(L"$L_x$"" [m]")
    ylabel("Concentration of DME in oil, "L"$c_o$"" "L"\left[\frac{mol}{m^3}\right]")
    savefig("1D Results/c_oil_init.png")

    figure()
    grid("on")
    c_water_init.value[:] = round.(c_water_init.value[:],3)
    visualizeCells(c_water_init)
    title("Concentration of DME in water")
    xlabel(L"$L_x$"" [m]")
    ylabel("Concentration of DME in water, "L"$c_w$"" "L"\left[\frac{mol}{m^3}\right]")
    savefig("1D Results/c_water_init.png")

    figure()
    grid("on")
    c_t_init.value[:] = round.(c_t_init.value[:],3)
    visualizeCells(c_t_init)
    title("Tracer")
    xlabel(L"$L_x$"" [m]")
    ylabel(L"$c_t$")
    savefig("1D Results/ct_init.png")

    figure()
    grid("on")
    plot(t_s/(60*60*24), rec_fact*100)
    ylim([0, 100])
    title("Recovery factor")
    xlabel(L"$t$"" [days]")
    ylabel(L"$RF$"" [%]")
    savefig("1D Results/recovery_factor.png")
end

function plot_DME_2D(k::CellValue, sor_c::Function, kro0_c::Function, KRO::Function, KRW::Function, BC_DME::BoundaryConditionDME, c_water_init::CellValue, c_oil_init::CellValue, p_init::CellValue, sw_init::CellValue, c_t_init::CellValue, t_s, rec_fact)
    figure()
    visualizeCells(k)
    colorbar()
    title("Permeability, "L"$k$"" "L"[m^2]")
    xlabel(L"$L_x$"" [m]")
    ylabel(L"$L_y$"" [m]")
    savefig("2D Results/perm_field.png")

    figure()
    sor = cellEval(sor_c, c_oil_init)
    sor.value[:] = round(sor.value[:],2)
    visualizeCells(sor)
    colorbar()
    title("Residual Oil Saturation, "L"$S_{or}$"" "L"[-]")
    xlabel(L"$L_x$"" [m]")
    ylabel(L"$L_y$"" [m]")
    savefig("2D Results/residual_oil.png")

    figure()
    grid("on")
    sw_plot = collect(linspace(0,1,1000))
    sor = sor_c(0)
    plot(sw_plot, KRO.(sw_plot, kro0_c(sor), sor), sw_plot, KRW.(sw_plot, sor))
    plot([swc, swc], [0, 1], color=:black, linestyle="--")
    plot([1-sor, 1-sor], [0, 1], color=:black, linestyle="--")
    axis([0,1,0,1])
    title("Relative permeability, "L"$k_{ro},\ k_{rw}$"" [-]")
    xlabel(L"$S_w$"" [-]")
    ylabel(L"$k_r$"" [-]")
    legend([L"$k_{ro}$", L"$k_{rw}$"])
    savefig("2D Results/rel_perm_start.png")

    figure()
    grid("on")
    sw_plot = collect(linspace(0,1,1000))
    sor = sor_c(maximum(BC_DME.BCc_oil.left.c[:]))
    plot(sw_plot, KRO.(sw_plot, kro0_c(sor), sor), sw_plot, KRW.(sw_plot, sor))
    plot([swc, swc], [0, 1], color=:black, linestyle="--")
    plot([1-sor, 1-sor], [0, 1], color=:black, linestyle="--")
    axis([0,1,0,1])
    title("Relative permeability, "L"$k_{ro},\ k_{rw}$"" [-]")
    xlabel(L"$S_w$"" [-]")
    ylabel(L"$k_r$"" [-]")
    legend([L"$k_{ro}$", L"$k_{rw}$"])
    savefig("2D Results/rel_perm_end.png")

    figure()
    visualizeCells(p_init/1e6) #MPa
    colorbar()
    title("Pressure, "L"$p$"" [MPa]")
    xlabel(L"$L_x$"" [m]")
    ylabel(L"$L_y$"" [m]")
    savefig("2D Results/p_init.png")

    figure()
    visualizeCells(sw_init)
    colorbar()
    title("Water saturation, "L"$S_w$"" [-]")
    xlabel(L"$L_x$"" [m]")
    ylabel(L"$L_y$"" [m]")
    savefig("2D Results/sw_init.png")

    figure()
    visualizeCells(c_oil_init)
    colorbar()
    title("Concentration of DME in oil, "L"$c_o$"" "L"\left[\frac{mol}{m^3}\right]")
    xlabel(L"$L_x$"" [m]")
    ylabel(L"$L_y$"" [m]")
    savefig("2D Results/c_oil_init.png")

    figure()
    visualizeCells(c_water_init)
    colorbar()
    title("Concentration of DME in water, "L"$c_w$"" "L"\left[\frac{mol}{m^3}\right]")
    xlabel(L"$L_x$"" [m]")
    ylabel(L"$L_y$"" [m]")
    savefig("2D Results/c_water_init.png")

    figure()
    visualizeCells(c_t_init)
    colorbar()
    title("Tracer")
    xlabel(L"$L_x$"" [m]")
    ylabel(L"$L_y$"" [m]")
    savefig("2D Results/ct_init.png")

    figure()
    plot(t_s/(60*60*24), rec_fact*100)
    ylim([0, 100])
    title("Recovery factor")
    xlabel(L"$t$"" [days]")
    ylabel(L"$RF$"" [%]")
    savefig("2D Results/recovery_factor.png")
end

function plot_DME_3D(k::CellValue, KRO::Function, KRW::Function, BC_DME::BoundaryConditionDME, c_water_init::CellValue, c_oil_init::CellValue, p_init::CellValue, sw_init::CellValue, c_t_init::CellValue, t_s, rec_fact)
    figure()
    visualizeCells(k)
    colorbar()
    title("Permeability, "L"$k$"" "L"[m^2]")
    xlabel(L"$L_x$"" [m]")
    ylabel(L"$L_y$"" [m]")
    savefig("3D Results/perm_field.png")

    figure()
    sw_plot = collect(linspace(0,1,1000))
    plot(sw_plot, KRO.(sw_plot, sor_c(0)), sw_plot, KRW.(sw_plot, sor_c(0)))
    plot([swc, swc], [0, 1], color=:black)
    plot([1-sor_c(0),1-sor_c(0)], [0, 1], color=:black)
    axis([0,1,0,1])
    title("Relative permeability, "L"$k_{ro},\ k_{rw}$"" [-]")
    xlabel(L"$S_w$"" [-]")
    ylabel(L"$k_r$"" [-]")
    legend([L"$k_{ro}$", L"$k_{rw}$"])
    savefig("3D Results/rel_perm_start.png")

    figure()
    sw_plot = collect(linspace(0,1,1000))
    plot(sw_plot, KRO.(sw_plot, sor_c(maximum(BC_DME.BCc_oil.left.c[:]))), sw_plot, KRW.(sw_plot, sor_c(maximum(BC_DME.BCc_oil.left.c[:]))))
    plot([swc, swc], [0, 1], color=:black)
    plot([1-sor_c(maximum(BC_DME.BCc_oil.left.c[:])),1-sor_c(maximum(BC_DME.BCc_oil.left.c[:]))], [0, 1], color=:black)
    axis([0,1,0,1])
    title("Relative permeability, "L"$k_{ro},\ k_{rw}$"" [-]")
    xlabel(L"$S_w$"" [-]")
    ylabel(L"$k_r$"" [-]")
    legend([L"$k_{ro}$", L"$k_{rw}$"])
    savefig("3D Results/rel_perm_end.png")

    figure()
    visualizeCells(p_init/1e6) #MPa
    colorbar()
    title("Pressure, "L"$p$"" [MPa]")
    xlabel(L"$L_x$"" [m]")
    ylabel(L"$L_y$"" [m]")
    savefig("3D Results/p_init.png")

    figure()
    visualizeCells(sw_init)
    colorbar()
    title("Water saturation, "L"$S_w$"" [-]")
    xlabel(L"$L_x$"" [m]")
    ylabel(L"$L_y$"" [m]")
    savefig("3D Results/sw_init.png")

    figure()
    visualizeCells(c_oil_init)
    colorbar()
    title("Concentration of DME in oil, "L"$c_o$"" "L"\left[\frac{mol}{m^3}\right]")
    xlabel(L"$L_x$"" [m]")
    ylabel(L"$L_y$"" [m]")
    savefig("3D Results/c_oil_init.png")

    figure()
    visualizeCells(c_water_init)
    colorbar()
    title("Concentration of DME in water, "L"$c_w$"" "L"\left[\frac{mol}{m^3}\right]")
    xlabel(L"$L_x$"" [m]")
    ylabel(L"$L_y$"" [m]")
    savefig("3D Results/c_water_init.png")

    figure()
    visualizeCells(c_t_init)
    colorbar()
    title("Tracer")
    xlabel(L"$L_x$"" [m]")
    ylabel(L"$L_y$"" [m]")
    savefig("3D Results/ct_init.png")

    figure()
    plot(t_s/(60*60*24), rec_fact*100)
    ylim([0, 100])
    title("Recovery factor")
    xlabel(L"$t$"" [days]")
    ylabel(L"$RF$"" [%]")
    savefig("3D Results/recovery_factor.png")
end

max_stars_repo_name: FrederikLehn/coupled-solvers

max_issues_repo_name: FrederikLehn/coupled-solvers

max_forks_repo_name: FrederikLehn/coupled-solvers


text:
struct MinEigenvalueShift <: AbstractShiftMethod end

function compute_diagonal_shift(Q::Matrix{Float64}, ::MinEigenvalueShift)
    @assert LinearAlgebra.issymmetric(Q)
    n = size(Q, 1)
    shift = -LinearAlgebra.eigmin(Q) .* ones(n)
    @assert LinearAlgebra.eigmin(Q + LinearAlgebra.diagm(shift)) >= -1e-4
    return shift
end

struct SemidefiniteShift <: AbstractShiftMethod
    sdp_optimizer_factory::Any
end

function compute_diagonal_shift(Q::Matrix{Float64}, shift::SemidefiniteShift)
    @assert LinearAlgebra.issymmetric(Q)
    n = size(Q, 1)
    sdp_model = JuMP.Model(shift.sdp_optimizer_factory)
    δ = JuMP.@variable(sdp_model, [i = 1:n])
    JuMP.@objective(sdp_model, Min, sum(δ))
    JuMP.@SDconstraint(sdp_model, Q + LinearAlgebra.diagm(δ .- 1e-4) >= 0)
    JuMP.optimize!(sdp_model)
    @assert JuMP.primal_status(sdp_model) == MOI.FEASIBLE_POINT
    shift = JuMP.value.(δ)
    @assert LinearAlgebra.eigmin(Q + LinearAlgebra.diagm(shift)) >= -1e-4
    return shift
end

ext: jl


text:
using Util
using Test

@testset "Util" begin
  @testset "dummy" begin
    @test true
  end
end


text:
using AWSAuth
using Test

include("aws4.jl")


text:
#
# Computes the total potential energy of the system
#

function utotal(p,data)
  utotal = 0.
  for ip in 1:data.N-1
    for jp in ip+1:data.N
      rij = rpbc(p[ip],p[jp],data.side)
      if rij < data.cutoff
        utotal += upair(rij,data)
      end
    end
  end
  return utotal
end


text:
function make_mat(dims=(3,4))
    mat = PETSc.Mat(ST, dims...)
    vt1 = RC(complex(3., 3.))
    vt2 = RC(complex(5., 5.))
    mat[1,1] = vt1
    mat[1,2] = vt2
    PETSc.assemble(mat)
    return mat
end

@testset "Mat{$ST}" begin

  @testset "Preallocator" begin
    for mt in [PETSc.C.MATMPIAIJ,PETSc.C.MATMPIBAIJ,PETSc.C.MATMPISBAIJ]
      @test_throws ArgumentError PETSc.Mat(ST, 3, 4, mtype=mt, nnz=collect(1:10))
      @test_throws ArgumentError PETSc.Mat(ST, 3, 4, mtype=mt, onnz=collect(1:10))
    end
    for mt in [PETSc.C.MATBLOCKMAT,PETSc.C.MATSEQAIJ,PETSc.C.MATSEQBAIJ,PETSc.C.MATSEQSBAIJ]
      @test_throws ArgumentError PETSc.Mat(ST, 3, 4, mtype=mt, nnz=collect(1:10))
    end
  end

  @testset "Shell Matrix" begin
#    if ST == Float64  # until Clang works correctly
      ctx = (2, 4)
      mat = MatShell(ST, 3, 3, ctx)
      ctx_ret = getcontext(mat)
#      println("typeof(ctx) = ", typeof(ctx))
#      println("typeof(ctx_ret) = ", typeof(ctx_ret))
      @test ctx_ret == ctx

      f_ptr = cfunction(mymult, PETSc.C.PetscErrorCode, (PETSc.C.Mat{ST}, PETSc.C.Vec{ST}, PETSc.C.Vec{ST}))
      setop!(mat, PETSc.C.MATOP_MULT, f_ptr)
      x = Vec(ST[1.0, 2, 3])
      b = mat*x
      @test b == ST[1.0, 4.0, 9.0]
#    end

  end  # end testset Shell Matrix

  vt1 = RC(complex(3., 3.))
  vt2 = RC(complex(5., 5.))
  @testset "Utility functions" begin
    mat = make_mat()

    @test size(mat) == (3,4)
    @test sizelocal(mat) == (3,4)
    @test lengthlocal(mat) == 12
    val_ret = mat[1,1]
    vt1 = RC(complex(3., 3.))
    @test val_ret ≈ vt1
    vt1 = RC(complex(4., 4.))
    mat[1,2] = vt1
    PETSc.assemble(mat)

    mtype = PETSc.gettype(mat)
    @test mtype == PETSc.C.MATMPIAIJ

    mat_copy = Mat(mat.p)
    @test mat_copy == mat

    # test set/get index
    vt1 = RC(complex(3., 3.))
    vt2 = RC(complex(5., 5.))
    mat[1,1] = vt1
    mat[1,2] = vt2
    PETSc.assemble(mat)
    val_ret = mat[1,1]
    @test val_ret ≈ vt1
    vt1 = RC(complex(4., 4.))
    mat[1,2] = vt1
    PETSc.assemble(mat)

    #test nnz
    @test nnz(mat) == 2

    rows, cols = localranges(mat)
    @test rows == 1:size(mat, 1)
    @test cols == 1:size(mat, 2)

    @test PETSc.isfinalized(mat) == false
    PETSc.PetscDestroy(mat)
    @test PETSc.isfinalized(mat) == true
  end
  @testset "real and imag" begin
    mat  = make_mat()
    rmat = PETSc.Mat(ST, 3, 4)
    rmat[1,1] = RC(complex(3., 0.))
    rmat[1,2] = RC(complex(5., 0.))
    PETSc.assemble(rmat)
    @test real(mat)[1,1] == rmat[1,1]
    @test real(mat)[1,2] == rmat[1,2]
    if ST <: Complex
        @test imag(mat)[1,1] == rmat[1,1]
        @test imag(mat)[1,2] == rmat[1,2]
    end
  end
  @testset "diag and trace" begin
    dmat = similar(make_mat(),3,3)
    dmat[1,1] = vt1
    assemble(dmat)
    d = diag(dmat)
    @test d[1] == vt1
    @test trace(dmat) == vt1
  end
  @testset "similar and resize" begin
    mat = similar(make_mat())
    @test size(mat) == (3,4)
    @test mat[1,1] != vt1
    @test size(similar(mat, ST)) == (3,4)
    @test size(similar(mat, 4, 4)) == (4,4)
    @test size(similar(mat, (4, 4))) == (4,4)
    @test mat[1,1] != make_mat()[1,1]
    @test_throws ArgumentError resize!(mat)
    @test_throws ArgumentError resize!(mat,5,mlocal=2)
  end
  @testset "copy and conj" begin
    mat = similar(make_mat((4,4)))
    @test size(mat) == (4,4)
    mat2 = copy(mat)
    @test mat2[1,1] ≈ mat[1, 1]
    @test conj(conj(mat))[1,1] ≈ mat[1,1]
  end
  @testset "getting Mat info, inserting and assembling" begin
    mat = make_mat()
    @test PETSc.getinfo(mat).block_size == 1
    @test isassembled(mat)

    mat2 = similar(mat, ST, 4, 4)
    mat2[1,2] = vt1
    mat2[1,3] = vt2
    PETSc.assemble(mat2)

    @test mat2[1,2] == vt1
    @test mat2[1,3] == vt2
    mat3 = Mat( ST, 3, 3)

    function increasing_diag()
      (m,n) = size(mat3)
      dim = min(m, n)
      for i=1:dim
        i_float = Float64(i)
        mat3[i,i] = RC(complex(i_float, i_float))
      end
    end

    assemble(increasing_diag, mat3)

    (m,n) = size(mat3)
    dim = min(m, n)

    for i=1:dim
      i_float = Float64(i)
      @test mat3[i,i] ≈ RC(complex(i_float, i_float))
    end
  end
  @testset "transpose and transpose!" begin
    mat = make_mat((3,3))
    ctmat = copy(mat)
    @test transpose!(transpose!(copy(ctmat))) == mat
    @test transpose(transpose(ctmat)) == mat
  end
  vt = RC(complex(3., 3.))
  @testset "array indexing" begin
    vals = RC(complex(rand(3, 2), rand(3,2)))
    idx = Array(1:3)
    idy = Array(1:2)
    @testset "sub indexing" begin
      mat = PETSc.Mat(ST, 3, 3)
      mat[idx, idy] = vals
      assemble(mat)
      matj = zeros(ST, 3,3)
      matj[1:3, 1:2] = vals
      @test mat == matj
      @test mat[idx,idy] ≈ vals
    end
    @testset "y indexing" begin
      mat = PETSc.Mat(ST, 3, 3)
      vals = RC( complex(rand(3), rand(3)))
      mat[1, idx] = vals
      assemble(mat)
      matj = zeros(ST, 3,3)
      matj[1, idx] = vals
      @test mat == matj

      vals_ret = mat[1, idx]
      @test vals_ret.' ≈ vals[idx]
    end
    @testset "x indexing" begin
      mat = PETSc.Mat(ST, 3, 3)
      vals = RC(complex(rand(3), rand(3)))
      mat[idx, 1] = vals
      assemble(mat)
      matj = zeros(ST, 3,3)
      matj[idx, 1] = vals
      @test mat == matj
      vals_ret = mat[idx, 1]
      @test vals_ret ≈ vals
    end
    @testset "x,y set and fetch" begin
      mat = PETSc.Mat(ST, 3, 3)
      mat[idx, idy] = vt
      assemble(mat)
      matj = zeros(ST, 3,3)
      matj[1:3, 1:2] = vt
      @test mat == matj

      mat = PETSc.Mat(ST, 3,3)
      vals = rand(ST, 3,3)
      mat[1:3, 1:3] = vals
      assemble(mat)
      @test mat == vals
    end
    @testset "x set and fetch" begin
      mat = PETSc.Mat(ST, 3, 3)
      mat[idx, 1] = vt
      assemble(mat)
      matj = zeros(ST, 3,3)
      matj[1:3, 1] = vt
      @test mat == mat
    end
    @testset "y set and fetch" begin
      mat = PETSc.Mat(ST, 3, 3)
      mat[1, idy] = vt
      assemble(mat)
      matj = zeros(ST, 3,3)
      matj[1, 1:2] = vt
      @test mat == matj
    end
  end

  @testset "SubMatrix" begin
    mat = PETSc.Mat(ST, 6, 6, nz=6)
    for i=1:6
      for j=1:6
        mat[i,j] = i*6 + j
      end
    end
    assemble(mat, PETSc.C.MAT_FLUSH_ASSEMBLY)
#    petscview(mat)
    isx = IS(ST, [1, 2, 3])
    isy = IS(ST, [1, 2])
    smat = PETSc.SubMat(mat, isx, isx)
    val = RC(complex(2.0, 2.0))

    smat[1,1] =val
    SubMatRestore(smat)
    assemble(mat)
    @test mat[1,1] == val
  end
  @testset "test ranges and colon" begin
    idy = Array(1:2)
    @testset "submatrix" begin
      mat = PETSc.Mat(ST, 3, 3)
      mat[1:3, 1:2] = vt
      assemble(mat)
      matj = zeros(ST, 3,3)
      matj[1:3, 1:2] = vt
      @test mat == matj
    end
    @testset "on an axis with range" begin
      mat = PETSc.Mat(ST, 3, 3)
      mat[:, idy] = vt
      assemble(mat)
      matj = zeros(ST, 3,3)
      matj[:, 1:2] = vt
      @test mat == matj
    end
    @testset "on a column" begin
      vals = [1, 2, 3]
      mat = PETSc.Mat(ST, 3, 3)
      mat[:, 1] = vt
      assemble(mat)
      matj = zeros(ST, 3,3)
      matj[:, 1] = vt
      @test mat == matj
    end
  end

  @testset "full and fill" begin
    vt = RC(complex(1.,1.))
    mat = PETSc.Mat(ST, 3, 3)
    fill!(mat, vt)
    assemble(mat)
    @test mat == fill(vt,(3,3))
    matjd = full(mat)
    @test mat == matjd
  end

  function check_mats(a, b)
    for i=1:size(a, 2)
      for j=1:size(a, 1)
        @test a[i, j] ≈ b[i,j]
      end
    end
  end
  @testset "0-based indexing " begin
    mat = PETSc.Mat(ST, 3, 3)
    matj = zeros(ST, 3, 3)
    idx = PetscInt[1, 2]
    idy = PetscInt[0, 1]
    vals = ST[1.0 2.0; 3.0 4.0]
    set_values!(mat, idx, idy, vals)
    set_values!(matj, idx, idy, vals)
    assemble(mat)
    check_mats(mat, matj)

    maprow, mapcol = local_to_global_mapping(mat)
    set_local_to_global_mapping(mat, maprow, mapcol)
    vals = 2*vals
    set_values_local!(mat, idx, idy, vals)
    set_values_local!(matj, idx, idy, vals)
    assemble(mat)
    check_mats(mat, matj)

    bs=2
    mat = PETSc.Mat(ST, 6, 6, bs=bs)
    matj = zeros(ST, 12, 12)
    idx = [0, 1]
    idy = [0, 1, 2]
    vals = rand(ST, bs*length(idx), bs*length(idy))
    set_values_blocked!(mat, idx, idy, vals)
    assemble(mat)
    for j=1:length(idy)
      start_idy = idy[j]*bs
      for i=1:length(idx)
        start_idx = idx[i]*bs
        for k=1:bs
          for p=1:bs
            colidx = start_idy + k
            rowidx = start_idx + p
            @test mat[rowidx, colidx] ≈ vals[(i-1)*bs + p, (j-1)*bs + k]
          end
        end
      end
    end

    vals = 2*vals
    maprow, mapcol = local_to_global_mapping(mat)
    set_local_to_global_mapping(mat, maprow, mapcol)
    set_values_blocked_local!(mat,  idx, idy, vals)
    assemble(mat)
    for j=1:length(idy)
      start_idy = idy[j]*bs
      for i=1:length(idx)
        start_idx = idx[i]*bs
        for k=1:bs
          for p=1:bs
            colidx = start_idy + k
            rowidx = start_idx + p
            @test mat[rowidx, colidx] ≈ vals[(i-1)*bs + p, (j-1)*bs + k]
          end
        end
      end
    end

  end

  @testset "Sparse Matrix conversion" begin
    A = sprand(10, 10, 0.1)
    B = Mat(A)
    assemble(B)
    @test A == B
    info = PETSc.getinfo(B)
    @test info.mallocs == 0

    A2 = full(A)
    B2 = Mat(B)
    assemble(B2)
    @test A2 == B2
    info = PETSc.getinfo(B2)
    @test info.mallocs == 0

    dtol = 1e-16
    A3 = ST[1. 2 3; 4 5 6; 7 8 dtol/10]
    B3 = Mat(A3, droptol=dtol)
    assemble(B3)
    A3b = copy(A3)
    A3b[3, 3] = 0
    @test A3b == B3
    info = PETSc.getinfo(B3)
    @test info.nz_used == 8
  end

  @testset "test conversion of values to a new type" begin
    mata = PETSc.Mat(ST, 3, 3)
    matb = PETSc.Mat(ST, 3, 3)
    mataj = zeros(ST, 3, 3)
    matbj = zeros(ST, 3, 3)
    vec = PETSc.Vec(ST, 3)
    vecj = zeros(ST, 3)
    cnt = 1
    for i=1:3
      for j=1:3
        cnt_f = RC(complex(Float64(cnt), Float64(cnt)))
        cnt_f2 = RC(complex(Float64(cnt + 9), Float64(cnt + 9)))

        mata[i,j] = cnt_f
        mataj[i,j] = cnt_f
        matb[i,j] = cnt_f2
        matbj[i,j] = cnt_f2
        cnt += 1
      end
      vec[i] = RC(complex(Float64(i), i))
      vecj[i] = RC(complex(Float64(i), i))
      assemble(vec)
    end

    assemble(mata)
    assemble(matb)

    @testset "matrix-vector product" begin
      result = mata*vec
      resultj = mataj*vecj
      @test result == resultj

      result = mata.'*vec
      resultj = mataj.'*vecj
      @test result == resultj

      result = mata'*vec
      resultj = mataj'*vecj
      @test result == resultj
    end
    @testset "binary matrix operations" begin
      result = mata + matb
      assemble(result)
      resultj = mataj + matbj
      @test result == resultj
      result = mata - matb
      assemble(result)
      resultj = mataj - matbj
      @test result == resultj
      result = mata * matb
      assemble(result)
      resultj = mataj * matbj
      @test result == resultj
    end
    @testset "matrix operations with numbers" begin
      result = 2*mata
      assemble(result)
      resultj = 2*mataj

      @test result == resultj
      result = mata/2
      assemble(result)
      resultj = mataj/2
      @test result == resultj

      result = 2.\mata
      assemble(result)
      resultj = 2.\mataj
      @test result == resultj

      result  = -mata
      resultj = -mataj
      @test result == resultj
    end
  end

  @testset "Testing {c}transpose mults" begin
    mat1  = PETSc.Mat(ST,3,3,mtype=PETSc.C.MATSEQAIJ)
    mat2  = PETSc.Mat(ST,3,3,mtype=PETSc.C.MATSEQAIJ)
    mat1j = zeros(ST,3,3)
    mat2j = zeros(ST,3,3)
    vt1 = RC(complex(3., 0.))
    vt2 = RC(complex(5., 5.))
    vt3 = RC(complex(4., 4.))
    vt4 = RC(complex(10., 0.))
    vt5 = RC(complex(1., 0.))
    mat1[1,1] = vt1
    mat1[1,2] = vt2
    mat1[2,1] = conj(vt2)
    mat1[1,3] = vt3
    mat1[3,1] = conj(vt3)
    mat1[2,2] = vt4
    mat1[3,3] = vt5
    mat1j[1,1] = vt1
    mat1j[1,2] = vt2
    mat1j[2,1] = conj(vt2)
    mat1j[1,3] = vt3
    mat1j[3,1] = conj(vt3)
    mat1j[2,2] = vt4
    mat1j[3,3] = vt5
    mat2[1,1] = vt1
    mat2[1,2] = vt2
    mat2[2,1] = vt2
    mat2[1,3] = vt3
    mat2[3,1] = vt3
    mat2[2,2] = vt4
    mat2[3,3] = vt5
    mat2j[1,1] = vt1
    mat2j[1,2] = vt2
    mat2j[2,1] = vt2
    mat2j[1,3] = vt3
    mat2j[3,1] = vt3
    mat2j[2,2] = vt4
    mat2j[3,3] = vt5
    PETSc.assemble(mat1)
    PETSc.assemble(mat2)
    @test ishermitian(mat1)
    @test issym(mat2)

    mat3 = mat1.'*mat2
    mat3j = mat1j.'*mat2j
    assemble(mat3)
    @test mat3 == mat3j

    mat4 = mat1*mat2.'
    mat4j = mat1j*mat2j.'
    assemble(mat4)
    @test mat4 == mat4j
  end

  @testset "MatRow" begin
    matj = [1. 0 0; 0 2 3; 4 5 6];
    mat = Mat(matj)
    assemble(mat)
    val = 1
    for i=1:3
      row_i = PETSc.MatRow(mat, i)
      @test length(row_i) == i
      for j=1:length(row_i)
        @test PETSc.getval(row_i, j) ≈ val
        val += 1
      end
      restore(row_i)
    end

    row = PETSc.MatRow(mat, 1)
    @test PETSc.getcol(row, 1) == 1
    restore(row)
    @test PETSc.count_row_nz(mat, 1) == 1

    row = PETSc.MatRow(mat, 2)
    @test PETSc.getcol(row, 1) == 2
    @test PETSc.getcol(row, 2) == 3
    restore(row)
    @test PETSc.count_row_nz(mat, 2) == 2

    row = PETSc.MatRow(mat, 3)
    @test PETSc.getcol(row, 1) == 1
    @test PETSc.getcol(row, 2) == 2
    @test PETSc.getcol(row, 3) == 3
    restore(row)
    @test PETSc.count_row_nz(mat, 3) == 3
  end

  @testset "kron" begin
    function testrun(Aj, Bj)
      # A and B are julia matrices of some kind
      A = Mat(Aj); B = Mat(Bj)
      assemble(A); assemble(B)
      Cj = kron(Aj, Bj)
      C = kron(A, B)
      assemble(C)
      m, n = size(C)
      C2 = C[1:m, 1:n]
      @test C2 ≈ Cj
    end

    # case 1
    Aj = ST[1. 0 0; 0 1 0; 0 0 1]
    testrun(Aj, Aj)

    # case 2
    Aj = ST[1. 2 0; 3 4 5; 0 6 7]
    testrun(Aj, Aj)

    n = 10
    Aj = rand(ST, n, n)
    testrun(Aj, Aj)

    m = 3
    n = 2
    Aj = rand(ST, m, n)
    testrun(Aj, Aj)

    m1 = 3
    n1 = 2
    m2 = 4
    n2 = 5
    Aj = rand(ST, m1, n1)
    Bj = rand(ST, m2, n2)
    testrun(Aj, Bj)

    m1 = 5
    m2 = 7
    n1 = 8
    n2 = 10
    Aj = sprand(m1, n1, 0.01)
    Bj = sprand(m2, n2, 0.01)
    testrun(Aj, Bj)

    function testkron(Aj, Bj)
      m1, n1 = size(Aj)
      m2, n2 = size(Bj)
      Cj = kron(Aj, Bj)
      C = PETSc.PetscKron(Aj, Bj)
      assemble(C)
      C_full = C[1:(m1*m2), 1:(n1*n2)]
      Cj_full = full(Cj)
    #  @test Cj_full ≈ C_full
    end

#    println("\n----- Case 1 -----")
    Aj = sparse(eye(3, 3))
    testkron(Aj, Aj)

#    println("\n----- Case 2 -----")
    Aj = sparse([1. 2 0; 3 4 5; 0 6 7])
    testkron(Aj, Aj)

#    println("\n----- Case 3 -----")
    Aj = sparse([1. 2 3; 4 5 6; 7 8 9])
    Bj = sparse([10. 11 12 13; 14 15 16 17; 18 19 20 21])
    testkron(Aj, Bj)

#    println("\n----- Case 4 -----")
    Aj = sparse([1. 2 ; 4 5 ; 7 8 ])
    Bj = sparse([10. 11 12 13 22; 14 15 16 17 23; 18 19 20 21 24])
    testkron(Aj, Bj)

#    println("\n----- Case 5 -----")
    Aj = sprand(10, 15, 0.01)
    Bj = sprand(7, 21, 0.01)
    testkron(Aj, Bj)

  end
end


text:
# Roots of the characteristic equation of an SDE

"""
$SIGNATURES

Split roots into real and complex-valued roots. For the complex-valued roots,
only provides one of a complex conjugate pair. An error is raised if there is
no single complex conjugate
"""
function splitroots(r)
  rw = copy(r)
  rrc = empty(rw); ri = empty(rw)
  while !isempty(rw)
    rc = popfirst!(rw)
    if isreal(rc)
      push!(rrc,rc)
    else
      push!(ri,rc)
      is_conj = rw.==conj(rc)
      if count(is_conj)==0  error("No complex conjugate found for $rc")  end
      rw = rw[.!is_conj]
    end
  end
  rr = real(rrc)
  if 2length(ri)+length(rr)≠length(r)
    error("Order mismatch; complex roots may not occur in conjugate pairs")
  end
  return rr,ri
end

function splitroots(θ::NamedTuple)
  r_re = θ.r_re
  ri = if haskey(θ,:r_im)
    # zip not supported by Zygote; using buffer.
    # r_re_c = r_re[1:length(θ.r_im)]
    # map(x->complex(x...),zip(r_re_c,θ.r_im))
    nc = length(θ.r_im)
    B = zeros(complex(eltype(θ.r_im)),nc)
    ri_buf = Zygote.bufferfrom(B)
    for i in 1:nc
      ri_buf[i] = complex(r_re[i],θ.r_im[i])
    end
    copy(ri_buf)
  else
    (complex(eltype(r_re)))[]
  end
  nc = length(ri)
  rr = r_re[nc+1:end]
  return rr,ri
end

"
$SIGNATURES

Combine roots
"
function combineroots(rr,ri)
  if  !(eltype(rr)<:Real && eltype(ri)<:Complex) error("r_re must be real, and r_im complex")  end
  # ru = vcat(rr,ri,conj(ri)) # Zygote does not support
  nr = length(rr); ni = length(ri)
  ru_buf = Zygote.bufferfrom(zeros(eltype(ri),nr+2ni))
  ru_buf[1:nr] = rr
  ru_buf[nr.+(1:ni)] = ri
  ru_buf[nr.+ni.+(1:ni)] = conj(ri)
  ru = copy(ru_buf)
  r = sortroots(ru)
  return  r
end

sortroots(r) = sort(r;by=x->(real(x), imag(x)))

issorted_roots(r) = issorted(r;by=x->(real(x), imag(x)))

function issorted_rootstuple(θ)
  rr,ri = splitroots(θ)
  return issorted(rr) && issorted_roots(ri)
end

"""
$SIGNATURES

Random roots
"""
function randroots(rng,r,f,N,T)
  rreal,rcomplex = splitroots(r)
  Tr,Tc = eltype(rreal),eltype(rcomplex)
  rreald = map(rreal) do rc
    sdα = normparstd(rc*T)
    s = max(1,min(sdα,10)/T)
    σ = f*s/√N
    rd = uniform_rand(rng,rc,σ)
  end |> Vector{Tr}
  rcomplexd = map(rcomplex) do rc
    # Real component
    x = real(rc)
    sdα = normparstd(x*T)
    s = max(1,min(sdα,10)/T)
    σ = f*s/√N
    rdr = uniform_rand(rng,x,σ)
    # Imaginary component
    y = imag(rc)
    σ = f*1/T/√N
    rdi = uniform_rand(rng,y,σ)
    Complex(rdr,rdi)
  end |> Vector{Tc}
  rd = combineroots(rreald,rcomplexd)
  return rd
end

"""
$FIELDS

Roots representation of a Stochastic Differential Equation `Sde`.
"""
struct SdeRoots{T<:Complex,U} <: ScalarSde{T}
  "Roots"
  r::Vector{T}
  "Driving noise standard deviation"
  σϵ::U
  function SdeRoots(r::Vector{T},σϵ::U) where {T,U<:Real}
    if !(σϵ>0)  error("σϵ must be positive")  end
    new{T,U}(r,complex(σϵ))
  end
end

function Base.show(io::IO,eq::SdeRoots)
  p = order(eq)
  print(io,"SDE($p) roots "*join(rootstr.(eq.r),"; "))
  print(io,@sprintf " | σϵ = %0.3g" eq.σϵ)
end

function Base.show(io::IO,::MIME"text/plain",eq::SdeRoots)
  print(io,"Roots of Scalar Stochastic Differential Equation\n",eq)
end

function SdeRoots(θ::NamedTuple,σϵ)
  rr,ri = splitroots(θ)
  r = combineroots(rr,ri)
  return SdeRoots(r,σϵ)
end

function SdeRoots(r::AbstractVector{T};σy=one(real(T))) where {T}
  σϵ1 = one(real(T))
  g = std(SdeRoots(r,σϵ1))
  σϵ = σy/g
  return SdeRoots(r,σϵ)
end

function Statistics.std(eq::SdeRoots)
  if order(eq)==0 return eq.σϵ end
  mv = convert(VectorSde,eq)
  Σs = statcovz(mv)
  σy = sum(Σs) |> real |> √
  return σy
end

function SdeRoots(eq::Sde{T}) where {T}
  p = Polynomial([eq.a; one(T)])
  r = complex(roots(p))
  return SdeRoots(sortroots(r),eq.σϵ)
end

Polynomials.roots(eq::AbstractSde) = SdeRoots(eq).r

function Base.convert(::Type{VectorSde},eq::SdeRoots)
  r = eq.r
  if length(unique(r))<length(r) error("Roots must be unique") end
  p = order(eq)
  A = Diagonal(r)
  C = ones(eltype(r),1,p)
  U = [rc^(j-1) for j in 1:p, rc in r]
  v = U\[zeros(p-1);1]
  Q = eq.σϵ^2*v*v'
  R = zeros(eltype(r),1,1)
  veq = VectorSde(A,C,Q,R)
  return veq
end

function Base.convert(::Type{AbstractSde{U}},eq::SdeRoots{T}) where {T,U<:Complex}
  r = convert(Vector{U},eq.r)
  σϵ = convert(real(U),eq.σϵ)
  return SdeRoots(r,σϵ)
end

function Statistics.mean(eqs::Vector{<:SdeRoots})
  r = mean(getproperty.(eqs,:r))
  σϵ = mean(getproperty.(eqs,:σϵ))
  return SdeRoots(r,σϵ)
end

"""
Matrix of eigenvectors `U` of the `A` matrix for the SDE roots `eq`
"""
function eigenvecA(eq::SdeRoots)
  r = eq.r
  p = order(eq)
  # Next line not supported by Zygote, using Buffer instead
  # U = [rc^(i-1) for i in 1:p, rc in r]
  U_buf = Zygote.Buffer(r,p,p)
  for i in 1:p, (j,rc) in enumerate(r)
    U_buf[i,j] = rc^(i-1)
  end
  U = copy(U_buf)
  return U
end

function statcovz(eq::SdeRoots,U)
  p = order(eq)
  v = U\[zeros(p-1);1]
  Q = eq.σϵ^2*v*v'
  r = eq.r
  # Next line not supported by Zygote, using Buffer instead
  # Σs = [-Q[i,j]/(r[i]+conj(r[j])) for i in 1:p,j in 1:p]
  Σs_buf = Zygote.Buffer(v,p,p)
  for i in 1:p, j in 1:p
    Σs_buf[i,j] = -Q[i,j]/(r[i]+conj(r[j]))
  end
  Σs = copy(Σs_buf)
  return Σs
end

statcovz(eq::SdeRoots) = statcovz(eq,eigenvecA(eq))

"Complex to real conversion tolerance"
const def_tol = 1e-8

function roots2sdepar(r,tol=0)
  T = real(eltype(r))
  p = coeffs(fromroots(r))[1:end-1]
  a = map(p) do pc
    if !(abs(imag(pc))≤tol)
      error("Imaginary part exceeds tolerance, abs(",imag(pc),") > ",tol)
    end
    T(real(pc))
  end
  return a
end

function Sde(eq::SdeRoots,tol=def_tol)
  a = roots2sdepar(eq.r,tol)
  return Sde(a,eq.σϵ)
end

SdeRoots(eq::SdeRoots) = eq

SdeRoots(m::WhiteNoise{T}) where {T} = SdeRoots(complex(T)[],m.σ)

SdeRoots(proc::SdeProcess) = SdeRoots(proc.eq)

order(eq::SdeRoots) = length(eq.r)

"print roots"
function printroots(r)
  rstr = map(r) do rr
    s = @sprintf "%0.3g" real(rr)
    if !isreal(rr)
      s *= @sprintf " ± 2π%0.3gim" abs(imag(rr))/2π
    end
    s
  end
  w = maximum(length,rstr)
  for (i,rr) in enumerate(r)
    any(conj(rr).==r[1:i-1]) && continue
    rc = rstr[i]
    p = repeat(' ',w-length(rc))
    irc = @sprintf "-1/%0.3g" -1/real(rr)
    if !isreal(rr)
      irc *= @sprintf " ± (2π/%0.3g)im" 2π/abs(imag(rr))
    end
    s = "$rc$p = $irc"
    if real(rr)>0
      printstyled(s,"\n",color=:yellow)
    else
      println(s)
    end
  end
end

"compact root string"
function rootstr(x)
  return if isreal(x)
    @sprintf "%0.3g" real(x)
  else
    @sprintf "%0.3g%+0.3gim" real(x) imag(x)
  end
end

"
$SIGNATURES

Truncate real part of roots to `-R`.
"
function truncate_roots(r,R;warn=true)
  c = 0
  rt = map(r) do x
    a = real(x)
    return if a<-R
      c+=1
      isreal(r) ? -R : Complex(-R,imag(x))
    else
      x
    end
  end
  if c>0 && warn
    printroots(r)
    @warn "Truncated $c roots"
  end
  return rt
end


text:
using Documenter, AlphaVantage

makedocs(
    sitename = "AlphaVantage.jl Documentation"
)

deploydocs(
    deps   = Deps.pip("mkdocs"),
    repo   = "github.com/ellisvalentiner/AlphaVantage.jl.git"
)


text:
@inline function with_temp_dir(f)
    original_dir = pwd()
    temp_dir = mktempdir(; cleanup = true)
    cd(temp_dir)
    result = f(temp_dir)
    cd(original_dir)
    rm(temp_dir; force = true, recursive = true)
    return result
end

@inline function with_cloned_repo(f, url)
    result = with_temp_dir() do temp_dir
        cd(temp_dir)
        run(`git clone $(url) CLONED_REPO`)
        cd("CLONED_REPO")
        return f(pwd())
    end
    return result
end


text:
using Parquet
using Test

function get_parcompat(parcompat=joinpath(dirname(@__FILE__), "parquet-compatibility"))
    # look for parquet-compatibility in test folder, clone to tempdir if not found
    if !isdir(parcompat)
        run(`git clone https://github.com/Parquet/parquet-compatibility.git $parcompat`)
    end
end

function get_juliaparcompat(juliaparcompat=joinpath(dirname(@__FILE__), "julia-parquet-compatibility"))
    # look for julia-parquet-compatibility in test folder, clone to tempdir if not found
    if !isdir(juliaparcompat)
        run(`git clone https://github.com/JuliaIO/parquet-compatibility.git $juliaparcompat`)
    end
end

get_parcompat()
get_juliaparcompat()


text:
module SequenceSpaceRepresentationScript

#nprocs()==1 && addprocs() # THREADING

using SynapseClient
using DISSEQT
using DISSEQT.SynapseTools
using DISSEQT.Plots
using DISSEQT.AnnotatedArrays
using LinearAlgebra
using DataFrames
using CSV
using Statistics
using JLD
using Gadfly
using Colors

function cachemetadata(syn,alignmentFolder,metadataID)
    metadataCacheFilename = "metadatacache.csv"
    metadata = DataFrame()
    if isfile(metadataCacheFilename)
        metadata = CSV.read(metadataCacheFilename; missingstrings=["","NA"])
    else
        # Get metadata (metadataID can be the synapse ID of a .csv file or of a folder containing .csv files).
        metadata = getsamplemetadata(syn, metadataID, :Rejected=>"No") # Always remove rejected samples

        # extend metadata
        appendswarmids!(syn,metadata,alignmentFolder)
        appendalignedids!(syn,metadata,alignmentFolder)

        # download swarms and consensus sequences
        downloadswarms!(syn,metadata,cacheDir="MutantSwarms")
        downloadconsensuses!(syn,metadata,cacheDir="MutantSwarms")

        CSV.write(metadataCacheFilename, metadata; missingstring="NA")
    end
    metadata
end

function main()
    syn = SynapseClient.login()

    # projectFolder should point to MyProject folder (containing subfolders Metadata and Analysis)
    projectFolder   = "syn11639899" # VignuzziLabPublic/Projects/FitnessLandscapes
    analysisFolder  = getchildbyname(syn, projectFolder, "Analysis")
    alignmentFolder = getchildbyname(syn, analysisFolder, "Alignment")
    referenceFolder = getchildbyname(syn, alignmentFolder, "ReferenceGenomes")

    isdir("MutantSwarms") || mkdir("MutantSwarms")
    isdir("plots") || mkdir("plots")

    # set to true to upload files
    doUpload = true
    uploadName = "fitness landscapes invitro"

    metadataID = getchildbyname(syn, projectFolder, "Metadata", "fitness landscapes invitro.csv")
    metadata = cachemetadata(syn, alignmentFolder, metadataID)
    metadata = metadata[ .~ismissing.(metadata[:Dose]), : ] # skip samples with unknown dose

    # Get Reference Genomes
    referenceNames = ["WT", "Less", "More", "Stop"]
    referenceIDs, referenceGenomes = getreferencegenomes(syn, referenceFolder, referenceNames)

    # get CVB3 feature list
    featureFileID = getchildbyname(syn,referenceFolder,"CVB3_features.csv")
    features = CSV.read(localpath(syn, featureFileID); missingstrings=["","NA"])
    positions = featurepositions(features, "ORF") # coding positions

    # load swarms
    swarms  = loadswarm(metadata[:SampleID],metadata[:SwarmPath],metadata[:ConsensusPath])

    # filter by positions
    swarms = swarms[:,:,:,inset(:position,positions),:]

    coverage = dropdims(swarms[:coverage]; dims=(1,2,3))
    meanCoverage = mean(coverage; dims=1)[:]

    # filter by mean read coverage at coding sites
    sampleMask = meanCoverage .>= 1000
    println("Removed ", count(.~sampleMask), " samples due to low coverage.")
    swarms  = swarms[:,:,:,:,sampleMask]
    metadata = metadata[sampleMask,:]

    # remove mixed samples
    mixes = estimatemix(swarms, referenceGenomes)
    mixMask = maximum(mixes; dims=2)[:] .>= 0.98
    println("Removed ", count(.~mixMask), " mixed samples.")
    swarms  = swarms[:,:,:,:,mixMask]
    metadata = metadata[mixMask,:]

    println("Total number of samples used: ", size(swarms,5))

    nbrDims = 13 # same as for the SMSSVD representation

    # setup dependencies
    dependencies = unique(vcat(featureFileID, referenceIDs, metadata[:SwarmID], metadata[:MetadataID]))

    # get rid of annotations to simplify math stuff
    # arrays are 4x4x4xPxN - nuc3 x nuc2 x nuc1 x positions x samples
    s = convert(Array, swarms)
    s = reshape(s, 64, size(swarms,4), size(swarms,5)) # reshape to VxN - variables x samples

    # Convert to consensus (1-of-K) encoding
    M = maximum(s; dims=1)
    M[M.<0.25] .= 1. # if no variant is above 0.25, set the maximum to 1 which will force all variables to zero.
    X = float(s.==M) # 1 if equals to the maximum and 0 otherwise

    # reshape to VxN - variables x samples
    X = reshape(X, 64*size(X,2), size(X,3))

    # center variables
    X = X .- mean(X; dims=2) # remove mean over samples

    # PCA keeping only nbrDims dimensions
    N = size(X,2)
    F = eigen(Symmetric(X'X), N-nbrDims+1:N)
    Σ = sqrt.(max.(0.,reverse(F.values)))
    V = F.vectors[:,end:-1:1]
    U = X*V ./ Σ'

    uploads = UploadList()

    # Low-Rank Representation, Sample names, positions and codon list
    seqSpaceRep = "plots/seqspacerepresentationConsensusPCA.jld"
    save(seqSpaceRep, "U",U, "Σ",Σ, "V",V,
                      "SampleID", convert(Vector{String},metadata[:SampleID]),
                      "positions", swarms[:position][:],
                      "codons",    swarms[:codon][:],
                      "consensus", dropdims(swarms[:consensus]; dims=(1,2,3)) )
    markforupload!(uploads, "SequenceSpaceRepresentation", seqSpaceRep, dependencies)

    # make PCA plots

    virusColors = Dict("WT"  =>colorant"black",
                       "Less"=>colorant"blue",
                       "More"=>colorant"green",
                       "Stop"=>colorant"red")
    virusColorMap = groupcolors(metadata[:Reference],virusColors)

    mutagenColors = Dict("5FU" =>RGB(1,0,0),
                         "AML" =>RGB(0,1,0),
                         "AZC" =>RGB(0,0,1),
                         "Mn"  =>RGB(0.8,0.8,0),
                         "Mock"=>RGB(1,0,1),
                         "RIBA"=>RGB(0,1,1))
    mutagenColorMap = groupcolors(metadata[:Mutagen],mutagenColors)

    P = pairwisescatterplot(V, metadata[:Reference], virusColorMap,
                               metadata[:Mutagen],   mutagenColorMap,
                               point_size=0.6mm)
    # markforupload!(uploads, "SequenceSpacePlots", saveplot(P, [:png,:svg,:pdf], "plots/pairwiseConsensusPCA"), dependencies)
    markforupload!(uploads, "SequenceSpacePlots", saveplot(P, [:png,:pdf], "plots/pairwiseConsensusPCA"), dependencies)

    doUpload && uploadfiles(syn, createchildfolder(syn, analysisFolder, "SequenceSpace", uploadName), uploads, @__FILE__, "Sequence Space Representation")

    # swarms, metadata, X, U, Σ, V
    nothing
end

end
SequenceSpaceRepresentationScript.main();


text:
using Plots
using Unitful
using UnitfulRecipes
using Dates
using UUIDs
using MPIMeasurements
#pyplot()
#plotly()
#inspectdr()
default(show = true)

ENV["JULIA_DEBUG"] = "all"

# Add test configurations to path
testConfigDir = normpath(string(@__DIR__), "../test/TestConfigs")
addConfigurationPath(testConfigDir)

study = MDFv2Study(
  description = "n.a.",
  name = "My simulated study",
  number = 1,
  time = now(),
  uuid = UUIDs.uuid4()
)

experiment = MDFv2Experiment(;
  description = "n.a.",
  isSimulation = true,
  name = "My simulated experiment",
  number = 1,
  subject = "Phantom of the Opera",
  uuid = UUIDs.uuid4()
)

scannerName_ = "TestRedPitayaScanner"
protocolName_ = "RedPitayaDAQMeasurement"
operator = "Jonas"
filename = joinpath("./tmp.mdf")

protocol = Protocol(protocolName_, scannerName_)
prepareMDF(protocol, filename, study, experiment, operator)
runProtocol(protocol)

data = protocol.mdf.measurement.data
#dataRMS = mapslices(x -> sqrt(1/length(x)*sum(x.^2)), data, dims = 1)
#plot(reshape(dataRMS, length(dataRMS)))
plot(reshape(data, length(data)))


text:
module AEMS

using POMDPs

using FIB: FIBSolver
using BeliefUpdaters
using POMDPModelTools
using POMDPPolicies
using POMDPTesting: TestSimulator
import POMDPs: Solver, Policy
import POMDPs: solve, action, value, update, initialize_belief, updater
import POMDPs: simulate

export
    AEMSSolver,
    AEMSPlanner,
    solve,
    action,
    value,

    clear_graph!,
    count_nodes,
    update_root

include("graph.jl")
include("bounds.jl")

include("solver.jl")
include("action.jl")

include("simulate.jl")

include("visualization.jl")

export visualize

end # module


text:
"""
    waterfall(x, y)

Creates an `EChart` representing the increase/decrease between two consecutive points.

## Methods
```julia
waterfall(x::AbstractVector, y::AbstractVector{<:Real})
```

## Arguments
* `legend::Bool` : display legend?
* `scale::Bool = false` : show full Y-axis or truncated
* `kwargs` : varargs to set any field of resulting `EChart` struct

## Notes

Reasonable defaults set for different methods of `area`, such as displaying a legend when two or more series present.

## Examples
```julia
x = 1:5
y = [2900, -1200, -300, -200, -900]
w = waterfall(x, y)
```
"""
function waterfall(x::AbstractVector, y::AbstractVector{<:Real};
 		legend::Bool = false,
 		scale::Bool = false,
 		kwargs...)

    #Need to add a value for total, since user passes in data values only
    labels = [string(x) for x in x]
    push!(labels, "total")

    #Calculate transparent series as base for stacking
    bottom = cumsum(y)
    bottom[1] = 0
    push!(bottom, 0)

    #Calculate actual waterfall
    top = abs.(y)
    push!(top, sum(y))

    ec = bar(labels, hcat(bottom, top), stack = true, legend = legend, scale = scale, kwargs...)
    #Make bottom series transparent
    ec.series[1].itemStyle = trans = ItemStyle(barBorderColor = "transparent", color = "transparent")

    return ec
end


text:
# ---
# title: 231. Power of Two
# id: problem231
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Math, Bit Manipulation
# link: <https://leetcode.com/problems/power-of-two/description/>
# hidden: true
# ---
#
# Given an integer `n`, return _`true` if it is a power of two. Otherwise,
# return `false`_.
#
# An integer `n` is a power of two, if there exists an integer `x` such that `n
# == 2x`.
#
#
#
# **Example 1:**
#
#
#
#     Input: n = 1
#     Output: true
#     Explanation: 20 = 1
#
#
# **Example 2:**
#
#
#
#     Input: n = 16
#     Output: true
#     Explanation: 24 = 16
#
#
# **Example 3:**
#
#
#
#     Input: n = 3
#     Output: false
#
#
# **Example 4:**
#
#
#
#     Input: n = 4
#     Output: true
#
#
# **Example 5:**
#
#
#
#     Input: n = 5
#     Output: false
#
#
#
#
# **Constraints:**
#
#   * `-231 <= n <= 231 - 1`
#
#
## @lc code=start
using LeetCode

function is_power_of_two(n::Int)
    return n > 0 && (n & (n - 1)) == 0
end
## @lc code=end