# %%
# Surface Tension on a Square Fluid Patch
# =======================================
# This example demonstrates the effect of **surface tension** on a square droplet
# of viscous fluid immersed in another immiscible fluid.
# The test shows how surface tension drives the interface to minimize curvature,
# progressively deforming the square patch into a circular droplet.
#
# Keywords
# --------
# Two-Phase Flow, Surface Tension, FEM
#
# Description
# -----------
# The test demonstrates:
# - How to initialize a two-phase flow with a square droplet region.
# - How to apply surface tension using MigFlow’s `sigma` parameter.
# - How to evolve the concentration field representing the interface.
# - How to export results for visualization and verification.
from migflow import fluid, time_integration, gmsh_io
import numpy as np
import os, sys, shutil, subprocess

# %%
# Output Directory
# ----------------
# Create a clean directory for simulation outputs.
outputdir = "output_2d_surface_tension" 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])

# %%
# Physical Parameters
# -------------------
# Define the properties of the two fluids and the surface tension coefficient.
g = np.array([0.0, 0.0])  # gravity [m/s²]
rhof = 1.0  # outer fluid density [kg/m³]
rhop = 10.0  # inner (droplet) density [kg/m³]
r = 0.01  # half-size of the square droplet [m]
compacity = 0.20  # solid volume fraction in the droplet
phi = 1.0 - compacity  # fluid volume fraction in droplet
nuf = 1.0  # kinematic viscosity of outer fluid [m²/s]
muf = nuf * rhof  # dynamic viscosity of outer fluid [Pa·s]
sigma = 0.1  # surface tension coefficient [N/m]

# Mixture properties (for the inner region)
rhom = (1 - phi) * rhop + phi * rhof  # mixture density [kg/m³]
num = 0.1  # mixture kinematic viscosity [m²/s]

print(f"Mixture properties: rhom = {rhom:g}, num = {num:g}")

# %%
# Numerical Parameters
# --------------------
tEnd = 2.0  # final simulation time [s]
dt = 5e-3  # time step [s]
outf = 1  # number of iterations between outputs
t = 0.0
i = 0

# %%
# Fluid Problem Initialization
# ----------------------------
# The domain and boundary conditions are loaded from a Gmsh mesh file.
# The concentration field is initialized such that `c=0` inside the square droplet.
f = fluid.FluidProblem2(g, nuf * rhof, rhof,
                        density_element="triangle_p1dg",
                        viscosity_element="triangle_p1dg")
gmsh_io.load_msh(f, mesh_filename)
# Define wall boundaries and mean pressure
for wall in ["Bottom", "Lateral", "Top"]:
    f.set_wall_boundary(wall)
f.set_mean_pressure(0)

# Initialize concentration field
x = f.coordinates()
c = np.ones(f.n_nodes())  # 1 = outer fluid
c[np.logical_and(np.abs(x[:, 0]) < r, np.abs(x[:, 1]) < r)] = 0.0  # inner droplet
f.set_concentration_cg(c)

# %%
# Simulation Loop
# ---------------
# Time integration is performed using the standard MigFlow time iterator.
# The concentration field evolves under advection and capillary forces.
while t < tEnd:
    print(f"{i:4d}, {t:.6g}/{tEnd:.6g}")
    # Write output files
    if i % outf == 0:
        f.write_mig(outputdir, t)
    # Advance one time step
    # the concentration is an ordinary field: transport it, then let the
    # material properties and the capillary force follow it
    f.advance_concentration(dt)
    f.set_two_fluid_properties([rhof, rhom], [nuf * rhof, num * rhom], sigma)
    time_integration.iterate(f, None, dt, check_residual_norm=1)
    i += 1
    t += dt

# %%
# Plot
# ----
# .. code-block:: shell
#
#  python3 -m migflow.plot.migplot output_2d_surface_tension --actors fluid --fluid-field concentration --fluid-vmin 0 --fluid-vmax 1
# %%
# Artifacts
# ---------
# - output_2d_surface_tension/animation.mp4