# %%
# 2D Stationary Bubble test case
# ==============================
# This test case simulates a 2D droplet using the Finite Element Method (FEM) with a ghost method to include the correct surface tension effects.
# The droplet is initially stationary and should remain so throughout the simulation.
# %%
# Keywords
# --------
# FEM, fluid, 2D, ghost fluid method, surface tension
#
# %%
# Description
# -----------
# The test demonstrates:
# - How to include a ghost fluid method to handle surface tension effects at the interface between two fluids.
# - How to find the front nodes and compute curvature for surface tension.
# - How to set up a stationary droplet in a fluid domain.

import os, sys, shutil, subprocess
import gmsh
import numpy as np
from migflow import fluid, gmsh_io

gmsh.initialize()


def get_front_nodes(el, concentration):
    n_nodes = el.max() + 1
    touched_pos = np.full(n_nodes, False)
    touched_neg = np.full(n_nodes, False)
    touched_pos[el[concentration > 0]] = True
    touched_neg[el[concentration <= 0]] = True
    return touched_pos & touched_neg


# naive p_jump
def compute_p_jump(pos, el, concentration, front_nodes, sigma):
    curvature = np.zeros(len(front_nodes))
    for i in range(len(front_nodes)):
        n1 = front_nodes[i - 1]
        n2 = front_nodes[i]
        n3 = front_nodes[(i + 1) % len(front_nodes)]
        p1 = pos[n1]
        p2 = pos[n2]
        p3 = pos[n3]
        a = np.linalg.norm(p2 - p1)
        b = np.linalg.norm(p3 - p2)
        c = np.linalg.norm(p3 - p1)
        area = 0.5 * np.cross(p2 - p1, p3 - p1)[2]
        if a < 1e-12 or b < 1e-12 or c < 1e-12 or abs(area) < 1e-12:
            curvature[i] = 0.0
        else:
            curvature[i] = (4.0 * area) / (a * b * c)
    p_jump = np.zeros(el.shape)
    for i, node in enumerate(front_nodes):
        id = np.where(el == node)
        for j in range(len(id[0])):
            if concentration[id[0][j]] > 0:
                p_jump[id[0][j], id[1][j]] += sigma * curvature[i]
            else:
                p_jump[id[0][j], id[1][j]] += -sigma * curvature[i]
    return p_jump.reshape(-1)


# Output Directory
# ----------------
# Create a clean output directory for simulation results.
outputdir = "output_2d_stationnary_ghost" if len(sys.argv) < 2 else sys.argv[1]
shutil.rmtree(outputdir, ignore_errors=True)
os.makedirs(outputdir)
meshfile = f"{outputdir}/mesh.msh"
subprocess.call(["gmsh", "-2", "2d_mesh.geo", "-o", meshfile])

# Fluid Properties
# ----------------
g = np.array([0.0, 0.0])
rho_0 = 1000
rho_1 = 10
mu_0 = 1
mu_1 = 0.01
sigma = 10

# 2-Fluid Problem
# ---------------
f = fluid.FluidProblem2(g, mu_0, rho_0,
                        density_element="triangle_p1dg",
                        viscosity_element="triangle_p1dg")
gmsh_io.load_msh(f, meshfile)
f.set_wall_boundary("Boundary")
f.set_mean_pressure(0.0)

# Numerical parameters
# --------------------
t = 0
ii = 0
dt = 0.01
tEnd = 0.1
outf = 1

# Find font nodes
gmsh.open(meshfile)
el_tags = np.array(gmsh.model.mesh.get_elements(2)[1]).reshape(-1)
el_in = np.array(gmsh.model.mesh.get_elements(2, 200)[1]).reshape(-1)
el_in_set = set(el_in)
concentration = np.zeros(f.n_elements())
concentration = np.array([1 if tag in el_in_set else 0 for tag in el_tags])
front_nodes = get_front_nodes(f.elements(), concentration)
front_nodes = np.where(front_nodes)[0]
# Compute parametric coordinates of front nodes
node_coords = f.coordinates()[front_nodes]
thetas = np.arctan2(node_coords[:, 1] - 0.5, node_coords[:, 0] - 0.5)
sorted_indices = np.argsort(thetas)
front_nodes = front_nodes[sorted_indices]

while t < tEnd:
    print("Time t = ", t)
    # write fluid output files
    if ii % outf == 0:
        f.write_mig(outputdir, t)
    # move mesh
    x = f.coordinates().copy()
    x_old = x.copy()
    v_front = f.solution_at_coordinates(x[front_nodes])[:, :2]
    x[front_nodes, :2] += v_front * dt
    f.set_coordinates(x)
    vmesh = (x[:, :2] - x_old[:, :2]) / dt
    f.mesh_velocity().set_from_host(vmesh[:, :2].copy())
    f.concentration_dg().set_from_host(np.repeat(concentration.reshape((-1, 1)), 3, axis=1))
    # p_jump() is a flat P1DG array of n_elements * (dim + 1) values
    f.p_jump().set_from_host(compute_p_jump(
        f.coordinates(), f.elements(), concentration, front_nodes, sigma
    ))
    # compute fluid
    f.full_implicit_euler(dt, itermax=100)
    ii += 1
    t += dt

# %%
# Plot
# ----
# .. code-block:: shell
#
#  python3 -m migflow.plot.migplot output_2d_stationnary_ghost --actors fluid --fluid-field pressure --show-edges 1
# %%
# Artifacts
# ---------
# - output_2d_stationnary_ghost/animation.mp4