# %%
# Static bubble with an enriched pressure space
# =============================================
# The fitted circular interface carries a Leblond membrane. The pressure is
# continuous inside each phase and duplicated only at interface nodes.

import json
import os
import shutil
import subprocess
import sys

import numpy as np

from migflow import fluid, gmsh_io


output_dir = "output" if len(sys.argv) < 2 else sys.argv[1]
shutil.rmtree(output_dir, ignore_errors=True)
os.makedirs(output_dir)
mesh_file = os.path.join(output_dir, "mesh.msh")
subprocess.run(
    ["gmsh", "-2", "2d_mesh.geo", "-o", mesh_file],
    check=True,
    stdout=subprocess.DEVNULL,
)

radius = 0.25
sigma = 1.0
problem = fluid.FluidProblem2(
    np.zeros(2),
    1.0,
    1.0,
    sigma=sigma,
    solver=fluid.ILinearSystem.create_pardiso(iterative_refinement=5),
    temporal=False,
    advection=True,
    pressure_enrichment=True,
    surface_tension_model="membrane",
)
gmsh_io.load_msh(problem, mesh_file)

elements = problem.elements()
centroids = problem.coordinates()[elements, :2].mean(axis=1)
phases = (np.linalg.norm(centroids - np.array([0.5, 0.5]), axis=1) < radius).astype(np.int32)

# Mapping regression: exactly one extra pressure unknown per interface node.
touched = np.zeros((2, problem.n_nodes()), dtype=bool)
for phase in (0, 1):
    touched[phase, elements[phases == phase]] = True
interface_nodes = np.flatnonzero(touched[0] & touched[1])
problem.set_materials(phases)
assert problem.velocity_index().shape == (problem.n_nodes(), 2)
assert problem.pressure_index().size == problem.n_nodes() + interface_nodes.size
assert problem.pressure_node_volume().shape == (problem.pressure_index().size,)

pressure_map = problem.global_map()[:, 6:9]
for node in interface_nodes:
    occurrences = np.argwhere(elements == node)
    phase_dofs = []
    for phase in (0, 1):
        dofs = {
            pressure_map[iel, ilocal]
            for iel, ilocal in occurrences
            if phases[iel] == phase
        }
        assert len(dofs) == 1
        phase_dofs.append(dofs.pop())
    assert phase_dofs[0] != phase_dofs[1]

problem.set_strong_boundary("Boundary", velocity=[0.0, 0.0])
problem.set_mean_pressure(0.0)
problem.implicit_euler(1.0, check_residual_norm=1e-8)
problem.write_mig(output_dir, 0.0)

pressure_dg = problem.pressure_dg()
phase_pressure = [pressure_dg[phases == phase].mean() for phase in (0, 1)]
pressure_jump = phase_pressure[1] - phase_pressure[0]
laplace_jump = sigma / radius
relative_jump_error = abs(pressure_jump - laplace_jump) / laplace_jump
max_velocity = np.linalg.norm(problem.velocity(), axis=1).max()
phase_pressure_std = max(
    pressure_dg[phases == 0].std(), pressure_dg[phases == 1].std()
)

assert relative_jump_error < 0.02, (pressure_jump, laplace_jump)
assert phase_pressure_std < 1e-7, phase_pressure_std
assert max_velocity < 1e-7, max_velocity

# ALE smoke test: the fixed connectivity and enriched mapping survive coordinate
# updates, while the equilibrium interface remains immobile.
initial_coordinates = problem.coordinates().copy()
coordinates = initial_coordinates.copy()
dt = 1e-3
for step in range(10):
    old_coordinates = coordinates.copy()
    coordinates[interface_nodes, :2] += problem.velocity()[interface_nodes] * dt
    problem.set_coordinates(coordinates)
    problem.mesh_velocity().set_from_host(
        np.ascontiguousarray((coordinates[:, :2] - old_coordinates[:, :2]) / dt))
    problem.implicit_euler(dt)
    problem.write_mig(output_dir, (step + 1) * dt)

max_displacement = np.linalg.norm(
    coordinates[interface_nodes, :2] - initial_coordinates[interface_nodes, :2],
    axis=1,
).max()
assert max_displacement < 1e-8, max_displacement

with open(os.path.join(output_dir, "fluid.migf"), encoding="utf-8") as stream:
    metadata = json.load(stream)
assert metadata["fields"]["pressure"] == "P1DG"
assert metadata["pressure_enrichment"] is True
assert metadata["surface_tension_model"] == "membrane"

print(
    "static membrane bubble:",
    f"jump={pressure_jump:.12g}",
    f"relative_error={relative_jump_error:.3e}",
    f"max_velocity={max_velocity:.3e}",
    f"max_ALE_displacement={max_displacement:.3e}",
)