import dolfinx.fem as fem
from dolfinx.fem.petsc import LinearProblem
# from dolfinx.io import XDMFFile
import dolfinx.io as io

from dolfinx.io.gmshio import model_to_mesh
from dolfinx.mesh import compute_midpoints, locate_entities_boundary
from dolfinx.plot import vtk_mesh

# from ufl import TestFunction, TrialFunction, as_vector, dot, dx, grad, inner
import ufl

from mpi4py import MPI
from petsc4py import PETSc

import gmsh
import numpy as np
import pyvista
import yaml
import tqdm.autonotebook

# Personal libraries
import set_bc
from set_bc import BoundaryConditions
from weak_forms.wf_pnp import PoissonNernstPlanck
from weak_forms.wf_ns import NavierStokes
import utils


rank = MPI.COMM_WORLD.rank


# //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
def scaling_fun(fun, scaled_fun, scale_factor):
    with fun.vector.localForm() as loc, scaled_fun.vector.localForm() as loc_scaled:
        loc.copy(loc_scaled)
        loc_scaled.scale(scale_factor)
    return scaled_fun


class FunctionalSpacePNPNS:
    def __init__(self, nondim=True):

        # READING dolfinx mesh, cell tags, facet tags
        self.file_to_read = "../data/mesh/mesh_1.msh"
        self.mesh_dx, self.cell_tags, self.facet_tags = io.gmshio.read_from_msh(
            self.file_to_read, MPI.COMM_WORLD, 0, gdim=2)

        # DEF. dt (time) variation
        self.dt = 1.0e-3

        # DEF. ds measure and facet normal
        self.dt = 1.0e-3
        self.ds = ufl.Measure("ds", domain=self.mesh_dx, subdomain_data=self.facet_tags)
        self.n_facet = ufl.FacetNormal(self.mesh_dx)

        # DEF. FunctionSpaces
        # s_, v_: scalar, vectorial
        # cg: Continuous Garlekin (method)
        # 1, 2: orders of Lagrange Polynomials
        # ---------------------------------------------------------------- #
        s_cg1 = ufl.FiniteElement("Lagrange", self.mesh_dx.ufl_cell(), 1)
        v_cg2 = ufl.VectorElement("Lagrange", self.mesh_dx.ufl_cell(), 2)
        # -----------------------------------------------------------------#
        self.space_elpot = fem.FunctionSpace(self.mesh_dx, s_cg1)
        self.space_vel = fem.FunctionSpace(self.mesh_dx, v_cg2)
        self.space_press = fem.FunctionSpace(self.mesh_dx, s_cg1)
        self.space_conc = fem.FunctionSpace(self.mesh_dx, s_cg1)

        # DEF. TrialFunctions
        self.elpot = ufl.TrialFunction(self.space_elpot)
        self.vel = ufl.TrialFunction(self.space_vel)
        self.press = ufl.TrialFunction(self.space_press)
        self.conc = ufl.TrialFunction(self.space_conc)

        # DEF. TestFunctions
        self.vtest_elpot = ufl.TestFunction(self.space_elpot)
        self.vtest_vel = ufl.TestFunction(self.space_vel)
        self.vtest_press = ufl.TestFunction(self.space_press)
        self.vtest_conc = ufl.TestFunction(self.space_conc)

        # DEF. Weak Formulation of PNPNS
        self.bil_form_poisson = 0
        self.lin_form_poisson = 0

        self.bil_form_navier_stokes_1 = 0
        self.lin_form_navier_stokes_1 = 0
        self.bil_form_navier_stokes_2 = 0
        self.lin_form_navier_stokes_2 = 0
        self.bil_form_navier_stokes_3 = 0
        self.lin_form_navier_stokes_3 = 0

        self.bil_form_nernst_planck = 0
        self.lin_form_nernst_planck = 0

        # Read the YAML file to import mesh parameters
        with open('prms_pnpns.yaml', 'r') as file:
            yaml_data = yaml.safe_load(file)

        # fill attributes of the class with yaml file parameters
        self.fl_den = yaml_data['PNPNS_PRMS']['flow_density']
        self.fl_dyn_vis = yaml_data['PNPNS_PRMS']['flow_dynamic_viscosity']
        self.dif_coeff_ics = yaml_data['PNPNS_PRMS']['dif_coeff_ics']
        self.dif_coeff_mem = yaml_data['PNPNS_PRMS']['dif_coeff_mem']
        self.dif_coeff_ecs = yaml_data['PNPNS_PRMS']['dif_coeff_ecs']
        self.boltz = yaml_data['PNPNS_PRMS']['boltzmann_const']
        self.temp = yaml_data['PNPNS_PRMS']['temp_kelvin']
        self.el_charge = yaml_data['PNPNS_PRMS']['electric_charge']
        self.diel_perm = yaml_data['PNPNS_PRMS']['dielectric_permittivity']
        self.z_ion = yaml_data['PNPNS_PRMS']['valence']

        self.len_ref = 1.0
        self.dif_coeff_ref = self.dif_coeff_ecs
        self.t_ref = pow(self.len_ref, 2) / self.dif_coeff_ref
        self.elpot_ref = self.boltz * self.temp / (self.z_ion * self.el_charge)
        self.vel_ref = self.dif_coeff_ref / self.len_ref
        self.press_ref = self.fl_dyn_vis * self.dif_coeff_ref / pow(self.len_ref, 2)
        self.conc_ref = 144
        self.flux_ref = self.conc_ref * self.dif_coeff_ref / self.len_ref

        # calculate the Debye's length, Schmidt's number, ant the elctric coupling constant
        self.len_debye = np.sqrt(self.diel_perm * self.boltz * self.temp
                                 / (2.0 * pow(self.el_charge * self.z_ion, 2) * self.conc_ref)
                                 )
        self.schmidt_num = self.fl_dyn_vis / (self.fl_den * self.dif_coeff_ref)
        self.eldyn_coupl_const = self.diel_perm / (self.fl_dyn_vis * self.dif_coeff_ref) \
                                 * pow(self.boltz * self.temp / (self.el_charge * self.z_ion), 2)

        if nondim is True:

            self.dt /= self.t_ref
            self.len_debye_adim = self.len_debye / self.len_ref

    def read_mesh_xdmf(self):
        # READING mesh.xdmf
        with io.XDMFFile(self.mesh_dx.comm, self.file_to_read, "r") as xdmf:
            mesh_dx = xdmf.read_mesh(name="Grid")
            cell_tags = xdmf.read_meshtags(mesh_dx, name="Grid")
        mesh_dx.topology.create_connectivity(mesh_dx.topology.dim, mesh_dx.topology.dim - 1)
        with io.XDMFFile(self.mesh_dx.comm, self.file_to_read, "r") as xdmf:
            facet_tags = xdmf.read_meshtags(mesh_dx, name="Grid")
        return mesh_dx, cell_tags, facet_tags


class SystemSolver:
    def __init__(self, mesh_tags, functional_space_config):

        # Import parent class attributes
        self.mtg = mesh_tags
        self.cfg_f = functional_space_config
        self.path = "../data/sol/"

        # DEF. number of steps of the simulation
        self.num_steps = 5

        # DEF. solvers (PETSc)
        self.solver_poisson = PETSc.KSP().create(self.cfg_f.mesh_dx.comm)
        self.solver_navier_stokes_1 = PETSc.KSP().create(self.cfg_f.mesh_dx.comm)
        self.solver_navier_stokes_2 = PETSc.KSP().create(self.cfg_f.mesh_dx.comm)
        self.solver_navier_stokes_3 = PETSc.KSP().create(self.cfg_f.mesh_dx.comm)
        self.solver_nernst_planck = PETSc.KSP().create(self.cfg_f.mesh_dx.comm)

        # DEF. system convergence and solver parameters
        self.TOL_SYS_CONV = 1.0e-05
        self.MAX_ITER_SYS_CONV = 2000
        self.ABS_TOL_SOLVER = 1.0e-12
        self.REL_TOL_SOLVER = 1.0e-12

        # Tool for visualizing the progress bar of the simulation
        self.progress = tqdm.autonotebook.tqdm(desc="Solving PDE", total=self.num_steps)

    def solve_pnpns_system(self, nondim=True):
        # DEF. Functions at step n
        elpot_n = fem.Function(self.cfg_f.space_elpot, name="elpot_n")
        vel_n = fem.Function(self.cfg_f.space_vel, name="vel_n")
        press_n = fem.Function(self.cfg_f.space_press, name="press_n")
        conc_n = fem.Function(self.cfg_f.space_conc, name="conc_n")

        # DEF. Functions candidate for the step n+1
        vel_tent = fem.Function(self.cfg_f.space_vel, name="vel_tent")
        press_tent = fem.Function(self.cfg_f.space_press, name="press_tent")
        conc_tent = fem.Function(self.cfg_f.space_conc, name="conc_tent")

        # DEF. Functions store solution of the precedent (n)-steps and check convergence in j-cycles
        elpot_old = fem.Function(self.cfg_f.space_elpot, name="elpot_old")
        vel_old = fem.Function(self.cfg_f.space_vel, name="vel_old")
        press_old = fem.Function(self.cfg_f.space_press, name="press_old")
        conc_old = fem.Function(self.cfg_f.space_conc, name="conc_old")

        # BOUNDARY CONDITIONS ::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
        bcs_class = set_bc.BoundaryConditions(mesh_tags=self.mtg, functional_space_config=self.cfg_f)
        bcs_elpot_list, bcs_vel_list, bcs_press_list, bcs_conc_list = bcs_class.set_pnpns_bcs()

        # We can now loop through the boundary condition and append them to L(v)
        # or the list of Dirichlet boundary conditions
        bcs_elpot = []
        for condition in bcs_elpot_list:
            if condition.type == "Dirichlet":
                bcs_elpot.append(condition.bc)
            else:
                self.cfg_f.lin_form_poisson += condition.bc

        bcs_vel = []
        for condition in bcs_vel_list:
            if condition.type == "Dirichlet":
                bcs_vel.append(condition.bc)
            else:
                self.cfg_f.lin_form_navier_stokes_1 += condition.bc

        bcs_press = []
        for condition in bcs_press_list:
            if condition.type == "Dirichlet":
                bcs_press.append(condition.bc)
            else:
                self.cfg_f.lin_form_navier_stokes_1 += condition.bc

        bcs_conc = []
        for condition in bcs_elpot_list:
            if condition.type == "Dirichlet":
                bcs_conc.append(condition.bc)
            else:
                self.cfg_f.lin_form_nernst_planck += condition.bc

        # INITIAL CONDITIONS :::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
        #
        #
        # ..............................................................................................................

        # fill the tent vector.localForm
        with vel_n.vector.localForm() as loc_n, vel_tent.vector.localForm() as loc_tent:
            loc_n.copy(loc_tent)

        with press_n.vector.localForm() as loc_n, press_tent.vector.localForm() as loc_tent:
            loc_n.copy(loc_tent)

        with conc_n.vector.localForm() as loc_n, conc_tent.vector.localForm() as loc_tent:
            loc_n.copy(loc_tent)

        # Call the class constructor of the Poisson, Nernst-Planck, and Navier-Stokes variational formulaiton
        pnp = PoissonNernstPlanck(self.cfg_f)
        ns = NavierStokes(self.cfg_f)

        # Construct variational formulation of Poisson
        pnp.weak_form_poisson_adim(conc_tent)

        # Assemble matrix of the bil_form_poisson (outside) and create vector of the lin_form_poisson
        bf_p = fem.form(self.cfg_f.bil_form_poisson)
        lf_p = fem.form(self.cfg_f.lin_form_poisson)
        mat_poi = fem.petsc.assemble_matrix(bf_p)
        mat_poi.assemble()
        vec_poi = fem.petsc.create_vector(lf_p)

        # Construct variational formulation of Navier-Stokes
        ns.weak_form_navier_stokes_adim(elpot_n, vel_tent, vel_tent, press_n, press_tent, conc_n)

        # Create matrix of the bil_form_navier_stokes_1 and create vector of the bil_form_navier_stokes_1
        bf_ns_1 = fem.form(self.cfg_f.bil_form_navier_stokes_1)
        lf_ns_1 = fem.form(self.cfg_f.lin_form_navier_stokes_1)
        mat_ns_1 = fem.petsc.create_matrix(bf_ns_1)
        vec_ns_1 = fem.petsc.create_vector(lf_ns_1)

        # Assemble the matrices of the pressure correction (2) and velocity correction (3)
        # related to the incremental pressure correction scheme (IPCS) used for solving the Navier-Stokes eqautions
        # mat_ns_2 = fem.petsc.assemble_matrix(self.cfg_f.bil_form_navier_stokes_2, bcs=bcs_press)
        bf_ns_2 = fem.form(self.cfg_f.bil_form_navier_stokes_2)
        lf_ns_2 = fem.form(self.cfg_f.lin_form_navier_stokes_2)
        mat_ns_2 = fem.petsc.assemble_matrix(bf_ns_2)
        mat_ns_2.assemble()
        vec_ns_2 = fem.petsc.create_vector(lf_ns_2)

        bf_ns_3 = fem.form(self.cfg_f.bil_form_navier_stokes_3)
        lf_ns_3 = fem.form(self.cfg_f.lin_form_navier_stokes_3)
        mat_ns_3 = fem.petsc.assemble_matrix(bf_ns_3)
        mat_ns_3.assemble()
        vec_ns_3 = fem.petsc.create_vector(lf_ns_3)

        # Construct variational formulation of Nernst-Planck
        pnp.weak_form_nernst_planck_adim(elpot_n, vel_tent, conc_n)

        # Assemble matrix of the bil_form_nernst_planck and create vector of the lin_form_nernst_planck
        bf_np = fem.form(self.cfg_f.bil_form_nernst_planck)
        lf_np = fem.form(self.cfg_f.lin_form_nernst_planck)
        mat_np = fem.petsc.create_matrix(bf_np)
        vec_np = fem.petsc.create_vector(lf_np)

        # Use PETSc as a linear algebra backend
        self.solver_poisson.setOperators(mat_poi)
        self.solver_poisson.setType(PETSc.KSP.Type.PREONLY)  # To check
        pc_poi = self.solver_poisson.getPC()
        pc_poi.setType(PETSc.PC.Type.LU)  # To check

        self.solver_navier_stokes_1.setOperators(mat_ns_1)
        self.solver_navier_stokes_1.setType(PETSc.KSP.Type.BCGS)
        pc_ns_1 = self.solver_navier_stokes_1.getPC()
        pc_ns_1.setType(PETSc.PC.Type.JACOBI)

        self.solver_navier_stokes_2.setOperators(mat_ns_2)
        self.solver_navier_stokes_2.setType(PETSc.KSP.Type.MINRES)
        pc_ns_2 = self.solver_navier_stokes_1.getPC()
        pc_ns_2.setType(PETSc.PC.Type.HYPRE)
        pc_ns_2.setHYPREType("boomeramg")

        self.solver_navier_stokes_3.setOperators(mat_ns_3)
        self.solver_navier_stokes_3.setType(PETSc.KSP.Type.CG)
        pc_ns_1 = self.solver_navier_stokes_1.getPC()
        pc_ns_1.setType(PETSc.PC.Type.SOR)

        self.solver_poisson.setOperators(mat_np)
        self.solver_poisson.setType(PETSc.KSP.Type.CGS)
        pc_poi = self.solver_poisson.getPC()
        pc_poi.setType(PETSc.PC.Type.JACOBI)

        ################################################################################################################
        # SOLVE CHECKING CONVERGENCE
        ################################################################################################################
        # Set the initial (real) time of the nernst-planck-navier-stokes process to simulate
        t = 0

        # Create XDMF/VTK Files to store solution data
        xdmf_elpot = io.XDMFFile(MPI.COMM_WORLD, self.path + 'elpot/elpot.xdmf', "w")
        vtk_vel = io.VTKFile(MPI.COMM_WORLD, self.path + 'vel/vel.vtk', "w")
        xdmf_press = io.XDMFFile(MPI.COMM_WORLD, self.path + 'press/press.xdmf', "w")
        xdmf_conc = io.XDMFFile(MPI.COMM_WORLD, self.path + 'conc/conc.xdmf', "w")

        # Write mesh
        xdmf_elpot.write_mesh(self.cfg_f.mesh_dx)
        vtk_vel.write_mesh(self.cfg_f.mesh_dx)
        xdmf_press.write_mesh(self.cfg_f.mesh_dx)
        xdmf_conc.write_mesh(self.cfg_f.mesh_dx)

        # Write initial data (n = 0)
        xdmf_elpot.write_function(elpot_n, t)
        vtk_vel.write_function(vel_tent, t)
        xdmf_press.write_function(press_tent, t)
        xdmf_conc.write_function(conc_tent, t)

        if nondim is True:
            # DEF. Functions to store dimensional value
            elpot_dim = fem.Function(self.cfg_f.space_elpot, name="elpot_dim")
            vel_dim = fem.Function(self.cfg_f.space_vel, name="vel_dim")
            press_dim = fem.Function(self.cfg_f.space_press, name="press_dim")
            conc_dim = fem.Function(self.cfg_f.space_conc, name="conc_dim")

            xdmf_elpot_dim = io.XDMFFile(MPI.COMM_WORLD, self.path + 'elpot_dim/elpot_dim.xdmf', "w")
            vtk_vel_dim = io.VTKFile(MPI.COMM_WORLD, self.path + 'vel_dim/vel_dim.vtk', "w")
            xdmf_press_dim = io.XDMFFile(MPI.COMM_WORLD, self.path + 'press_dim/press_dim.xdmf', "w")
            xdmf_conc_dim = io.XDMFFile(MPI.COMM_WORLD, self.path + 'conc_dim/conc_dim.xdmf', "w")

            xdmf_elpot_dim.write_mesh(self.cfg_f.mesh_dx)
            vtk_vel_dim.write_mesh(self.cfg_f.mesh_dx)
            xdmf_press_dim.write_mesh(self.cfg_f.mesh_dx)
            xdmf_conc_dim.write_mesh(self.cfg_f.mesh_dx)

            t_dim = t * self.cfg_f.t_ref
            elpot_dim = scaling_fun(elpot_n, elpot_dim, self.cfg_f.elpot_ref)
            vel_dim = scaling_fun(vel_tent, vel_dim, self.cfg_f.elpot_ref)
            press_dim = scaling_fun(press_tent, press_dim, self.cfg_f.elpot_ref)
            conc_dim = scaling_fun(conc_tent, conc_dim, self.cfg_f.elpot_ref)

            xdmf_elpot_dim.write_function(elpot_dim, t_dim)
            vtk_vel_dim.write_function(vel_dim, t_dim)
            xdmf_press_dim.write_function(press_dim, t_dim)
            xdmf_conc_dim.write_function(conc_dim, t_dim)

        print("\n")
        # . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
        for n in range(1, self.num_steps + 1):
            self.progress = tqdm.autonotebook.tqdm(desc="Solving PDE", total=self.num_steps)
            # Update time step
            t += self.cfg_f.dt

            for j in range(1, self.MAX_ITER_SYS_CONV + 1):

                # save old values
                with elpot_n.vector.localForm() as loc, elpot_old.vector.localForm() as loc_old:
                    loc.copy(loc_old)

                with vel_tent.vector.localForm() as loc, vel_old.vector.localForm() as loc_old:
                    loc.copy(loc_old)

                with press_tent.vector.localForm() as loc, press_old.vector.localForm() as loc_old:
                    loc.copy(loc_old)

                with conc_tent.vector.localForm() as loc, conc_old.vector.localForm() as loc_old:
                    loc.copy(loc_old)

                # solve Poisson equation * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
                pnp.weak_form_poisson_adim(conc_tent)
                bf_p = fem.form(self.cfg_f.bil_form_poisson)
                lf_p = fem.form(self.cfg_f.lin_form_poisson)

                mat_poi = fem.petsc.assemble_matrix(bf_p)
                mat_poi.assemble()
                vec_poi = fem.petsc.create_vector(lf_p)
                self.solver_poisson.setOperators(mat_poi)
                self.solver_poisson.setType(PETSc.KSP.Type.PREONLY)  # To check
                pc_poi = self.solver_poisson.getPC()
                pc_poi.setType(PETSc.PC.Type.LU)  # To check

                with vec_poi.localForm() as loc:
                    loc.set(0)

                fem.petsc.assemble_vector(vec_poi, lf_p)

                # Apply Dirichlet boundary condition to the vector
                fem.petsc.apply_lifting(vec_poi, [bf_p], [bcs_elpot])
                vec_poi.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
                fem.set_bc(vec_poi, bcs_elpot)

                # Solve the linear problem
                self.solver_poisson.solve(vec_poi, elpot_n.vector)
                elpot_n.x.scatter_forward()

                # solve Navier-Stokes * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *

                # IPCS incremental pressure correction scheme:
                ns.weak_form_navier_stokes_adim(elpot_n, vel_n, vel_tent, press_n, press_tent, conc_n)

                # Step 1: Tentative velocity
                bf_ns_1 = fem.form(self.cfg_f.bil_form_navier_stokes_1)
                lf_ns_1 = fem.form(self.cfg_f.lin_form_navier_stokes_1)
                mat_ns_1.zeroEntries()
                fem.petsc.assemble_matrix(mat_ns_1, bf_ns_1, bcs=bcs_vel)
                mat_ns_1.assemble()
                with vec_ns_1.localForm() as loc:
                    loc.set(0)
                fem.petsc.assemble_vector(vec_ns_1, lf_ns_1)
                fem.apply_lifting(vec_ns_1, [bf_ns_1], [bcs_vel])
                vec_ns_1.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
                fem.set_bc(vec_ns_1, bcs_vel)
                self.solver_navier_stokes_1.solve(vec_ns_1, vel_tent.vector)

                # Step 2: Pressure correction step
                bf_ns_2 = fem.form(self.cfg_f.bil_form_navier_stokes_2)
                lf_ns_2 = fem.form(self.cfg_f.lin_form_navier_stokes_2)
                with vec_ns_2.localForm() as loc:
                    loc.set(0)

                fem.petsc.assemble_vector(vec_ns_2, lf_ns_2)
                fem.apply_lifting(vec_ns_2, [bf_ns_2], [bcs_press])
                vec_ns_2.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
                fem.set_bc(vec_ns_2, bcs_press)
                self.solver_navier_stokes_2.solve(vec_ns_2, press_tent.vector)
                # todo here the dolfinx demo on Navier-Stokes DFG 2D-3 benchmark introduces novelties
                press_tent.x.scatter_forward()

                # Step 3: Velocity correction step
                bf_ns_3 = fem.form(self.cfg_f.bil_form_navier_stokes_3)
                lf_ns_3 = fem.form(self.cfg_f.lin_form_navier_stokes_3)
                with vec_ns_3.localForm() as loc:
                    loc.set(0)

                fem.petsc.assemble_vector(vec_ns_3, lf_ns_3)
                # fem.apply_lifting(vec_ns_3, [bf_ns_3], [bcs_press])  # todo is it necessary?
                vec_ns_3.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
                fem.set_bc(vec_ns_3, bcs_press)
                self.solver_navier_stokes_3.solve(vec_ns_3, vel_tent.vector)

                # solve Nernst-Planck * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *  * * * * * * * * * *
                pnp.weak_form_nernst_planck_adim(elpot_n, vel_tent, conc_n)
                bf_np = fem.form(self.cfg_f.bil_form_nernst_planck)
                lf_np = fem.form(self.cfg_f.lin_form_nernst_planck)
                mat_np.zeroEntries()
                mat_np = fem.petsc.assemble_matrix(bf_np)
                mat_np.assemble()
                with vec_np.localForm() as loc:
                    loc.set(0)
                fem.petsc.assemble_vector(vec_np, lf_np)
                fem.apply_lifting(vec_np, [bf_np], [bcs_conc])
                vec_poi.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
                self.solver_poisson.solve(vec_np, conc_tent.vector)

                # ......................................................................................................
                # CALCULUS OF THE RESIDUALS
                res_elpot = np.linalg.norm(elpot_n.x.array - elpot_old.x.array)
                res_vel = np.linalg.norm(vel_tent.x.array - vel_old.x.array)
                res_press = np.linalg.norm(press_tent.x.array - press_old.x.array)
                res_conc = np.linalg.norm(conc_tent.x.array - conc_old.x.array)

                utils.print_local_min_max_res(
                    n, j, elpot_n, res_elpot, vel_tent, res_vel, press_tent, res_press, conc_tent, res_conc
                )

                # Check for convergence
                if res_elpot < self.TOL_SYS_CONV and \
                        res_vel < self.TOL_SYS_CONV and \
                        res_press < self.TOL_SYS_CONV and \
                        res_conc < self.TOL_SYS_CONV:

                    break

                # END of the j-for-cycle (convergence cycle) ...........................................................

            xdmf_elpot.write_function(elpot_n, t)
            vtk_vel.write_function(vel_tent, t)
            xdmf_press.write_function(press_tent, t)
            xdmf_conc.write_function(conc_tent, t)

            if nondim is True:
                t_dim = t * self.cfg_f.t_ref
                elpot_dim = scaling_fun(elpot_n, elpot_dim, self.cfg_f.elpot_ref)
                vel_dim = scaling_fun(vel_tent, vel_dim, self.cfg_f.elpot_ref)
                press_dim = scaling_fun(press_tent, press_dim, self.cfg_f.elpot_ref)
                conc_dim = scaling_fun(conc_tent, conc_dim, self.cfg_f.elpot_ref)

                xdmf_elpot_dim.write_function(elpot_dim, t_dim)
                vtk_vel_dim.write_function(vel_dim, t_dim)
                xdmf_press_dim.write_function(press_dim, t_dim)
                xdmf_conc_dim.write_function(conc_dim, t_dim)

                elpot_n.vector.localForm.copy(elpot_dim)
                vel_tent.vector.localForm.copy(vel_dim)
                press_tent.vector.localForm.copy(press_dim)
                conc_tent.vector.localForm.copy(conc_dim)

                elpot_dim *= self.cfg_f.elpot_ref
                vel_dim *= self.cfg_f.vel_ref
                press_dim *= self.cfg_f.press_ref
                conc_dim *= self.cfg_f.conc_ref
                t_dim = t * self.cfg_f.t_ref

                # Write to xdmffile
                xdmf_elpot.write_function(elpot_dim, t_dim)
                vtk_vel.write_function(vel_dim, t_dim)
                xdmf_press.write_function(press_dim, t_dim)
                xdmf_conc.write_function(conc_dim, t_dim)

            # updating for the next dt-step
            vel_n.vector.localForm.copy(vel_tent)
            press_n.vector.localForm.copy(press_tent)
            conc_n.vector.localForm.copy(conc_tent)

            # END of the n-for-cycle ...................................................................................

        return

give me a feedback