"""Stokes on the unit square, and what the inf-sup condition actually does.

    -nu Delta u + grad p = f,   div u = 0,   u = 0 on the boundary

The first saddle-point system here: the pressure has no equation of its own,
it is the multiplier enforcing ``div u = 0``, and it enters the two equations
with opposite signs. The matrix is symmetric with eigenvalues of both signs,
so conjugate gradient does not apply -- hence ``EllipticSaddleScheme``, whose
only job is to keep the default solver from being the wrong one.

The two runs below are the point of the file. One pair satisfies the inf-sup
(LBB) condition and one does not, and what separates them is *not* accuracy:

* Q2/Q1 -- Taylor-Hood, the classical stable pair.
* Q1/Q1 -- unstable. The checkerboard pressure mode.

Read the pressure column, not the velocity one. In the unstable case the
velocity still converges perfectly well; only the pressure stagnates, two
orders of magnitude above the stable pair. A run that reported the velocity
alone would call both a success.

(Two more runs used to put the Q1 pressure on a COARSER mesh -- stable -- and
on a FINER one -- unstable, the same element pair. Spaces on two refinements
are no longer supported: one mesh per domain, one basis per variable, see
``check_same_refinement``.)
"""

import itertools

import jax
import jax.numpy as jnp
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 (
    LagrangeDofMap,
    UnstructuredNodalDofMap,
)
from scimba_jax.linear_approximation.basis.general_bases import AnalyticBasis
from scimba_jax.linear_approximation.galerkin.fem.elliptic_saddle_scheme import (
    EllipticSaddleScheme,
)
from scimba_jax.linear_approximation.meshes.mesh import Mesh
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.variables.variables_fe import VariablesFE
from scimba_jax.mapping.mapping import InvertibleFunction, Mapping
from scimba_jax.physical_models.abstract_physical_weak_model import (
    AbstractPhysicalWeakModel,
)
from scimba_jax.physical_models.classical_weakform.stokes_weak_form import (
    StokesWeakForm,
)
from scimba_jax.physical_models.weak_boundary_conditions import Dirichlet

DIM = 2
QUAD_ORDER = 6


def u_exact(x):
    """Divergence-free and vanishing on the whole boundary."""
    sx, cx = jnp.sin(jnp.pi * x[0]), jnp.cos(jnp.pi * x[0])
    sy, cy = jnp.sin(jnp.pi * x[1]), jnp.cos(jnp.pi * x[1])
    return jnp.pi * jnp.array([sx * sx * sy * cy, -sy * sy * sx * cx])


def p_exact(x):
    """Zero mean on the square, since the pressure is fixed only up to one."""
    return jnp.array([jnp.cos(jnp.pi * x[0]) * jnp.sin(jnp.pi * x[1])])


def source(x):
    """f = -Delta u + grad p, differentiated rather than written out."""
    laplacian = jnp.trace(jax.jacfwd(jax.jacfwd(u_exact))(x), axis1=1, axis2=2)
    return -laplacian + jax.jacfwd(lambda y: p_exact(y)[0])(x)


def make_mesh(n_cells):
    return Mesh(
        dim=DIM,
        n_cells=(n_cells, n_cells),
        ref_quad=UnitSquareTensorized(dim=DIM, order=QUAD_ORDER),
        mapping=Mapping(mappings=[InvertibleFunction(lambda x: x, lambda y: y)]),
    )


def unstructured_grid(n_cells, order):
    """The very same grid, written as a node list and a connectivity.

    Node ordering follows the tensor convention the Lagrange basis uses, and
    cell ordering ``unravel_index(idx, (n, n))``, so the two meshes agree cell
    for cell -- without which "the same answer" would not be testable.
    """
    side = n_cells * order + 1
    axis = np.linspace(0.0, 1.0, side)
    nodes = np.array([[axis[i], axis[j]] for i in range(side) for j in range(side)])
    cells = [
        [
            (r * order + a) * side + (c * order + b)
            for a, b in itertools.product(range(order + 1), repeat=2)
        ]
        for r in range(n_cells)
        for c in range(n_cells)
    ]
    return UnstructuredMesh(
        nodes=nodes,
        cells=np.array(cells),
        ref_quad=UnitSquareTensorized(dim=DIM, order=QUAD_ORDER),
        order=order,
    )


def make_space(mesh, order, out_dim, basis_type):
    unstructured = isinstance(mesh, UnstructuredMesh)
    basis = AnalyticBasis(
        nb_basis=(order + 1) ** DIM,
        out_dim=out_dim,
        mesh=mesh,
        local_basis=lambda y, i, m: local_lagrange_basis(
            y, i, m, order=order, out_dim=out_dim
        ),
        local_basis_by_logical=lambda y, i, m: local_lagrange_basis_by_logical(
            y, i, m, order=order, out_dim=out_dim
        ),
        basis_type=basis_type,
    )
    return VariablesFE(
        basis=basis,
        nb_variables=out_dim,
        # The nodal map BUILDS its numbering, so the element degree is free of
        # the mesh degree -- which is what lets Taylor-Hood put Q2 and Q1 on the
        # same mesh object. The isoparametric map, which reads the file's
        # connectivity, could not.
        dof_map=UnstructuredNodalDofMap if unstructured else LagrangeDofMap,
    )


def run(n_cells, order_velocity=2, order_pressure=1, mesh=None):
    """Solve once and return the two L2 errors."""
    mesh = make_mesh(n_cells) if mesh is None else mesh
    velocity = make_space(mesh, order_velocity, DIM, "field")
    pressure = make_space(mesh, order_pressure, 1, "scalar")

    model = AbstractPhysicalWeakModel(dim=DIM)
    model.add_weak_form("main", StokesWeakForm(dim=DIM, f=source))
    # On space 0 only: the velocity is prescribed, the pressure is not -- it is
    # a multiplier, and constraining it would over-determine the system.
    model.add_boundary_condition("0/boundary", Dirichlet(lambda x: jnp.zeros(DIM)))

    scheme = EllipticSaddleScheme(model, [velocity, pressure])
    scheme = EllipticSaddleScheme.solve(
        scheme, tol=1e-11, max_iter=1, max_iter_linear=20000
    )

    grid = (np.arange(60) + 0.5) / 60
    points = jnp.asarray(
        np.stack(np.meshgrid(grid, grid, indexing="ij"), axis=-1).reshape(-1, 2)
    )
    u_h = scheme.variables_list[0].evaluate(points)
    p_h = scheme.variables_list[1].evaluate(points)
    u_ref = jax.vmap(u_exact)(points)
    p_ref = jax.vmap(p_exact)(points)

    error_u = float(jnp.sqrt(jnp.mean(jnp.sum((u_h - u_ref) ** 2, axis=-1))))
    # The pressure is determined up to a constant (u is prescribed all around),
    # so both are centred before comparing.
    error_p = float(
        jnp.sqrt(jnp.mean(((p_h - p_h.mean()) - (p_ref - p_ref.mean())) ** 2))
    )
    return error_u, error_p


CASES = [
    ("Q2/Q1 (Taylor-Hood)", dict(), True),
    ("Q1/Q1", dict(order_velocity=1), False),
]

print("Stokes, saddle point.  Read the PRESSURE column.\n")
print(f"{'pair':38}{'n':>4}{'L2(u)':>12}{'L2(p)':>12}{'rate(p)':>9}{'inf-sup':>10}")
print("-" * 85)

# Three meshes rather than two: an unstable pair is not one whose pressure is
# merely large, it is one whose pressure does not IMPROVE. That needs a rate,
# and a rate needs a third point.
for label, options, stable in CASES:
    previous = None
    for n in (4, 8, 16):
        eu, ep = run(n, **options)
        rate = "" if previous is None else f"{np.log2(previous / ep):9.2f}"
        previous = ep
        print(
            f"{label:38}{n:>4}{eu:>12.3e}{ep:>12.3e}{rate:>9}"
            f"{'ok' if stable else 'FAILS':>10}"
        )
    print()

print(
    "The velocity converges in all four. Only the pressure tells the stable\n"
    "pairs from the unstable ones -- which is what the inf-sup condition is."
)


# ── The same, on an unstructured mesh ─────────────────────────────────────────
#
# A saddle point is where the unstructured port is most exposed: it is the only
# multi-space case here, so it exercises the cross-space coupling as well as the
# assembly. Run on a grid *described as* a node list and a connectivity, it must
# reproduce the Cartesian answer above -- which is known to be right
# independently.
#
# Taylor-Hood needs two orders on ONE mesh, which is what a DOF map that builds
# its numbering (rather than reading the file's connectivity) allows. Both
# spaces then share a mesh object, the scheme's `shared_mesh` stays true, and
# the cross-space terms are a connectivity gather -- no `find_cell_index`, so no
# Newton inverse anywhere on the path.
#
# That gain is structural, not numerical: the figures below do not improve, and
# would not. They are set by the saddle solve's own tolerance (1e-11), which is
# why Q1/Q1 -- which never went through find_cell_index either -- sits at 1e-13
# rather than 1e-15, the pressure being the multiplier and the worst
# conditioned. What the single mesh buys is that a distorted cell can no longer
# put a Newton solve on the critical path.

print("\n\nThe same, unstructured: a grid given as nodes + connectivity.\n")
print(f"{'pair':38}{'n':>4}{'|du|':>12}{'|dp|':>12}")
print("-" * 66)

# Q1/Q1 starts at 4: the pair is unstable AND 2x2 leaves it too few pressure
# DOFs, so the saddle system is singular there and both runs return NaN.
for label, order_v, meshes in (
    ("Q1/Q1, same mesh", 1, (4, 8)),
    ("Q2/Q1 Taylor-Hood, same mesh", 2, (2, 4)),
):
    for n in meshes:
        mesh_u = unstructured_grid(n, 1)
        cartesian = run(n, order_velocity=order_v)
        unstructured = run(n, order_velocity=order_v, mesh=mesh_u)
        print(
            f"{label:38}{n:>4}"
            f"{abs(cartesian[0] - unstructured[0]):>12.2e}"
            f"{abs(cartesian[1] - unstructured[1]):>12.2e}"
        )
    print()

print(
    "Both columns are differences against the Cartesian run, not errors: the\n"
    "point is that the two describe the same discrete problem."
)