# %%
# 2D Bubbling Fluidized Bed
# =========================
# This example simulates a two-dimensional bubbling fluidized bed in which
# solid granular particles are suspended by an upward viscous fluid flow.
# Particles are injected through a narrow nozzle at the base of the domain
# and interact with the fluid through unresolved drag coupling.
# %%
# Keywords
# --------
# FEM, DEM, Fluidized Bed, Unresolved
import os, sys, shutil
import gmsh
import numpy as np
from migflow import advdiff, fluid, scontact, time_integration, gmsh_io, volume_coupling as vc

gmsh.initialize()

# %%
# Output Directory
# ----------------
# Create a clean output directory for simulation results.
outputdir = "output" if len(sys.argv) < 2 else sys.argv[1]
shutil.rmtree(outputdir, ignore_errors=True)
os.makedirs(outputdir)

# %%
# Geometrical parameters and mesh generation
# ------------------------------------------
# A rectangular 2D mesh is generated using Gmsh with named boundaries.
cm = 0.01
np.random.seed(0)


def gen_mesh(h, w, n, lc, o):
    o = np.asarray(o)
    gmsh.model.add("box")
    p0 = gmsh.model.occ.add_point(o[0], o[1], 0, lc)
    p1 = gmsh.model.occ.add_point(o[0] + (w - n) / 2, o[1], 0, lc)
    p2 = gmsh.model.occ.add_point(o[0] + (w + n) / 2, o[1], 0, lc)
    p3 = gmsh.model.occ.add_point(o[0] + w, o[1], 0, lc)
    p4 = gmsh.model.occ.add_point(o[0] + w, o[1] + h, 0, lc)
    p5 = gmsh.model.occ.add_point(o[0], o[1] + h, 0, lc)
    l0 = gmsh.model.occ.add_line(p0, p1)
    l1 = gmsh.model.occ.add_line(p1, p2)
    l2 = gmsh.model.occ.add_line(p2, p3)
    l3 = gmsh.model.occ.add_line(p3, p4)
    l4 = gmsh.model.occ.add_line(p4, p5)
    l5 = gmsh.model.occ.add_line(p5, p0)
    cloop = gmsh.model.occ.add_curve_loop([l0, l1, l2, l3, l4, l5])
    gmsh.model.occ.add_plane_surface([cloop])
    gmsh.model.occ.synchronize()
    gmsh.model.add_physical_group(1, [l0, l2], name="Bottom")
    gmsh.model.add_physical_group(1, [l4], name="Top")
    gmsh.model.add_physical_group(1, [l3, l5], name="Wall")
    gmsh.model.add_physical_group(1, [l1], name="Nozzle")
    gmsh.option.setNumber("Mesh.Algorithm", 2)
    gmsh.model.mesh.generate(2)


index_v = 2
index_m = 0
# ========= Input parameters =========
v_b = [1.2, 1.54, 1.71][index_v]  # bottom velocity               [m/s]
mass = [75e-3, 125e-3][index_m]  # mass of the particles         [kg]
friction = 0.3  # friction coefficient          [/]
print(f"{v_b = }, {mass = }, {friction = }")
# ========= Geometry =========
height = 25 * cm  # domain height                 [m]
width = 8 * cm  # domain width                  [m]
depth = 1.5 * cm  # domain depth                  [m]
nozzle = 1.3 * cm  # nozzle width                  [m]
origin = [-width / 2, 0]  # leftmost bottom corner        [m]
lc = 0.15 * cm  # mesh size                     [m]
gen_mesh(height, width, nozzle, lc, origin)
# ========= Fluid properties =========
g = np.array([0, -9.81])  # gravity                       [m/s**2]
rho = 1.204  # nitrogen density              [kg/m**3]
mu = 2.0e-5  # nitrogen dynamic viscosity    [kg/(ms)]
nu = mu / rho  # nitrogen kinematic viscosity  [m**2/s]
k = 2e-2  # nitrogen thermal conductivity [W/(m K)]
cp = 1010  # nitrogen heat capacity        [J/(Kg K)]
beta = 3e-3  # nitrogen thermal dilatation   [/]
u_b, v_b = 0, v_b  # bottom velocity               [m/s]
# ========= Grains properties =========
d = 0.1 * cm  # glass diameter                [m]
r = d / 2  # glass radii                   [m]
rhog = 2500  # glass density                 [kg/m**3]
cpg = 840  # glass heat capacity           [J/(kg K)]
kg = 1.4  # glass thermal conductivity    [W/(m K)]
young_modulus = 6e7  # glass young modulus           [kg/(m s**2)]
poisson_ratio = 0.22  # glass poisson ratio           [/]
T0 = -20  # offset temperature            [K]
Tf0 = T0 + 20  # initial fluid temperature     [K]
Tp0 = T0 + 90  # initial particle temperature  [K]
hw = 350  # wall convective heat transfer [W/(m**2 K)]
volume_ratio = 0.82

# ========= Particle Problem =========
p = scontact.ParticleProblem(2)
volume = mass / rhog
n_p3d = int(volume / (4 / 3 * np.pi * r**3))
n_p1d = int(depth / d)
n_p = int(n_p3d / n_p1d)
h = 1.0
read_input = False

eps = r
# Generate a hexagonal dense packing of particles
step = 2 * r + eps
lx, ly = width, height
x = np.arange(-lx / 2, lx / 2 - 2 * r, step)
y = np.arange(r + eps, h - r - eps, np.sqrt(3) * step)
x2 = np.arange(-lx / 2 + r, lx / 2 - 2 * r, step)
y2 = np.arange(r + eps, h - r - eps, np.sqrt(3) * step) + np.sqrt(3) * step / 2
x, y = np.meshgrid(x, y)
x2, y2 = np.meshgrid(x2, y2)
x = np.concatenate([x.ravel(), x2.ravel()])
y = np.concatenate([y.ravel(), y2.ravel()])
dx = x.max() - x.min()
offset_x = (width - dx) / 2
offset_y = y.min() - origin[1] - r
x += offset_x
y -= offset_y

# Sort particles by height for body ordering, for better visualization
order = np.argsort(y)
x, y = x[order], y[order]
x = x[:n_p]
y = y[:n_p]
for xi, yi in zip(x, y):
    p.add_particle((xi, yi), r, np.pi * r**2 * rhog, material="particle")

bl = np.array([origin[0], origin[1]])
br = np.array([origin[0] + width, origin[1]])
tl = np.array([origin[0], origin[1] + height])
tr = np.array([origin[0] + width, origin[1] + height])
p.add_boundary_segment(bl, br, "bottom", "wall")
p.add_boundary_segment(bl, tl, "left", "wall")
p.add_boundary_segment(tl, tr, "top", "wall")
p.add_boundary_segment(br, tr, "right", "wall")

p.set_friction_coefficient(1.0, "particle", "wall")
p.set_friction_coefficient(friction, "particle", "particle")

mcp = cpg / p.delassus()[:, 0, 0]
ks = np.full(p.n_particles(), kg)
e = np.full(p.n_particles(), young_modulus)
v = np.full(p.n_particles(), poisson_ratio)

# ========= Fluid Problem =========
# No drag_in_stab here: the closure ships a PER-ENTRY stabilisation weight and
# the kernel prefers it over the global flag. These grains are sub-cell, so the
# weight is 1 and the drag enters the SUPG/PSPG residual by itself -- which is
# what carries the bed. Verified 2026-08-21 with the flag explicitly FALSE:
# drag/weight peaks at 2.50 and the bed rises 4.67 -> 13.28 cm, against 0.60 and
# a collapse to 2.24 cm before the ramp existed.
f = fluid.FluidProblem2(g, mu, rho, volume_drag=12 * mu / (depth**2),
                        density_element="triangle_p1")
gmsh_io.load_msh(f, None)
f.set_open_boundary("Bottom", velocity=[u_b, v_b])
f.set_open_boundary("Top", pressure=0)
f.set_wall_boundary("Nozzle")
f.set_wall_boundary("Wall")
f.set_strong_boundary("Bottom", velocity=[u_b, v_b])
f.set_strong_boundary("Nozzle", velocity=[0, 0])

# %%
# Advection-Diffusion Problem Initialization
# ------------------------------------------
hp_mean = 2.0 * hw / depth
c = advdiff.AdvDiffProblem2(0, k, cp * rho, rho, mu, velocity_element="triangle_p1", volumetric_drag=hp_mean)
gmsh_io.load_msh(c, None)
c.set_open_boundary("Top")
c.set_strong_boundary("Bottom", solution=Tf0)
c.set_neumann_boundary("Nozzle", solution=Tf0, h_w=hw)
c.set_neumann_boundary("Wall", solution=Tf0, h_w=hw)
c.solution().fill(Tf0)


# %%
# Output Fields Function
# ----------------------
def get_fields(fluid, advdiff=None):
    """Return derived output fields for visualization."""
    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),
        "density": (fluid.density().get(), p1_element),
        "temperature": (advdiff.solution().get().reshape(-1, 1), p1_element),
    }


print(f"Start simulation with input parameters : {v_b=} ---- {mass=} ---- {friction=}")
# %%
# Simulation Loop
# ---------------
# Time integration of coupled fluid–particle motion.
# ========= Numerical parameters =========
t = 0
i = 0
dt_fluid = 1e-3
dt_heat = 1e-3
n_substeps = int(dt_heat / dt_fluid)
dt = dt_heat
tEnd = 1.0
outf = 10
# ========= Time integration =========
mass = np.pi * p.r() ** 2 * rhog
Tp = np.full_like((p.volume()), Tp0, dtype=np.float64)  # initial particle temperature
cpp = np.full_like(
    (p.volume()), cpg * rhog, dtype=np.float64
)  # heat capacity times density
kpp = np.full_like((p.volume()), kg, dtype=np.float64)  # thermal conductivity (unused)
nonzero = p.volume() > 0

while t < tEnd:
    print(f"{i:4d}, {t:.6g}/{tEnd:.6g}")
    c.set_particles(
        volume_ratio * p.volume(), p.position(), p.velocity(), cpp * p.volume(), kpp, Tp
    )
    # The 2d-3d ratio scales the body volume, so tiers A and B see the corrected
    # body: its radius and its density both come from that volume. The grains
    # are free and the implicit_euler below is the Patankar step that advances
    # them, so this generation carries a datum.
    vol = volume_ratio * p.volume()
    geo = vc.get_particles(f, p.position(), vc.body_radii(vol, 2),
                           density=vc.body_density(p, volume=vol))
    clo = vc.get_particles_closure(f, geo, mu, rho, p.velocity(),
                                   omega=p.omega(),
                                   contact_forces=p.contact_forces())
    dat = vc.get_patankar(geo, clo, dt_fluid, g, p.velocity())
    vc.set_coupling_geometry(f, geo)
    vc.set_coupling_closure(f, clo)
    vc.set_coupling_datum(f, dat)

    if i % outf == 0:
        f.write_mig(outputdir, t, get_fields(f, c))
        p.write_mig(outputdir, t, {"temperature": Tp.reshape(-1, 1)})

    c.velocity().set_from_host(f.velocity())
    c.implicit_euler(dt_heat)
    Tp[nonzero] += (
        c.get_flux_on_bodies()[nonzero] * dt / (cpp[nonzero] * p.volume()[nonzero])
    )
    dT = c.solution().get().reshape(-1, 1) - Tf0
    f.density().set_from_host(rho * (1 - np.fmin(beta * dT, 0.3)))

    # Solve fluid-particle coupling
    f.implicit_euler(dt_fluid)
    fext = g * mass + f.compute_node_force()
    time_integration._advance_particles(p, fext, dt_fluid, 1, 1e-3 * r)

    t += dt
    i += 1

# %%
# Plot
# ----
# .. code-block:: shell
#
#  python3 -m migflow.plot.migplot output --actors fluid particles --fluid-field density --element-type triangle_p1
# %%
# Artifacts
# ---------
# - output/animation.mp4