# %%
# Three-dimensional fall of two clouds made of particles in a viscous fluid
# =========================================================================
# This examples illustrates how to define the fluid problem and the particles one.
# Setup is based on "Implementation of an unresolved stabilised FEM–DEM model to solve immersed granular flows" by M.Constant
# We consider the Set 3 case.
# All the relevant mesh information is directly loaded from the GMSH api.
# %%
# Keywords
# --------
# FEM, DEM, Sedimentation, Mesh Adaptation
import gmsh
import os, time, shutil, sys
import numpy as np
from scipy.spatial import KDTree
from migflow import fluid, scontact, gmsh_io, time_integration
from migflow import volume_coupling as vc

np.random.seed(0)
use_diagonal_offset = False

# %%
# Output Directory
# ----------------
# The output directory is (re)created
import argparse
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("outputdir", nargs="?", default="output_3d_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
# -------------------
rout = 2.7e-3  # drop radius
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.03  # solid volume fraction in the drop
g = np.array([0, -9.81, 0])  # gravity
rho = 1220  # fluid density
rhop = 2400  # particles density
nu = 1e-4  # kinematic viscosity
mu = rho * nu  # dynamic viscosity

# %%
# Geometrical parameters
# ----------------------
height = 1  # domain height
width = 10 * rout  # domain width
depth = 10 * rout  # domain width
origin = [-width / 2, -height / 2, -depth / 2]  # leftmost bottom corner


# %%
# 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.
gmsh.initialize()

def gen_geometry(width, height, origin=np.array([0, 0, 0])):
    origin = np.asarray(origin)
    gmsh.model.add("box")
    gmsh.model.occ.add_box(origin[0], origin[1], origin[2], width, height, depth)
    gmsh.model.occ.synchronize()

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

    h, w, d = height, width, depth
    lateral = (
        get_surface(origin + [0, 0, 0], origin + [0, h, d])
        + get_surface(origin + [0, 0, 0], origin + [w, h, 0])
        + get_surface(origin + [w, 0, 0], origin + [w, h, d])
        + get_surface(origin + [0, 0, d], origin + [w, h, d])
    )
    gmsh.model.add_physical_group(
        2, get_surface(origin + [0, h, 0], origin + [w, h, d]), name="Top"
    )
    gmsh.model.add_physical_group(
        2, get_surface(origin + [0, 0, 0], origin + [w, 0, d]), name="Bottom"
    )
    gmsh.model.add_physical_group(2, lateral, name="Lateral")
    gmsh.model.add_physical_group(3, [1], name="Domain")


def gen_mesh_callback(xp):
    tree = KDTree(xp)

    def size_f(x):
        dist, _ = tree.query(x[:, :3])
        distmin = 0.005
        mylc = 0.002
        alpha = np.clip((dist[0] - distmin) / 0.1, 0, 1)
        size = mylc * (1 - alpha) + 10 * mylc * alpha
        return size

    gmsh.model.mesh.set_size_callback(
        lambda dim, tag, x, y, z, lc: size_f(np.array([[x, y, z]]))
    )
    gmsh.model.mesh.clear()
    gmsh.model.mesh.generate(3)
    gmsh.model.mesh.remove_size_callback()


# %%
# The particle positions are initialised randomly in a circular domain to obtain a given compacity.
def genInitialPosition(p: scontact.ParticleProblem, 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
    """
    N_p = int(compacity * (rout / r) ** 3)
    bodies = np.asarray([], int)
    n = 0
    while n < N_p:
        print("...add particle : ", n, "/", N_p, end="\r")
        xyp = np.random.uniform(-rout, rout, 3)
        if np.linalg.norm(xyp) < 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, 4 / 3 * r**3 * np.pi * rhop)
                bodies = np.append(bodies, body)
                n += 1
    # Shift of the particles to the top of the box
    actual_compacity = np.sum(p.volume()) / ((4 / 3) * np.pi * rout**3)
    print(f"Actual compacity: {actual_compacity:.3f} ({p.n_particles()} particles)")
    p.body_position()[bodies, :] += origin
    return bodies


# %%
# Particle Problem
# ----------------
# The particle problem is created and its dimension is provided.
p = scontact.ParticleProblem(3)
offset0 = np.array([-2 * rout if use_diagonal_offset else 0, 1.9 * height / 4, 0])
offset1 = offset0 - np.array([-2 * rout if use_diagonal_offset else 0, 4 * rout, 0])
genInitialPosition(p, r, rout, rhop, compacity, offset0)
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.
gen_geometry(width, height, origin)
gen_mesh_callback(p.position()[p.r()[:, 0] > 0])
# The mesh can be loaded through the GMSH api if None is given or through a mesh filename.
gmsh_io.load_msh_boundaries(p, None, ["Bottom", "Top", "Lateral"])

# %%
# Fluid Problem
# -------------
# The fluid is described by its dimension, the external volume force applied, its dynamic viscosity and its density.
f = fluid.FluidProblem3(g, nu * rho, 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 = 20  # number of iterations between output files
remeshing = 20  # time step to adapt the mesh
dt = 5e-3  # time step
tEnd = args.tend  # final time [s], see --tend
t = 0  # initial time
ii = 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
# ------------------
tic = time.process_time()
mass = 4 / 3 * np.pi * p.r() ** 3 * rhop
while t < tEnd:
    print("%i : %.2g/%.2g (cpu %.6g)" % (ii, t, tEnd, time.process_time() - tic))
    # Mesh adaptation. The coupled step below rebuilds tier A against whatever
    # mesh the fluid now carries, so the SOLVE needs nothing here -- but
    # adapt_msh leaves the porosity field at 1, and a frame written before the
    # next solve would show an empty domain. Tiers A and B are therefore re-set
    # for the OUTPUT; no datum, because nothing is being advanced.
    if ii % remeshing == 0 and ii != 0:
        gen_mesh_callback(p.position()[p.r()[:, 0] > 0])
        gmsh_io.adapt_msh(f, None)
        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)
    if ii % outf == 0:
        f.write_mig(outputdir, t, get_fields(f))
        p.write_mig(outputdir, t)
    # 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=1e-3 * r,
        external_particles_forces=mass * g,
        mu=mu, rho_f=rho, gravity=g)
    t += dt
    ii += 1

# %%
# Plot
# ----
# .. code-block:: shell
#
#  python3 -m migflow.plot.migplot output_3d_drop_duo --actors particles --follow particles --grain-radius-scale 3 --bounds -0.012 0.012 0.4575 0.4815 -0.012 0.012
# %%
# Artifacts
# ---------
# - output_3d_drop_duo/animation.mp4