# %%
# Water–Salt Water Avalanche
# ==========================
# This test case simulates the collapse of a dense granular column immersed
# in a two-fluid environment composed of fresh water and sea water.
# The grain column is initially surrounded by denser sea water, while the rest
# of the domain is filled with lighter fresh water.
#
# The case cannot be run standalone: it requires the output of the
# *depot* testcase (``../depot/depot.py``) to initialize a compact granular
# column. The script can, however, generate a synthetic initial packing
# using a hexagonal arrangement for testing purposes.
#
# Keywords
# --------
# FEM, DEM, Unresolved, Friction, Two-Fluids

import sys, os, shutil, subprocess
from migflow import fluid, scontact, time_integration, gmsh_io, volume_coupling as vc
import numpy as np
import time

# %%
# Output Directory
# ----------------
# Create an output directory for results and generate a 2D mesh using Gmsh.
outputdir = "output_2d_avalanche_two_fluids" if len(sys.argv) < 2 else sys.argv[1]
geomesh_filename = f"mesh_two_fluids.geo"
mesh_filename = f"{outputdir}/mesh.msh"
shutil.rmtree(outputdir, ignore_errors=True)
os.makedirs(outputdir)
subprocess.call(["gmsh", "-2", geomesh_filename, "-o", mesh_filename])

# %%
# Particle Packing Function
# -------------------------
# Define a helper function to create an initial hexagonal particle arrangement
# when no pre-deposited configuration is available. The packing density is
# controlled by the range of particle radii.
use_pre_deposit = False


def hexagonal_packing(p, x0, y0, lx, ly, rmin, rmax):
    """Set all the particles’ center positions using a hexagonal arrangement.

    Parameters
    ----------
    p : migflow.scontact.ParticleProblem
        Particle problem object.
    x0, y0 : float
        Coordinates of the bottom-left corner of the packing region.
    lx, ly : float
        Width and height of the region to fill with particles.
    rmin, rmax : float
        Minimum and maximum particle radii.
    """
    x = np.arange(x0 + rmax, x0 + lx - rmax, 2 * rmax)
    y = np.arange(y0 + rmax, y0 + ly - rmax, 2 * (3**0.5) * rmax)
    x, y = np.meshgrid(x, y)
    x1 = np.arange(2 * rmax + x0, x0 + lx - rmax, 2 * rmax)
    y1 = np.arange(y0 + 3**0.5 * rmax + rmax, y0 + ly - rmax, 2 * (3**0.5) * rmax)
    x1, y1 = np.meshgrid(x1, y1)
    x = np.concatenate([x.reshape(-1), x1.reshape(-1)])
    y = np.concatenate([y.reshape(-1), y1.reshape(-1)])
    order = np.argsort(y)
    x = x[order]
    y = y[order]
    for xi, yi in zip(x, y):
        z = np.random.random()
        r = rmin + (rmax - rmin) * z
        p.add_particle((xi, yi), r, r**2 * np.pi * rhop, material="particle")


# %%
# Physical Parameters
# -------------------
# Define the physical properties for the particles, fresh water, and sea water.
# The density contrast between the two fluids drives buoyancy and stratification.
g = np.array([0, -9.81])  # gravity vector [m/s²]
rhop = 1500  # particle density [kg/m³]
r = 2e-3  # particle radius [m]
# Fresh water (light phase)
rhof = 1000  # density [kg/m³]
nuf = 1e-6  # kinematic viscosity [m²/s]
# Sea water (dense phase)
rhom = 1050  # density [kg/m³]
num = 1e-6  # kinematic viscosity [m²/s]

# %%
# Particle Problem
# ----------------
# Create the particle problem, load boundaries from the mesh, and initialize
# the grain assembly either from a pre-deposited configuration or by generating
# a regular packing.
p = scontact.ParticleProblem(2)
p.set_fixed_contact_geometry(0)  # workaround for contact geometry flag
gmsh_io.load_msh_boundaries(p, mesh_filename, ["Top", "Bottom", "Left", "Right"], material="Wall")
if use_pre_deposit:
    p1 = scontact.ParticleProblem(2)
    p1.read_mig("../depot/output", iteration=-1)
    first = p.n_bodies()
    for xi, ri in zip(p1.position(), p1.r()):
        p.add_particle(xi, ri, np.pi * ri**2 * rhop)
    p.body_velocity()[first:] = p1.velocity()
else:
    hexagonal_packing(p, 0, 0, 0.1, 0.199, 0.9 * r, 1.1 * r)
p.set_friction_coefficient(1.0, "Wall", "Particle")
p.set_friction_coefficient(0.5, "Particle", "Particle")

# %%
# Fluid Problem
# -------------
# Define the fluid problem with two phases (fresh water and sea water). The
# density and viscosity fields are specified as lists for the two components.
# The coupling is built in tiers: tier A is the body geometry on the mesh, tier
# B the legacy mixture closure on it.
f = fluid.FluidProblem2(g, nuf * rhof, rhof,
                        density_element="triangle_p1dg",
                        viscosity_element="triangle_p1dg")
gmsh_io.load_msh(f, mesh_filename)
# The properties follow the concentration, so they need the mesh to exist.
f.set_two_fluid_properties([rhof, rhom], [nuf * rhof, num * rhom])
f.set_wall_boundary("Bottom")
f.set_wall_boundary("Left")
f.set_wall_boundary("Right")
f.set_wall_boundary("Top")
f.set_mean_pressure(0)
# Tiers A and B only: this primes the coupling so the first frame carries a
# porosity field, and time_integration.iterate re-sets all three tiers -- datum
# included -- at every step.
geo = vc.get_particles(f, p.position(), p.r(), density=vc.body_density(p))
clo = vc.get_particles_closure(f, geo, nuf * rhof, rhof, p.velocity(),
                               omega=p.omega(),
                               contact_forces=p.contact_forces())
vc.set_coupling_geometry(f, geo)
vc.set_coupling_closure(f, clo)

# %%
# Initial Concentration Field
# ---------------------------
# The concentration field defines the spatial distribution of sea water
# (c = 1) and fresh water (c = 0). Initially, the left region is filled with
# sea water to represent the salt water column surrounding the grains.
x = f.coordinates()
c = np.zeros(f.n_nodes())
c[x[:, 0] < 0.1] = 1
f.set_concentration_cg(c)


# %%
# Time Integration
# ----------------
# The simulation runs for a short time interval with a fixed time step.
# At each iteration, the fluid phase is solved and the fluid-grain forces are applied to the particles,
# and then system evolves according to the contact dynamics solver.
# The simulation runs until the avalanche stabilizes.
t = 0  # initial time
i = 0  # iteration counter
tic = time.time()
outf = 10  # number of iterations between outputs
dt = 1e-3  # time step [s]
tEnd = 2  # final time [s]
mass = np.pi * p.r() ** 2 * rhop  # particle masses
while t < tEnd:
    print(f"{i = :4d} ----- {t = :g}")
    if i % outf == 0:
        p.write_mig(outputdir, t)
        f.write_mig(outputdir, t)
    # Transport the concentration, then let the two-fluid density and viscosity
    # follow it (in that order: the properties are derived FROM it). With grains
    # present advance_concentration runs its CONSERVATIVE form -- d(eps a)/dt +
    # div((u - umesh) a) = 0 -- so a compacting bed no longer carries the
    # concentration integral away with it.
    f.advance_concentration(dt)
    f.set_two_fluid_properties([rhof, rhom], [nuf * rhof, num * rhom])
    time_integration.iterate(f, p, dt, min_nsub=10, external_particles_forces=g * mass)
    t += dt
    i += 1

# %%
# Plot
# ----
# .. code-block:: shell
#
#  python3 -m migflow.plot.migplot output_2d_avalanche_two_fluids --actors fluid particles --fluid-field concentration
# %%
# Artifacts
# ---------
# - output_2d_avalanche_two_fluids/animation.mp4