# %%
# Bidimensional fall of two clouds made of particles in a viscous fluid
# ======================================================================
# This example illustrates the interaction of a cloud made of solid particles falling in a viscous
# fluid in two dimensions.
#
# Keywords
# --------
# FEM, DEM, Generation, Mesh Adaptation
#
# Description
# -----------
# The test demonstrates:
# - How to adapt the mesh according to the particles position.
# - How to define a gravity-driven sedimentation of dense particles.
# - How to couple the particle solver (`scontact`) and the fluid solver in MigFlow
# - How to export derived fields (pressure, velocity, porosity) for post-processing.
import os, sys, shutil
import numpy as np
import gmsh
from migflow import fluid, scontact, gmsh_io, time_integration
from migflow import volume_coupling as vc
from scipy.spatial import KDTree

gmsh.initialize()

# %%
DIAG = False
MULTIPLE_CLOUDS = True


# %%
# Output Directory
# ----------------
# The output directory is (re)created
import argparse
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("outputdir", nargs="?", default="output_2d_drop_duo")
# 6 s is enough for the cloud to be released, to fall and for the wake
# to develop; --tend 30 continues into the long-time spreading.
parser.add_argument("--tend", type=float, default=6.0,
                    help="final time [s] (default 6; long run: 30)")
args = parser.parse_args()
outputdir = args.outputdir
shutil.rmtree(outputdir, ignore_errors=True)
os.makedirs(outputdir)

# %%
# Physical Parameters
# -------------------
rhop = 2450  # particles density
g = np.array([0, -9.81])  # gravity
rho = 1030  # fluid density
nu = 1.17 / (2 * rho)  # kinematic viscosity
mu = rho * nu  # dynamic viscosity


# %%
# Geometrical parameters
# ----------------------
height = 2  # domain height
width = 0.1  # domain width
r = 100e-6  # particles radius (4x the original 25e-6: at 25 um a
            # grain is subpixel in the movie and the cloud reads as a smear)
compacity = 0.2  # solid volume fraction in the drop
rout = 3.3e-3  # drop radius
offset0 = np.array([0, 1.9 * height / 4])  # drop center position
offset1 = offset0 - (
    np.array([3 * rout, 3 * rout]) if DIAG else np.array([0, 3 * rout])
)  # second drop center position

# %%
# Particle Problem Initialization
# -------------------------------
# The particle positions are initialised randomly in a circular domain to obtain a given compacity.
p = scontact.ParticleProblem(2)


def genInitialPosition(p, r, rout, rhop, compacity, origin):
    """Set all the particles centre positions and create the particles objects

    Keyword arguments:
    p -- Particle Problem
    r -- max radius of the particles
    rout -- outer radius of the cloud
    rhop -- particles density
    compacity -- initial compacity in the cloud
    """
    # Space between the particles to obtain the expected compacity
    N = compacity * (rout / r) ** 2
    bodies = np.asarray([], int)
    n = 0
    while n < N:
        xyp = np.random.uniform(-rout, rout, 2)
        if np.hypot(xyp[0], xyp[1]) < rout:
            if p.n_particles() == 0:
                d = 1
            else:
                d = np.min(np.linalg.norm(p.position() - xyp[None, :], axis=1))
            if d > 2.1 * r:
                body = p.add_particle(xyp, r, r**2 * np.pi * rhop)
                bodies = np.append(bodies, body)
                n += 1
    # Shift of the particles to the top of the box
    p.body_position()[bodies, :] += origin
    return bodies


genInitialPosition(p, r, rout, rhop, compacity, offset0)
if MULTIPLE_CLOUDS:
    genInitialPosition(p, r, rout, rhop, compacity, offset1)

# %%
# Mesh generation
# ---------------
# The mesh is created with the python GMSH api.
# The refinement is given by the function set_size_call_back. In this example, a refinement close to the the particles position is chosen.
origin = [-width / 2, -height / 2]  # leftmost bottom corner


class MeshSize:
    def __call__(self, dim, tag, x, y, z, lc):
        dist, _ = self.tree.query([[x, y]])
        distmin = 0.01
        lcmin = 0.0015
        lcmax = 0.015
        alpha = np.clip((dist[0] - distmin) / 0.1, 0, 1)
        return lcmin * (1 - alpha) + lcmax * alpha

    def set_points(self, points):
        self.tree = KDTree(points)


def gen_mesh(width, height, mesh_size, origin=np.array([0, 0])):
    origin = np.asarray(origin)
    gmsh.model.add("box")
    gmsh.model.occ.add_rectangle(origin[0], origin[1], 0, width, height)
    gmsh.model.occ.synchronize()

    def get_line(x0, x1, eps=1e-6):
        r = gmsh.model.get_entities_in_bounding_box(
            x0[0] - eps, x0[1] - eps, -eps, x1[0] + eps, x1[1] + eps, eps, 1
        )
        return list(tag for dim, tag in r)

    h, w = height, width
    gmsh.model.add_physical_group(
        1, get_line(origin + [0, 0], origin + [w, 0]), name="Bottom"
    )
    gmsh.model.add_physical_group(
        1, get_line(origin + [0, h], origin + [w, h]), name="Top"
    )
    gmsh.model.add_physical_group(
        1,
        get_line(origin + [0, 0], origin + [0, h])
        + get_line(origin + [w, 0], origin + [w, h]),
        name="Lateral",
    )
    gmsh.model.add_physical_group(2, [1], name="domain")
    gmsh.model.mesh.set_size_callback(mesh_size)
    gmsh.model.mesh.generate(2)


mesh_size = MeshSize()
mesh_size.set_points(p.position()[p.r()[:, 0] > 0])
gen_mesh(width, height, mesh_size, origin)
# Load the mesh boundaries into the particle problem
gmsh_io.load_msh_boundaries(p, None, ["Bottom", "Top", "Lateral"])

# %%
# Fluid Problem Initialization
# ----------------------------
# The fluid solver is initialized using the generated mesh. All walls are no-slip.
f = fluid.FluidProblem2(g, mu, rho)
gmsh_io.load_msh(f, None)
f.set_wall_boundary("Bottom")
f.set_wall_boundary("Top")
f.set_wall_boundary("Lateral")
f.set_mean_pressure(0)

# %%
# Simulation Parameters
# ---------------------
outf = 10  # number of iterations between output files
remeshing = 10  # time step to adapt the mesh
dt = 1e-2  # time step
tEnd = args.tend  # final time [s], see --tend
t = 0  # initial time
i = 0  # initial iteration


# %%
# Write specific fields into the output
# -------------------------------------
def dynamic_pressure(f, rho, g):
    y = f.coordinates_fields()[f.field_indices(f.dimension())][:, 1]
    return f.pressure() - rho * g[1] * y


def get_fields(fluid):
    y = fluid.coordinates_fields()[fluid.field_indices(fluid.dimension())][:, 1]
    p1_element = fluid.get_p1_element()
    return {
        "pressure": (fluid.pressure(), p1_element),
        "velocity": (fluid.velocity(), p1_element),
        "porosity": (fluid.porosity().get(), p1_element),
        "u_solid": (fluid.u_solid(), p1_element),
        "dynamic_pressure": (fluid.pressure() - rho * g[1] * y, p1_element),
    }


# %%
# The coupling, before the first frame
# ------------------------------------
# THE COUPLED STEP IS ONE LIBRARY CALL below; this is only tier A + tier B so
# that the frame written at t = 0 already carries a porosity field. No tier C:
# nothing is being advanced here, and "no datum" is the absence of a
# set_coupling_datum call rather than a flag.
geo = vc.get_particles(f, p.position(), p.r(), density=vc.body_density(p))
clo = vc.get_particles_closure(f, geo, mu, rho, p.velocity(),
                               omega=p.omega(),
                               contact_forces=p.contact_forces())
vc.set_coupling_geometry(f, geo)
vc.set_coupling_closure(f, clo)

# %%
# Computational Loop
# ------------------
# The computational loop can start.
# The particles phase will use sub-iterations to keep a stable method. The minimal number of subiterations is given by min_nsub.
# The external forces given to the particles have to be given at each time step.
mass = np.pi * p.r() ** 2 * rhop
while t < tEnd:
    print(f"{i:4d}, {t:.6g}/{tEnd:.6g}")
    # Write outputs
    if i % outf == 0:
        p.write_mig(outputdir, t)
        f.write_mig(outputdir, t, get_fields(f))

    # The coupled step: the mixture operator (legacy retrieval), its Patankar
    # datum and the DEM advance, all inside the library call. The scheme sets
    # tiers A, B and C itself, so nothing about the coupling is written here.
    time_integration.iterate_patankar(
        f, p, dt, min_nsub=1, contact_tol=0.1 * r,
        external_particles_forces=mass * g,
        mu=mu, rho_f=rho, gravity=g)

    # Mesh adaptation. The overlap geometry is keyed on the fluid's CURRENT
    # mesh, so the next step rebuilds tier A against the new one by itself --
    # there is nothing to re-set here.
    if i % remeshing == 0 and i != 0:
        mesh_size.set_points(p.position()[p.r()[:, 0] > 0])
        gmsh.model.mesh.clear()
        gmsh.model.mesh.generate(2)
        gmsh_io.adapt_msh(f, None)  # project the fluid solution on the new mesh

    # Time stepping
    t += dt
    i += 1

# %%
# Plot
# ----
# .. code-block:: shell
#
#  python3 -m migflow.plot.migplot output_2d_drop_duo --actors fluid particles --fluid-field dynamic_pressure --follow particles --bounds -0.015 0.015 -0.015 0.015
# %%
# Artifacts
# ---------
# - output_2d_drop_duo/animation.mp4