"""Flow past a cylinder, Schäfer-Turek benchmark 2D-1 (Re = 20, stationary).

Channel [0, 2.2] x [0, 0.41], cylinder of radius 0.05 centred at (0.2, 0.2),
parabolic inflow of maximum 0.3, no-slip on the walls and the cylinder,
do-nothing outflow (the natural condition of the weak form), nu = 1e-3, so
Re = U_mean D / nu = 20. Reference values (Schäfer & Turek 1996, and the
high-accuracy values of the FEATFLOW benchmark page): C_D = 5.57953523384,
C_L = 0.010618948146, Delta p = p(0.15, 0.2) - p(0.25, 0.2) = 0.11752016697.

Taylor-Hood Q2/Q1 on a curved gmsh mesh (cylinder as an exact circle), nested
1-to-4 refinements for the multigrid. Forces are computed by the RESIDUAL,
not by integrating the stress over the cylinder: the momentum residual of the
solution, assembled with the cylinder left free, summed over the cylinder's
nodes, IS the force the fluid exerts on it (the test function 1 on the
cylinder, 0 elsewhere) -- and it converges at the rate of the energy norm
squared, where a boundary integral of gradients loses an order.

Solver: Newton, FGMRES, block-triangular preconditioner, velocity by one
Schwarz V-cycle, Schur complement by LSC (LSCInverse). Measured, on
20 282 / 79 948 DOFs: 5 / 5 Newton steps, 107 / 115 FGMRES iterations in
total, 17.5 / 72 s; C_D 5.57929 / 5.57925 (4e-5), C_L 0.00941 / 0.01053
(11 % -> 0.8 %), Delta p 0.11837 / 0.11775 (0.7 % -> 0.2 %).

Why these choices (5 418 DOFs, Newton steps / FGMRES iterations):

* PCD, which serves the lid-driven cavity, is SINGULAR here: with an outflow
  the pressure has no free constant, while its pure-Neumann F_p still
  kills the constants -- an eigenvalue at -4e-10 in S_hat^-1 S, and no
  convergence (20 / 11 401). LSC reads its boundary behaviour off B:
  spectrum in [0.098, 2.21];
* the velocity multigrid holds (with the exact Schur complement: 6 / 19);
* the linear steps are solved (forcing="none"), not Eisenstat-Walker's
  0.9 ||F|| at the first step: that one-iteration first step leaves an
  iterate where the LSC-preconditioned system stalls (116, then 600
  iterations per step, no convergence); solved, 5 / 186;
* the multigrid must know which sides are CONSTRAINED: before
  Level.constrained read the Dirichlet sides instead of the whole
  boundary, the 15 outflow nodes were zeroed in the cycle, and one Stokes
  solve took 1 005 FGMRES iterations instead of 64.

Usage: python solve_cylinder_schaefer_turek_2d.py [mesh_size] [n_levels] [--no-plot]
(default 0.06, 2). Needs the optional mesh extra (gmsh).
"""

import sys
import time

import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np

from scimba_jax.linear_approximation.basis.analytic_bases import (
    local_lagrange_basis,
    local_lagrange_basis_by_logical,
)
from scimba_jax.linear_approximation.basis.dof_map import UnstructuredNodalDofMap
from scimba_jax.linear_approximation.basis.general_bases import AnalyticBasis
from scimba_jax.linear_approximation.galerkin.fem.elliptic_fe_scheme import (
    EllipticFEscheme,
)
from scimba_jax.linear_approximation.galerkin.fem.elliptic_saddle_scheme import (
    EllipticSaddleScheme,
)
from scimba_jax.linear_approximation.meshes.unstructured_mesh import UnstructuredMesh
from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized
from scimba_jax.linear_approximation.solvers import NewtonKrylov
from scimba_jax.linear_approximation.solvers.multigrid import MG
from scimba_jax.linear_approximation.solvers.schur import (
    BlockSchurPreconditioner,
    ChebyshevInverse,
    LSCInverse,
    MultigridInverse,
)
from scimba_jax.linear_approximation.solvers.smoothers import SchwarzSmoother
from scimba_jax.linear_approximation.transfer.hierarchy import build_hierarchy_nested
from scimba_jax.linear_approximation.variables.variables_fe import VariablesFE
from scimba_jax.mapping.macro_mesh import ExactCircle, macro_mesh_from_loops
from scimba_jax.mapping.mesh_hierarchy import mesh_hierarchy
from scimba_jax.physical_models.abstract_physical_weak_model import (
    AbstractPhysicalWeakModel,
)
from scimba_jax.physical_models.classical_weakform.mass_weak_form import MassWeakForm
from scimba_jax.physical_models.classical_weakform.navier_stokes_weak_form import (
    NavierStokesWeakForm,
)
from scimba_jax.physical_models.classical_weakform.stokes_weak_form import (
    VectorLaplacianWeakForm,
)
from scimba_jax.physical_models.weak_boundary_conditions import Dirichlet
from scimba_jax.plots.plots_galerkin import plot_flow_2d

L, H, CX, CY, R = 2.2, 0.41, 0.2, 0.2, 0.05
U_MAX, NU = 0.3, 1e-3
ARGS = [a for a in sys.argv[1:] if not a.startswith("-")]
MESH_SIZE = float(ARGS[0]) if len(ARGS) > 0 else 0.06
N_LEVELS = int(ARGS[1]) if len(ARGS) > 1 else 2
U, P = 0, 1
REF = dict(drag=5.57953523384, lift=0.010618948146, dp=0.11752016697)


def size(x, y):
    return MESH_SIZE * (0.35 + 0.65 * min(1.0, np.hypot(x - CX, y - CY) / 0.5))


macro = macro_mesh_from_loops(
    np.array([[0, 0], [L, 0], [L, H], [0, H]], float),
    holes=[ExactCircle(CX, CY, R)],
    order=2,
    mesh_size=size,
    smooth=False,
)
hierarchy = mesh_hierarchy(macro, n_levels=N_LEVELS)
QUAD = UnitSquareTensorized(dim=2, order=4)
cache = {}


def mesh_of(level_macro):
    if id(level_macro) not in cache:
        m = UnstructuredMesh.from_macro_mesh(level_macro, QUAD)
        m.label_boundary("inlet", lambda p: p[:, 0] < 1e-6)
        m.label_boundary("outlet", lambda p: p[:, 0] > L - 1e-6)
        m.label_boundary("walls", lambda p: (p[:, 1] < 1e-6) | (p[:, 1] > H - 1e-6))
        m.label_boundary(
            "cylinder", lambda p: np.hypot(p[:, 0] - CX, p[:, 1] - CY) < 1.5 * R
        )
        cache[id(level_macro)] = m
    return cache[id(level_macro)]


def space(mesh, order, out_dim, kind):
    kw = dict(order=order, out_dim=out_dim)
    b = AnalyticBasis(
        nb_basis=(order + 1) ** 2,
        out_dim=out_dim,
        mesh=mesh,
        basis_type=kind,
        local_basis=lambda y, i, m: local_lagrange_basis(y, i, m, **kw),
        local_basis_by_logical=lambda y, i, m: local_lagrange_basis_by_logical(
            y, i, m, **kw
        ),
    )
    return VariablesFE(basis=b, nb_variables=out_dim, dof_map=UnstructuredNodalDofMap)


def inflow(x):
    return jnp.array([4.0 * U_MAX * x[1] * (H - x[1]) / H**2, 0.0])


def zero(x):
    return jnp.zeros(2)


mesh = mesh_of(hierarchy.finest)
velocity, pressure = space(mesh, 2, 2, "field"), space(mesh, 1, 1, "scalar")


def ns_scheme(with_cylinder=True):
    model = AbstractPhysicalWeakModel(dim=2)
    model.add_weak_form("main", NavierStokesWeakForm(dim=2, viscosity=NU))
    model.add_boundary_condition("0/inlet", Dirichlet(inflow))
    model.add_boundary_condition("0/walls", Dirichlet(zero))
    if with_cylinder:
        model.add_boundary_condition("0/cylinder", Dirichlet(zero))
    return EllipticSaddleScheme(model, [velocity, pressure])


def velocity_scheme(level_macro):
    model = AbstractPhysicalWeakModel(dim=2)
    model.add_weak_form("main", VectorLaplacianWeakForm(dim=2))
    for side in ("inlet", "walls", "cylinder"):
        model.add_boundary_condition(side, Dirichlet(zero))
    return EllipticFEscheme(model, space(mesh_of(level_macro), 2, 2, "field"))


scheme = ns_scheme()
print(
    f"mesh: {mesh.n_cells_total} curved Q2 cells, {N_LEVELS} levels, "
    f"{scheme._initial_dofs().size} DOFs"
)
t0 = time.perf_counter()
mg_u = MG(build_hierarchy_nested(velocity_scheme, hierarchy), SchwarzSmoother())
velocity_inverse = MultigridInverse(mg_u, component=U)
mass_u = EllipticFEscheme(
    AbstractPhysicalWeakModel.from_weak_form(MassWeakForm(dim=2)), velocity
)
mass_u_diag = (
    1.0 / ChebyshevInverse(mass_u, degree=1, bounds=(1.0, 2.0)).inverse_diagonal
)
pressure_inverse = LSCInverse(mass_u_diag, rtol=1e-4, max_iter=500)
pc = BlockSchurPreconditioner(
    sweep=[P, U], inverses={U: velocity_inverse, P: pressure_inverse}
)
print(f"preconditioner setup: {time.perf_counter() - t0:.1f} s")
solver = NewtonKrylov(
    max_iter=20,
    tol=1e-9,
    cg_solver="fgmres:60",
    preconditioner=pc,
    max_iter_linear=600,
    rebuild_preconditioner=True,
    forcing="none",
)
t0 = time.perf_counter()
dofsl, report = EllipticSaddleScheme.solve_pure(scheme, solver=solver)
dofsl.block_until_ready()
print(
    f"Newton {int(report.n_iter)}, FGMRES {int(report.n_linear)}, "
    f"||F|| {float(report.residual):.1e}, {time.perf_counter() - t0:.1f} s (compilation included)"
)

# Forces by the residual: the momentum residual on the cylinder's nodes, cylinder left free.
free = ns_scheme(with_cylinder=False)
residual = type(free)._assembly_scheme_pure(free, dofsl)
res_u = free._split_dofs(residual)[U]  # (n_nodes, 2)
cyl_nodes = np.asarray(velocity.boundary_dofs("cylinder"))
force = -np.asarray(res_u)[cyl_nodes].sum(axis=0)
u_mean = 2.0 * U_MAX / 3.0
coef = 2.0 / (u_mean**2 * 2 * R)
velocity.dofsl, pressure.dofsl = free._split_dofs(dofsl)
p_front, p_back = np.asarray(
    pressure.evaluate(jnp.array([[CX - R, CY], [CX + R, CY]]))
)[:, 0]
got = dict(drag=coef * force[0], lift=coef * force[1], dp=p_front - p_back)
for k in ("drag", "lift", "dp"):
    gap = abs(got[k] - REF[k]) / abs(REF[k])
    print(f"  {k:5s} {got[k]: .6f}   reference {REF[k]: .6f}   relative gap {gap:.2e}")

if "--no-plot" not in sys.argv:
    fig, axes = plot_flow_2d(
        velocity,
        pressure,
        bounds=(0.0, 2.2, 0.0, 0.41),
        n_plot=440,
        inside=lambda x, y: (x - CX) ** 2 + (y - CY) ** 2 > R**2,
        title=(
            f"Schäfer-Turek 2D-1, Re = 20: C_D {got['drag']:.4f}, "
            f"C_L {got['lift']:.4f}, Δp {got['dp']:.4f}"
        ),
        vertical=True,
        show=False,
    )
    for ax in axes:
        ax.add_patch(plt.Circle((CX, CY), R, color="0.6"))
    plt.show()