"""Stokes on spline spaces: which pairs satisfy inf-sup, and how far that goes.

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

The spline counterpart of ``solve_stokes_saddle_2d``. Write ``S^{p,r}`` for the
spline space of degree ``p`` and continuity ``r``; the four runs below are three
different things, and the measurement -- not the textbook -- says which.

* ``(S^{p,r})^2 / S^{p-1,r-1}`` -- the isogeometric Taylor-Hood analogue: one
  degree down AND one continuity down. Stable, at both degrees run.
* ``(S^{p,r})^2 / S^{p,r}`` -- equal order. **Degraded, not broken**: the
  pressure still converges, at roughly the same rate, with a constant about
  100x worse. Q1/Q1 fails harder than this, and the reason the spline pair
  survives is the continuity -- the checkerboard mode that sinks Q1/Q1 is not
  in ``S^{2,1}``. Written from the table below, not from the theory.
(A fourth case used to put the pressure on a FINER mesh than the velocity, and
failed by LOCKING -- too many pressure constraints, a velocity forced to nearly
zero, ``div u_h`` at 1e-10 and the velocity error at the size of ``u``. Spaces
on two refinements are no longer supported -- one mesh per domain, one basis
per variable, see ``check_same_refinement`` -- so that case is gone; the same
lesson survives in ``solve_stokes_saddle_2d.py`` as Q1/Q1.)

What the maximal continuity buys, and what it does not
------------------------------------------------------
A spline velocity space is ``C^{p-1}``, so the pair costs far fewer DOFs than
Taylor-Hood at the same accuracy -- the table prints both counts. It does not,
however, make the discrete velocity divergence-free: for the stable pairs the
``div u_h`` column converges to zero with the mesh instead of being zero.

Getting it EXACTLY zero is a different construction, and the pieces for it are
now in the repository: take the velocity in the spline H(div) space
(the ``fem_convergence`` benchmark) and the pressure in ``S^{p-1,r-1}``, the two
ends of the de Rham sequence, where ``div`` maps one ONTO the other. That pair
is pointwise solenoidal. It is not run here because an H(div) velocity has a
discontinuous tangential component, so the viscous term needs interior-penalty
faces -- a DG term inside a saddle-point FEM assembly, which does not exist yet.
Stated rather than approximated: this file measures Taylor-Hood, and the
``div u_h`` column is what the div-conforming pair would improve on.
"""

import jax
import jax.numpy as jnp
import numpy as np

from scimba_jax.linear_approximation.basis.analytic_bases import local_bspline_basis
from scimba_jax.linear_approximation.basis.dof_map import BSplineDofMap
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.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 = 7


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 make_space(mesh, order, regularity, out_dim, basis_type):
    basis = AnalyticBasis(
        nb_basis=(order + 1) ** DIM,
        out_dim=out_dim,
        mesh=mesh,
        local_basis=lambda y, i, m: local_bspline_basis(
            y, i, m, order=order, regularity=regularity, out_dim=out_dim
        ),
        basis_type=basis_type,
        spline_regularity=regularity,
    )
    return VariablesFE(basis=basis, nb_variables=out_dim, dof_map=BSplineDofMap)


def run(n_cells, velocity_space, pressure_space):
    """Solve once; return the two L2 errors, the DOF count and ||div u_h||."""
    mesh = make_mesh(n_cells)
    velocity = make_space(mesh, *velocity_space, DIM, "field")
    pressure = make_space(mesh, *pressure_space, 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))
    )

    # The divergence of the discrete velocity, sampled: not zero, and that is
    # the point of the closing note. Differentiated through `evaluate`, so it
    # is the divergence of the actual discrete field, not of an interpolant.
    def divergence(point):
        return jnp.trace(
            jax.jacfwd(
                lambda z: scheme.variables_list[0].evaluate(z[jnp.newaxis, :])[0]
            )(point)
        )

    div_norm = float(jnp.sqrt(jnp.mean(jax.vmap(divergence)(points) ** 2)))
    ndof = int(velocity.ndof_linear + pressure.ndof_linear)
    return error_u, error_p, ndof, div_norm


CASES = [
    ("S(2,1)^2 / S(1,0)  Taylor-Hood", (2, 1), (1, 0), "ok"),
    ("S(3,2)^2 / S(2,1)  Taylor-Hood", (3, 2), (2, 1), "ok"),
    ("S(2,1)^2 / S(2,1)  equal order", (2, 1), (2, 1), "poor"),
]


def main():
    print(__doc__.split("\n\n")[0])
    print("\nStokes on spline spaces.  Read the PRESSURE column.\n")
    print(
        f"{'pair':34}{'n':>4}{'ndof':>7}{'L2(u)':>11}{'L2(p)':>11}"
        f"{'rate(p)':>9}{'|div u|':>10}{'inf-sup':>9}"
    )
    print("-" * 95)

    # 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, velocity_space, pressure_space, verdict in CASES:
        previous = None
        for n in (4, 8, 16):
            eu, ep, ndof, div = run(n, velocity_space, pressure_space)
            rate = "" if previous is None else f"{np.log2(previous / ep):9.2f}"
            previous = ep
            print(
                f"{label:34}{n:>4}{ndof:>7}{eu:>11.3e}{ep:>11.3e}{rate:>9}"
                f"{div:>10.2e}{verdict:>9}"
            )
        print()

    print(
        "Three outcomes, not two.\n"
        "\n"
        "The two Taylor-Hood pairs converge in both fields. Equal order is\n"
        "DEGRADED rather than broken -- the pressure still improves at about\n"
        "the same rate, with a constant near 100x worse; maximal continuity\n"
        "removes the checkerboard mode that makes Q1/Q1 fail outright, so the\n"
        "spline pair is not the same kind of wrong as its nodal namesake.\n"
        "\n"
        "For the stable pairs div u goes to zero WITH the mesh rather than\n"
        "being zero. Velocity in the spline H(div) space with pressure in\n"
        "S(p-1,r-1) -- the two ends of the de Rham sequence, where div maps one\n"
        "ONTO the other -- would be pointwise solenoidal, and needs\n"
        "interior-penalty faces this assembly has not got."
    )


if __name__ == "__main__":
    main()