# %%
# Particle tracking from baroclinic instability testcase
# ======================================================
# This example performs offline Lagrangian particle tracking using the 3D
# hydrodynamic output of the baroclinic instability testcase. Particles are
# continuously seeded along a vertical column near the center of the domain
# and are advected by the stored velocity field.
#
# The script reads the vertical coordinates and 3D velocity from the hydro
# output, advances the particles in time, exports active and deleted particles,
# and reseeds a new column of particles after each export interval.

import sys
import numpy as np
from slim4 import slim, slim3d, particles3d as p3d
from slim4.slim import si_format
from time import perf_counter as clock

# %%
# We first load the 3D hydrodynamic output used to drive the particle tracking.
# The free-surface coordinates and the velocity field are read from the eddies
# testcase output.

print("Loading 3D tracker")
mesh = slim3d.Mesh.from_output("run3d/", "mesh3dSlim.msh")
z = slim.DataOutput("run3d", "z", mesh)
velocity = slim.DataOutput("run3d", "u", mesh)


# %%
# The particle tracker uses prescribed horizontal and vertical diffusivities in
# addition to the resolved velocity field.

horizontal_diffusivity = slim.DataStatic.create_constant(mesh, 3.0)
vertical_diffusivity = slim.DataStatic.create_constant(mesh, 1e-5)
tracker = p3d.Tracker3d(mesh, z, velocity, horizontal_diffusivity, vertical_diffusivity)
tracker.set_boundary("Wall", "settle")

# %%
# We define the particle array. Custom attributes can be attached to particles;
# here we use "spawn_time", but any number (or none) can be used. A random seed
# may be set for reproducibility.

particles = p3d.ParticleArrays3d(mesh, ["spawn_time"], random_seed=1)


# %%
# We set the time step and the export frequency. The export interval is adjusted
# so that it is an integer multiple of the tracking time step.

dt = 300
export_dt = 3600

t0 = z.time[0]
tf = z.time[-1]
nsub = int(np.ceil(export_dt / dt))
dt = export_dt / nsub
nsteps = int(np.ceil(tf / dt))


# %%
# Particles are seeded along a vertical column in the middle of the domain. The
# number of particles can be provided on the command line; otherwise a default
# value is used.

nparticles = int(sys.argv[1]) if len(sys.argv) > 1 and sys.argv[1].isdigit() else 100
seed_position = np.zeros((3, nparticles))
seed_position[0, :] = 80e3
seed_position[1, :] = 250e3
seed_position[2, :] = np.linspace(0, -1000, nparticles)


def seed_particles(particle_arrays, time):
    """Insert a new column of particles and record their seeding time."""
    attributes = {"spawn_time": np.full(nparticles, time)}
    tracker.add_particles(particle_arrays, seed_position, time, attributes)


seed_particles(particles, t0)
print(f"number of particles: {particles.size:g}")


# %%
# We prepare the particle outputs. Active particles and deleted particles are
# written to separate XDMF files.

particles_output = "particles"
deleted_output = "deleted"
p3d.ParticleArrays3d.create_xdmf(particles_output)
p3d.ParticleArrays3d.create_xdmf(deleted_output)
particles.append_to_xdmf(particles_output, t0)


# %%
# We can now run the particle tracking. After each export interval, inactive
# particles are removed (and exported separately), active particles are written
# to disk, and a new column of particles is seeded. Particle sorting is important
# for performance, but it is also expensive and should not be done every time
# step. Here, remove_inactive() already performs sorting, so calling sort()
# separately would be redundant.

t = t0
iexport = 0

print("starting computation")
tic = clock()
for istep in range(nsteps):
    t = tracker.move(particles, t, dt)

    if (istep + 1) % nsub == 0:
        # particles.sort()  # useful for performance if not already handled
        deleted_particles = particles.remove_inactive()
        deleted_particles.append_to_xdmf(deleted_output, t, ["deletion_time"])
        particles.append_to_xdmf(particles_output, t)

        periter = (clock() - tic) / nsub
        tic = clock()
        iexport += 1
        print(f"Export {iexport}, Time {t}, periter {si_format(periter, 3)}s")

        seed_particles(particles, t)