#!/usr/bin/env python
# coding: utf-8

# ### Crameri2012Case1_FreeSurface
# 
# This is the case1 (featuring a cosine perturbation of the surface) from the reference (Crameri et al., 2012), results see Fig.2.
# 
# <center>
# Numerical setups used in this work. (motified from table 1 in (Crameri et al., 2012))
# </center>
# 
# |Top boundary|Grid points | Tracers   |
# | ----------- | ----------- |----------- |
# |Stick air | 560 $\times$ 320|36 part./elem. |Tracers  
# |Free surface  | 256 $\times$ 64  | 36 part./elem.|
# 
# ### Reference
# 
# - Crameri, F., Schmeling, H., Golabek, G. J., Duretz, T., Orendt, R., Buiter, S. J. H., ... & Tackley, P. J. (2012). A comparison of numerical surface topography calculations in geodynamic modelling: an evaluation of the ‘sticky air’method. Geophysical Journal International, 189(1), 38-54.

# In[1]:


import petsc4py
from petsc4py import PETSc

import underworld3 as uw
from underworld3.systems import Stokes
from underworld3 import function

import numpy as np
import sympy
import os
from datetime import datetime
import sys
from underworld3.cython.petsc_discretisation import petsc_dm_find_labeled_points_local

comm = uw.mpi.comm
rank = uw.mpi.rank
size = uw.mpi.size
if size == 1:
    import matplotlib.pyplot as plt


# In[2]:


u = uw.scaling.units
ndim = uw.scaling.non_dimensionalise
dim = uw.scaling.dimensionalise

# # build reference units
# KL = 700 * u.kilometer
# K_viscosity = 1e21  * u.pascal * u.second
# K_density   = 3300 * u.kilogram / u.meter**3

# # compute dependent scaling units 
# #KT_seconds = 1000*u.kiloyear
# #KM_kilograms = KT_seconds * KL_meters * K_viscosit


# KM = K_density * KL**3
# KT = KM / ( KL  * K_viscosity )

# scaling_coefficients = uw.scaling.get_coefficients()
# scaling_coefficients["[length]"]      = KL.to_base_units()
# scaling_coefficients["[time]"]        = KT.to_base_units()
# scaling_coefficients["[mass]"]        = KM.to_base_units()


H = 700.  * u.kilometer
velocity = 1.0 * u.centimeter / u.year
#g    =   10.0 * u.meter / u.second**2  
#bodyforce    = 3300  * u.kilogram / u.metre**3 * g 
mu           = 1e21  * u.pascal * u.second

KL = H
Kt = KL / velocity
KM = mu * KL * Kt

scaling_coefficients = uw.scaling.get_coefficients()
scaling_coefficients["[length]"] = KL
scaling_coefficients["[time]"] = Kt
scaling_coefficients["[mass]"]= KM


# In[3]:


#case_bc = "freeslip"  
#case_bc = "noslip"
case_bc = "freesurf"

max_time  = ndim(1e5*u.year)
dt_set    = ndim(5e2*u.year)
save_every = int(ndim(1e3*u.year)/dt_set)

xres,yres = 256,64
xmin, xmax = ndim(0 * u.kilometer), ndim(2800 * u.kilometer)
ymin, ymax = ndim(0 * u.kilometer), ndim(700 * u.kilometer)

ND_gravity = ndim(9.81 * u.meter / u.second**2)

outputPath = "op_Crameri2012Case1_" + case_bc + "_yres{:n}/".format(yres)
if uw.mpi.rank == 0:
    # delete previous model run
    if os.path.exists(outputPath):
        for i in os.listdir(outputPath):
            os.remove(outputPath+ i)
            
    ### create folder if not run before
    if not os.path.exists(outputPath):
        os.makedirs(outputPath)


# In[4]:


def plot_mesh(title,mesh,showFig=True):
    import numpy as np
    import pyvista as pv
    import vtk
    
    pv.global_theme.background = "white"
    pv.global_theme.window_size = [500, 500]
    pv.global_theme.jupyter_backend = "static"
    pv.global_theme.smooth_shading = True
    pv.global_theme.show_edges = True
    pv.global_theme.axes.show = True

    mesh.vtk("tmp_box_mesh.vtk")
    pvmesh = pv.read("tmp_box_mesh.vtk")
    pl = pv.Plotter()
    pl.add_mesh(pvmesh,'Black', 'wireframe')

    pl.add_title(title,font_size=11)
    if showFig:
        pl.show(cpos="xy")
    
    pl.screenshot(outputPath+title+".png",window_size=pv.global_theme.window_size,return_img=False) 
    pvmesh.clear_data() 
    pvmesh.clear_point_data()


def plot_meshswarm(title,mesh,swarm,material,showSwarm=False,showFig=True):
    import numpy as np
    import pyvista as pv
    import vtk
    
    pv.global_theme.background = "white"
    pv.global_theme.window_size = [500, 500]
    pv.global_theme.jupyter_backend = "static"
    pv.global_theme.smooth_shading = True
    pv.global_theme.show_edges = True
    pv.global_theme.axes.show = True

    mesh.vtk("tmp_box_mesh.vtk")
    pvmesh = pv.read("tmp_box_mesh.vtk")
    pl = pv.Plotter()
    pl.add_mesh(pvmesh,'Black', 'wireframe')

    if showSwarm:
        with swarm.access():
            points = np.zeros((swarm.data.shape[0],3))
            points[:,0] = swarm.data[:,0]
            points[:,1] = swarm.data[:,1]
        point_cloud = pv.PolyData(points)
        with swarm.access():
            point_cloud.point_data["M"] = material.data.copy()
        #pl.add_points(point_cloud, color="red",render_points_as_spheres=False, point_size=3, opacity=0.5)
        pl.add_mesh(point_cloud, cmap="coolwarm", edge_color="Black", show_edges=False, scalars='M',
                        use_transparency=False, opacity=0.95, point_size= 3)
    pl.add_title(title,font_size=11)
    if showFig:
        pl.show(cpos="xy")
    
    pl.screenshot(outputPath+title+".png",window_size=pv.global_theme.window_size,return_img=False) 
    pvmesh.clear_data()
    pvmesh.clear_point_data()


# In[5]:


from underworld3.utilities._api_tools import uw_object
from scipy.spatial import distance
class PopulationControl_DeformMesh(uw_object):
    def __init__(self, swarm):
        self._swarm = swarm
        with swarm.access():
            self._swarm_coords     = np.copy(np.ascontiguousarray( self._swarm.data))
            self._swarm_cellid     = np.copy(np.ascontiguousarray( self._swarm.particle_cellid.data[:,0]))
            self._particleNumber   = self._swarm.data.shape[0]
            self._mesh_cID, self._cIDcount  = np.unique(self._swarm_cellid, return_counts=True)
            self._swarm_ppc = self._cIDcount[0]

    def repopulate(self,mesh,updateField=None):
    # similar method as PopulationControl.redisturibute, 
    # but update the owning cell frist, using new swarm build on the deformed mesh to add particles, 
    # and delete the over_populated particles 
        self._updateField = updateField
        self._mesh = mesh
        
        # update the owning cell
        cellid = self._swarm.dm.getField("DMSwarm_cellid")
        coords = self._swarm.dm.getField("DMSwarmPIC_coor").reshape((-1, swarm.dim))
        cellid[:] = self._swarm.mesh.get_closest_cells(coords).reshape(-1)
        self._swarm.dm.restoreField("DMSwarmPIC_coor")
        self._swarm.dm.restoreField("DMSwarm_cellid")
        self._swarm.dm.migrate(remove_sent_points=True)
        self._swarm._index = None
        self._swarm._nnmapdict = {}

        self._updateField = updateField

        ### get the particle positions in the cell from the deformed mesh
        _newswarm = uw.swarm.Swarm(self._mesh)
        _newswarm.populate_petsc(fill_param=fill_parameter,layout=uw.swarm.SwarmPICLayout.GAUSS)
        with _newswarm.access():
            original_swarm_coords = np.copy(np.ascontiguousarray(_newswarm.data))

        ### get the current swarm positions in the cell
        with self._swarm.access():
            current_swarm_coords = self._swarm.data
            current_swarm_cells = self._swarm.particle_cellid.data[:,0]

        current_particle_cellID, current_ppc  = np.unique(current_swarm_cells, return_counts=True)

        ### find cells that are empty
        empty_cells = self._mesh_cID[np.isin(self._mesh_cID, current_particle_cellID, invert=True)]
        empty_cell_coords = original_swarm_coords[np.isin(self._swarm_cellid,  empty_cells)]

        ### find under-populated cells
        underpopulated_cells = current_particle_cellID[current_ppc < self._swarm_ppc]
        ### number of particles missing from each cell
        underpopulated_cell_coords = []

        for cell in underpopulated_cells:
            ### get the current number of particles in the cell
            current_particles = current_ppc[current_particle_cellID == cell][0]
            ### get the number of particles to add
            no_particles_to_add = self._swarm_ppc - current_particles

            original_cell_coords = original_swarm_coords[self._swarm_cellid == cell]
            current_cell_coords  = current_swarm_coords[current_swarm_cells == cell]

            ### caculate the distances between the original and current particle positions
            point_distance = distance.cdist(original_cell_coords, current_cell_coords, 'euclidean')
            ### add the ones with the largest distance
            underpopulated_cell_coords.append(original_cell_coords[np.min(point_distance, axis=1).argsort()[::-1]][:no_particles_to_add])
 

        if underpopulated_cell_coords:  # Check if the list is not empty
            underpopulated_cell_coords = np.concatenate(underpopulated_cell_coords)
        else:
            underpopulated_cell_coords = np.empty(shape=(0,2))  # Create an empty numpy array

        ### Combination of the empty cells and the underpopulated cells
        particles_to_add = np.vstack([underpopulated_cell_coords, empty_cell_coords])

        uw.mpi.barrier()
        new_cells = self._swarm.mesh.get_closest_local_cells(particles_to_add)
        
        if isinstance(new_cells, int):
            new_cells = np.array([new_cells])

        # Ensure that 'valid' has the same length as 'particles_to_add'
        if len(particles_to_add) != len(new_cells):
            new_cells = np.resize(new_cells, len(particles_to_add))

        valid = new_cells != -1

        valid_coords = particles_to_add[valid]
        valid_cells = new_cells[valid]

        all_local_coords = np.vstack([current_swarm_coords, valid_coords])

        all_local_cells = np.hstack([current_swarm_cells, valid_cells])

        ### do interp before adding particles, otherwise interp happens from the newly added particles
        if (self._updateField != None):
            ### defualt to the nnn value
            new_particle_mat = np.rint( self._updateField.rbf_interpolate(all_local_coords) )#, nnn=self._swarm_ppc))
        
        swarm_new_size = all_local_coords.data.shape[0]

        self._swarm.dm.addNPoints(swarm_new_size - current_swarm_coords.shape[0])

        cellid = self._swarm.dm.getField("DMSwarm_cellid")
        coords = self._swarm.dm.getField("DMSwarmPIC_coor").reshape((-1, self._swarm.dim))

        coords[...] = all_local_coords[...]
        cellid[:]   = all_local_cells[:]

        self._swarm.dm.restoreField("DMSwarmPIC_coor")
        self._swarm.dm.restoreField("DMSwarm_cellid")

        if (self._updateField != None):
            with self._swarm.access(self._updateField):
                self._updateField.data[:,0] = new_particle_mat[:,0]

        uw.mpi.barrier()
        
        ### get the current swarm positions in the cell
        with self._swarm.access():
            current_swarm_coords = self._swarm.data
            current_swarm_cells = self._swarm.particle_cellid.data[:,0]

        current_particle_cellID, current_ppc  = np.unique(current_swarm_cells, return_counts=True)

        ### find over-populated cells
        overpopulated_cells = current_particle_cellID[current_ppc > self._swarm_ppc]

        overpopulated_cell_coords = []

        for cell in overpopulated_cells:
            ### get the current particles in the cell
            current_particles = current_ppc[current_particle_cellID == cell][0]
            ### number of particles to delete
            no_particles_to_remove = current_particles - self._swarm_ppc
            ### randomly select particles to remove from the current cell
            #rng = np.random.default_rng()
    
            #### Selection of coords doesn't delete all particles (?)
            current_cell_coords  = current_swarm_coords[current_swarm_cells == cell]
            current_cell_coords_local = self._swarm.mesh._centroids[cell].reshape(1,-1)
            point_distance = distance.cdist(current_cell_coords,current_cell_coords_local, 'euclidean')
    
            np.min(point_distance, axis=1).argsort()[::-1][:int(no_particles_to_remove)]
            overpopulated_cell_coords.append(current_cell_coords[np.min(point_distance, axis=1).argsort()[::-1][:int(no_particles_to_remove)]])

        if overpopulated_cell_coords:  # Check if the list is not empty
            particles_to_remove = np.concatenate(overpopulated_cell_coords)
        else:
            particles_to_remove = np.empty(shape=(0,2))  # Create an empty numpy array

        # Check if all_current_coords and particles_to_remove are not empty
        with self._swarm.access(self._swarm.particle_coordinates):
            ### get current coords
            all_current_coords = self._swarm.data

            if particles_to_remove.size > 0:
                ### get number of particles to remove
                no_particles_to_remove = particles_to_remove.shape[0]

                    
                # Calculate point_distance
                point_distance = distance.cdist(all_current_coords, particles_to_remove, 'euclidean')

                # Get the index of particles to remove
                index_condition = np.min(point_distance, axis=1).argsort()[:no_particles_to_remove]

                # Set coords to 1.0e100 so they are deleted by the DM, following the same logic as the remeshing example
                self._swarm.data[index_condition] = 1.0e100

        return


# In[6]:


def _adjust_time_units(val):
    """ Adjust the units used depending on the value """
    if isinstance(val, u.Quantity):
        mag = val.to(u.years).magnitude
    else:
        val = dim(val, u.years)
        mag = val.magnitude
    exponent = int("{0:.3E}".format(mag).split("E")[-1])

    if exponent >= 9:
        units = u.gigayear
    elif exponent >= 6:
        units = u.megayear
    elif exponent >= 3:
        units = u.kiloyears
    elif exponent >= 0:
        units = u.years
    elif exponent > -3:
        units = u.days
    elif exponent > -5:
        units = u.hours
    elif exponent > -7:
        units = u.minutes
    else:
        units = u.seconds
    return val.to(units)


# In[7]:


from scipy.interpolate import interp1d
from scipy.interpolate import CloughTocher2DInterpolator

class FreeSurfaceProcessor(object): 
    def __init__(self,v,dt):
        """
        Parameters
        ----------
        _init_mesh : the original mesh
        mesh : the updating model mesh
        vel : the velocity field of the model
        dt : dt for advecting the surface
        """
        if mesh.dim == 2:
            self.init_mesh = uw.meshing.StructuredQuadBox(elementRes=(int(xres), int(yres)), minCoords=(xmin, ymin), maxCoords=(xmax, ymax))
        else:
            self.init_mesh = uw.meshing.StructuredQuadBox(elementRes=(int(xres), int(yres),int(zres)), minCoords=(xmin, ymin, zmin), maxCoords=(xmax, ymax, zmax))
        self.Tmesh = uw.discretisation.MeshVariable("Tmesh", self.init_mesh, 1, degree=1)
        self.Bmesh = uw.discretisation.MeshVariable("Bmesh", self.init_mesh, 1, degree=1)
        self.mesh_solver = uw.systems.Poisson(self.init_mesh , u_Field=self.Tmesh, solver_name="FreeSurf_solver")
        self.mesh_solver.constitutive_model = uw.constitutive_models.DiffusionModel
        self.mesh_solver.constitutive_model.Parameters.diffusivity = 1. 
        self.mesh_solver.f = 0.0
        self.mesh_solver.add_dirichlet_bc(self.Bmesh.sym[0], "Top",0)
        self.mesh_solver.add_dirichlet_bc(self.Bmesh.sym[0], "Bottom",0)

        self.v = v
        self._dt   = dt
        self.top = topwall
        self.bottom = botwall

    def _solve_sle(self):
        self.mesh_solver.solve()
     
    def _advect_surface(self):
        if self.top.size > 0:
            if mesh.dim == 2:
                coords = mesh.data[self.top]
                x = coords[:,0]
                y = coords[:,-1]
                vx = uw.function.evalf(self.v.sym[0], coords)
                vy = uw.function.evalf(self.v.sym[1], coords)
                
                # Advect top surface
                x2 = x + vx * self._dt
                y2 = y + vy * self._dt
        
                # Spline top surface
                f = interp1d(x2, y2, kind='cubic', fill_value='extrapolate')
    
                with self.init_mesh.access(self.Bmesh):
                    self.Bmesh.data[self.bottom, 0] = mesh.data[self.bottom, -1]
                    self.Bmesh.data[self.top, 0] = f(x)      
            else:
                coords = mesh.data[self.top]
                x = coords[:,0]
                y = coords[:,1]
                z = coords[:,-1]
                vx = uw.function.evalf(self.v.sym[0], coords)
                vy = uw.function.evalf(self.v.sym[1], coords)
                vz = uw.function.evalf(self.v.sym[2], coords)
                
                # Advect top surface
                x2 = x + vx * self._dt
                y2 = y + vy * self._dt
                z2 = z + vz * self._dt
        
                # Spline top surface
                f = CloughTocher2DInterpolator((x2, y2), z2)
    
                with self.init_mesh.access(self.Bmesh):
                    self.Bmesh.data[self.bottom, 0] = mesh.data[self.bottom, -1]
                    self.Bmesh.data[self.top, 0] = f((x,y)) 
        uw.mpi.barrier()
        #self.Bmesh.syncronise()
         
    def _update_mesh(self):
        with self.init_mesh.access():
            new_mesh_coords = self.init_mesh.data
            new_mesh_coords[:,-1] = self.Tmesh.data[:,0]   
        #mesh.deform_mesh(new_mesh_coords)
        return new_mesh_coords
        
    
    def solve(self):
        self._advect_surface()
        self._solve_sle()
        new_mesh_coords = self._update_mesh()
        return new_mesh_coords


# In[8]:


mesh = uw.meshing.StructuredQuadBox(elementRes=(int(xres), int(yres)), minCoords=(xmin, ymin), maxCoords=(xmax, ymax))
init_mesh = uw.meshing.StructuredQuadBox(elementRes=(int(xres), int(yres)), minCoords=(xmin, ymin), maxCoords=(xmax, ymax))
#meshbox = uw.meshing.UnstructuredSimplexBox(minCoords=(xmin, ymin), maxCoords=(xmax, ymax),cellSize=1.0/res,regular=False,qdegree=2,)
#mesh.view()

botwall = petsc_dm_find_labeled_points_local(mesh.dm,"Bottom")
topwall = petsc_dm_find_labeled_points_local(mesh.dm,"Top")


# In[9]:


hmax =  1.0  #nd(700* u.kilometer)
hLith = ndim(100.* u.kilometer)

#coord = fn.coord()
wavelength = ndim(2800.* u.kilometer)
amplitude  = ndim(7*u.kilometer)
k = 2. * np.pi / wavelength
#perturbationFn = amplitude*fn.math.cos(k*coord[0] )

def perturbation(x):
    return amplitude*np.cos(k*x)+hmax


Tmesh = uw.discretisation.MeshVariable("Tmesh", init_mesh, 1, degree=1)
Bmesh = uw.discretisation.MeshVariable("Bmesh", init_mesh, 1, degree=1)

mesh_solver = uw.systems.Poisson(init_mesh, u_Field=Tmesh, solver_name="FreeSurf_solver")
mesh_solver.constitutive_model = uw.constitutive_models.DiffusionModel
mesh_solver.constitutive_model.Parameters.diffusivity = 1. 
mesh_solver.f = 0.0
mesh_solver.add_dirichlet_bc(Bmesh.sym[0], "Top",0)
mesh_solver.add_dirichlet_bc(Bmesh.sym[0], "Bottom",0)

x = init_mesh.data[topwall,0]
with init_mesh.access(Bmesh):
    Bmesh.data[topwall, 0] = perturbation(x)
    Bmesh.data[botwall, 0] = init_mesh.data[botwall,-1]
mesh_solver.solve()

def update_mesh():
    with init_mesh.access():
        new_mesh_coords = init_mesh.data
        new_mesh_coords[:,-1] = Tmesh.data[:,0]
    return new_mesh_coords
new_mesh_coords = update_mesh()
mesh.deform_mesh(new_mesh_coords)
#update_mesh(mesh)


# # dq2dq1
# v = uw.discretisation.MeshVariable("V", mesh, mesh.dim, degree=2)
# p = uw.discretisation.MeshVariable("P", mesh, 1, degree=1)

# q1dq0
v = uw.discretisation.MeshVariable("V", mesh, mesh.dim, degree=1,continuous=True)
p = uw.discretisation.MeshVariable("P", mesh, 1, degree=0,continuous=False)

timeField = uw.discretisation.MeshVariable("time", mesh, 1, degree=1)


# In[10]:


plot_mesh("mesh01",mesh)


# In[11]:


swarm  = uw.swarm.Swarm(mesh)
material  = uw.swarm.IndexSwarmVariable("M", swarm, indices=2, proxy_degree=1)  
fill_parameter= 3 # swarm fill parameter
swarm.populate_petsc(fill_param=fill_parameter,layout=uw.swarm.SwarmPICLayout.GAUSS)
#pop_control = uw.swarm.PopulationControl(swarm)
#pop_control_dm = PopulationControl_DeformMesh(swarm,swarm)
pop_control = PopulationControl_DeformMesh(swarm)

MantleIndex = 0
LithIndex   = 1

with swarm.access(material):
    #perturbation = offset + amplitude * np.cos(k * swarm.particle_coordinates.data[:, 0])
    material.data[:, 0] = np.where(swarm.particle_coordinates.data[:, 1] < hmax-hLith, MantleIndex, LithIndex)
    
viscM = ndim(1e21 * u.pascal * u.second)
viscL = ndim(1e23 * u.pascal * u.second)    
    
densityM = ndim(3300 * u.kilogram / u.metre**3)   # for an
densityL = ndim(3300 * u.kilogram / u.metre**3)   # for litho    
    
density_fn = material.createMask([densityM, densityL])
visc_fn = material.createMask([viscM,viscL])


# In[12]:


plot_meshswarm("swarm_01",mesh,swarm,material,showSwarm=True,showFig=True)


# In[13]:


def build_stokes_solver(mesh,v,p):
    stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p)
    stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel
    stokes.bodyforce = sympy.Matrix([0, -1 * ND_gravity * density_fn])
    stokes.constitutive_model.Parameters.shear_viscosity_0 = visc_fn
    stokes.saddle_preconditioner = 1.0 / stokes.constitutive_model.Parameters.shear_viscosity_0
    stokes.add_dirichlet_bc((0.0,0.0), "Left", (0,))
    stokes.add_dirichlet_bc((0.0,0.0), "Right", (0,))
    stokes.add_dirichlet_bc((0.0,0.0), "Bottom", (0,1))
        
    if uw.mpi.size == 1:
        stokes.petsc_options['pc_type'] = 'lu'
    
    stokes.tolerance = 1.0e-6
    stokes.petsc_options["ksp_rtol"] = 1.0e-6
    stokes.petsc_options["ksp_atol"] = 1.0e-6
    stokes.petsc_options["snes_converged_reason"] = None
    stokes.petsc_options["snes_monitor_short"] = None
    return stokes


# ### Test swarm advection

# In[14]:


stokes = build_stokes_solver(mesh,v,p)
stokes.solve(zero_init_guess=False)

dt_solver = stokes.estimate_dt()
dt = min(dt_solver,dt_set)


# In[15]:


with swarm.access():
     print (swarm.particle_coordinates.data.shape)


# In[16]:


swarm.advection(V_fn=v.fn, delta_t=dt,evalf=True)
#swarm.advection(V_fn=v.fn, delta_t=dt)


# In[ ]:


swarm.advection(V_fn=v.fn, delta_t=dt)


# In[ ]:


# step      = 0
# max_steps = 2
# time      = 0
# dt        = 0

# w = []
# dwdt = []
# times = []

# while time < max_time:
# #while step < max_steps:
    
#     if uw.mpi.rank == 0:
#         string = """Step: {0:5d} Model Time: {1:6.1f} dt: {2:6.1f} ({3})\n""".format(
#         step, _adjust_time_units(time),
#         _adjust_time_units(dt),
#         datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
#         sys.stdout.write(string)
#         sys.stdout.flush()
    
#     stokes = build_stokes_solver(mesh,v,p)
#     stokes.solve(zero_init_guess=False)


#     if step%save_every ==0:
#         if uw.mpi.rank == 0:
#             print(f'\nSave data:')
#         with mesh.access(timeField):
#             timeField.data[:,0] = dim(time, u.megayear).m

#         mesh.petsc_save_checkpoint(meshVars=[v, p, timeField], index=step, outputPath=outputPath)
#         swarm.petsc_save_checkpoint(swarmName='swarm', index=step, outputPath=outputPath) 

#         if uw.mpi.rank == 0:
#             print(f'\nSave data done.')

#     times.append(time)
#     dt_solver = stokes.estimate_dt()
#     dt = min(dt_solver,dt_set)
    
#     swarm.advection(V_fn=v.fn, delta_t=dt,evalf=True)
#     #swarm.advection(V_fn=v.fn, delta_t=dt)

#     freesuface = FreeSurfaceProcessor(v,dt)
#     new_mesh_coords=freesuface.solve()
#     mesh.deform_mesh(new_mesh_coords)
#     if uw.mpi.rank == 0:
#         print(f'\nrepopulate start:')
#     pop_control.repopulate(mesh,material)
#     if uw.mpi.rank == 0:
#         print(f'\nrepopulate end:')

#     step += 1
#     time += dt


# In[ ]:





# In[ ]:





# In[ ]:





# In[ ]:





# In[ ]: