# %%
# Die Swell of a Newtonian Fluid
# ==============================
# This example demonstrates the *die swell* of a Newtonian viscous fluid using
# MigFlow’s Arbitrary Lagrangian–Eulerian (ALE) formulation for free surfaces.
#
# A 2D fluid jet exits a die and expands freely due to the relaxation of normal
# stresses and surface tension effects. The example also shows how to use
# MigFlow’s `FreeSurface` class to track and evolve the free boundary.
#
# Keywords
# --------
# FEM, ALE, Free Surface, Surface Tension
#
# Description
# -----------
# This test demonstrates:
# - How to define a 2D fluid domain in Gmsh with symmetry and open boundaries.
# - How to apply a parabolic inflow velocity profile and zero-pressure outlet.
# - How to evolve the free surface using the ALE method.
# - How to include capillary forces at the free surface via a weak formulation.
from migflow import scontact, fluid, time_integration, free_surface, gmsh_io
import numpy as np
import os, sys, shutil, subprocess

# %%
# Output Directory and Mesh Generation
# ------------------------------------
# The output directory is created, and a 2D annular mesh is generated using Gmsh.
outputdir = "output" if len(sys.argv) < 2 else sys.argv[1]
geomesh_filename = "2d_mesh.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])

# %%
# Geometrical Parameters
# ----------------------
# The computational domain is a rectangular 2D section representing a
# Newtonian jet leaving a die.
h = 1.0  # domain height [m]
l = 20.5  # domain length [m]
Ox, Oy = 0, 0  # bottom-left corner of the domain


# %%
# Physical Parameters
# -------------------
# Set the material and flow parameters for a Newtonian fluid.
g = np.array([0.0, 0.0])  # gravity [m/s²]
rho = 5.9  # density [kg/m³]
mu = 1.0  # dynamic viscosity [Pa·s]
nu = mu / rho  # kinematic viscosity [m²/s]
gamma = 1.0  # surface tension coefficient [N/m]


# %%
# Surface Tension Function
# ------------------------
# This helper function computes the curvature and applies a surface tension
# force along the free surface, using a weak formulation.
def apply_surface_tension_weak(x):
    fs_nodes = fs.fs_nodes
    lapl_h = fs.compute_lapl_fs(fs_nodes)
    n_e = len(lapl_h) - 1
    kappa = np.zeros_like(x[:, 0])
    X = f.coordinates()[fs_nodes, 0]
    h = f.coordinates()[fs_nodes, 1]
    dX = X[1:] - X[:-1]
    dhdx = (h[1:] - h[:-1]) / dX[:]
    dl = np.sqrt(dhdx[:] * dhdx[:] + 1) ** (-3)
    ind = np.argsort(x[:, 0])
    for i in range(0, n_e):
        ind_i = ind[2 * i : 2 * i + 2]
        xi = (2 * x[ind_i, 0] - (X[i + 1] + X[i])) / (X[i + 1] - X[i])
        phi = np.column_stack([(1 - xi) / 2, (1 + xi) / 2])
        kappa[ind_i] = (phi[0, :] * lapl_h[i] + phi[1, :] * lapl_h[i + 1]) * dl[i]
    tension = -gamma * kappa[:]
    return tension


# %%
# Fluid Problem Setup
# -------------------
# Initialize the fluid problem and define the boundary conditions.
f = fluid.FluidProblem2(g, mu, rho)
gmsh_io.load_msh(f, mesh_filename)
inlet_velocity = [lambda x: 2 * (1 - x[:, 1] ** 2), 0]
f.set_open_boundary("Left", velocity=inlet_velocity)
f.set_strong_boundary("Left", velocity=inlet_velocity)
f.set_wall_boundary("TopLeft", velocity=[0, 0])
f.set_symmetry_boundary("BottomLeft")
f.set_symmetry_boundary("BottomRight")
f.set_open_boundary("TopRight", pressure=apply_surface_tension_weak, viscous_flag=0)
f.set_open_boundary("Right", pressure=0)

# %%
# Free Surface Initialization
# ---------------------------
# Initialize the free surface tracking between the top and bottom right edges.
fs = free_surface.FreeSurface(f, "TopRight", "BottomRight", isCentered=False)
f.write_mig(outputdir, 0)

# %%
# Numerical Parameters
# --------------------
# Define time step and simulation duration.
outf = 20  # number of iterations between output files
dt = 0.025  # time step [s]
tEnd = 75.0  # final time [s]

# %%
# Time Integration Loop
# ---------------------
# The main simulation loop advances the ALE mesh and updates the free surface.
t = 0.0
i = 0
swelling = 0

while t < tEnd:
    print(f"{i:4d}, {t:.6g}/{tEnd:.6g}")
    if i % outf == 0:
        f.write_mig(outputdir, t)
    time_integration.iterate(f, None, dt, check_residual_norm=1)
    if t >= 25.0:
        fs.iterate(dt)
        swelling = f.coordinates()[fs.fs_nodes[-1], 1] / h

    t += dt
    i += 1

# %%
# Plot
# ----
# .. code-block:: shell
#
#  python3 -m migflow.plot.migplot output --actors fluid --fluid-field velocity --fluid-vmin 0.0 --fluid-vmax 2.0
# %%
# Artifacts
# ---------
# - output/animation.mp4