"""Boundary conditions on a spline space: which ones keep the order, and which
one does not.

    -Delta u = f  on the unit square
    u = sin(pi(x + 1/4)) sin(pi(y + 1/4))

The solution vanishes on **no** side, and neither does its normal derivative.
That is the point: run against a solution that is zero on the boundary, a
Neumann or a Robin term could be wrong by any factor and the table would not
notice, because the data it multiplies would be zero.

Four configurations:

* **dirichlet** -- strong, on all four sides;
* **mixed** -- Dirichlet west/east, Neumann south/north;
* **robin** -- Robin on all four, so NO constrained DOF at all: the problem is
  held together by the boundary term alone;
* **nitsche** -- the SAME Dirichlet condition as the first line, imposed weakly
  instead of strongly.

The expected order is p+1 in L2 for a degree-p space. Neumann, Robin and Nitsche
meet it; strong Dirichlet does not. The first and last lines are the point of
the file: same condition, same data, same space -- only the way of imposing it
differs, and it costs two orders of convergence.
"""

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
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_fe_scheme import (
    EllipticFEscheme,
)
from scimba_jax.linear_approximation.meshes.mesh import Mesh, side_to_axis
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.laplacian_weak_form import (
    LaplacianWeakForm,
)
from scimba_jax.physical_models.weak_boundary_conditions import (
    Dirichlet,
    Neumann,
    Nitsche,
    Robin,
)

DIM = 2
ALPHA = 2.0
MESHES = (4, 8, 16, 32)
SIDES = ("west", "east", "south", "north")


def u_exact(x):
    return jnp.sin(jnp.pi * (x[0] + 0.25)) * jnp.sin(jnp.pi * (x[1] + 0.25))


def source(x):
    return 2 * jnp.pi**2 * u_exact(x)


def dirichlet_data(x):
    return jnp.atleast_1d(u_exact(x))


def normal_derivative(side):
    """``grad u . n`` on one side, as a one-argument function.

    A factory rather than ``lambda x, side=side: ...``. The default-argument
    idiom works now (the field classification counts required parameters only),
    but a closure says what is meant and cannot be misread.
    """
    axis, is_max = side_to_axis(side)
    sign = 1.0 if is_max else -1.0

    def g(x):
        return jnp.atleast_1d(sign * jax.grad(u_exact)(x)[axis])

    return g


def robin_data(side):
    """``grad u . n + alpha u``, the Robin datum on that side."""
    flux = normal_derivative(side)

    def g(x):
        return flux(x) + ALPHA * jnp.atleast_1d(u_exact(x))

    return g


def nitsche_penalty(order):
    """``sigma`` must beat the inverse-inequality constant, of order ``p^2``.

    Below that constant coercivity is not guaranteed, so the theory says nothing
    about the answer. Ten times ``p(p+1)`` clears it comfortably.

    ⚠ On THIS problem the requirement is not visible: :func:`penalty_sweep`
    measures a flat error from ``sigma = 0`` to ``10^4``. The default comes from
    theory, not from a measured threshold -- see that function.
    """
    return 10.0 * order * (order + 1)


def boundary_conditions(kind, order):
    if kind == "nitsche":
        sigma = nitsche_penalty(order)
        return {side: Nitsche(dirichlet_data, sigma=sigma) for side in SIDES}
    if kind == "dirichlet":
        return {side: Dirichlet(dirichlet_data) for side in SIDES}
    if kind == "mixed":
        return {
            "west": Dirichlet(dirichlet_data),
            "east": Dirichlet(dirichlet_data),
            "south": Neumann(normal_derivative("south")),
            "north": Neumann(normal_derivative("north")),
        }
    return {side: Robin(ALPHA, robin_data(side)) for side in SIDES}


def solve(n_cells, order, regularity, kind):
    mesh = Mesh(
        dim=DIM,
        n_cells=(n_cells,) * DIM,
        ref_quad=UnitSquareTensorized(dim=DIM, order=order + 2),
        mapping=Mapping(mappings=[InvertibleFunction(lambda x: x, lambda y: y)]),
    )
    basis = AnalyticBasis(
        nb_basis=(order + 1) ** DIM,
        out_dim=1,
        mesh=mesh,
        basis_type="scalar",
        local_basis=lambda y, i, m: local_bspline_basis(
            y, i, m, order=order, regularity=regularity, out_dim=1
        ),
        spline_regularity=regularity,
    )
    model = AbstractPhysicalWeakModel(dim=DIM)
    model.add_weak_form("main", LaplacianWeakForm(dim=DIM, f=source))
    for name, condition in boundary_conditions(kind, order).items():
        model.add_boundary_condition(name, condition)
    scheme = EllipticFEscheme(
        model, VariablesFE(basis=basis, nb_variables=1, dof_map=BSplineDofMap)
    )
    return EllipticFEscheme.solve(scheme, max_iter=1)


def l2_error(scheme, n_side=60):
    grid = (np.arange(n_side) + 0.5) / n_side
    points = jnp.asarray(
        np.stack(np.meshgrid(grid, grid, indexing="ij"), axis=-1).reshape(-1, DIM)
    )
    got = scheme.variables.evaluate(points).reshape(-1)
    want = jax.vmap(u_exact)(points)
    return float(jnp.linalg.norm(got - want) / jnp.linalg.norm(want))


def main():
    print(__doc__.split("\n\n")[0])
    print("\nRelative L2 error [observed order]. Maximal continuity (r = p-1).\n")

    results = {}
    for order in (2, 3):
        print(
            f"{f'p={order}, r={order - 1}':22}"
            + "".join(f"{f'n={n}':>18}" for n in MESHES)
        )
        print("-" * (22 + 18 * len(MESHES)))
        for kind in ("dirichlet", "mixed", "robin", "nitsche"):
            errors, row, previous = [], "", None
            for n_cells in MESHES:
                scheme = solve(n_cells, order, order - 1, kind)
                error = l2_error(scheme)
                errors.append(error)
                rate = "" if previous is None else f" [{np.log2(previous / error):.2f}]"
                previous = error
                row += f"{f'{error:.2e}{rate}':>18}"
            results[(order, kind)] = errors
            print(f"{kind:22}{row}")
        print()

    print(
        "Read the first and last lines together: SAME condition, same data, same\n"
        "space. Strong Dirichlet caps at order 2; Nitsche holds p+1, and at n=32\n"
        "with p=3 the gap is four orders of magnitude.\n"
        "\n"
        "The cause is not the condition but how a spline space imposes it. A\n"
        "spline coefficient is NOT a nodal value, so a strong datum is imposed at\n"
        "the Greville abscissae -- an INTERPOLATION, second order accurate, whose\n"
        "error then pollutes the whole domain. It is exact only for an affine\n"
        "datum, which is why the unit tests (affine) see nothing at all.\n"
        "\n"
        "Nitsche never evaluates g at a point: it only integrates it. That is the\n"
        "whole reason it escapes, and it is why Robin -- which constrains no DOF\n"
        "either -- was already at the right order.\n"
        "\n"
        "Strong Dirichlet remains the right default for a NODAL space on a\n"
        "conforming boundary: exact at the nodes, no penalty to choose. It is the\n"
        "wrong one exactly when the DOFs are not point values.\n"
        "\n"
        "The penalty sweep below is a caveat, not a result: on this problem the\n"
        "error does not depend on sigma at all, down to sigma = 0. What carries\n"
        "Nitsche here is its consistency and symmetry terms; the penalty is\n"
        "insurance this problem happens not to need."
    )

    penalty_sweep()
    _plot(results)


def penalty_sweep(order=3, n_cells=16):
    """How much room there is on ``sigma`` -- measured, and it contradicted me.

    I wrote that too small a penalty would be unstable rather than merely
    inaccurate, which is the textbook statement. The measurement says otherwise
    HERE: the error is flat from ``sigma = 0`` to ``10^4``, and ``sigma = 0`` is
    as good as the default.

    That does not make the theory wrong, it places it. Symmetric Nitsche without
    a penalty is not guaranteed coercive; it merely happens to give a
    non-singular system on this one -- a smooth spline space, few functions
    reaching the boundary, hence a small inverse-inequality constant. What
    carries the method here is the CONSISTENCY and SYMMETRY terms, not the
    penalty.

    Read the table as "the constant is not delicate on this problem", not as
    "the penalty is optional". A nodal space at high order, or a mesh with a
    sliver cell, would not be so forgiving.
    """
    print(
        f"\nPenalty sweep, p={order} r={order - 1}, n={n_cells}. "
        f"Default is sigma = {nitsche_penalty(order):.0f}.\n"
    )
    print(f"{'sigma':>10}{'relative L2':>16}")
    print("-" * 26)
    for sigma in (0.0, 0.001, 0.01, 0.1, 1.0, 120.0, 10000.0):
        model = AbstractPhysicalWeakModel(dim=DIM)
        model.add_weak_form("main", LaplacianWeakForm(dim=DIM, f=source))
        for side in SIDES:
            model.add_boundary_condition(side, Nitsche(dirichlet_data, sigma=sigma))
        mesh = Mesh(
            dim=DIM,
            n_cells=(n_cells,) * DIM,
            ref_quad=UnitSquareTensorized(dim=DIM, order=order + 2),
            mapping=Mapping(mappings=[InvertibleFunction(lambda x: x, lambda y: y)]),
        )
        basis = AnalyticBasis(
            nb_basis=(order + 1) ** DIM,
            out_dim=1,
            mesh=mesh,
            basis_type="scalar",
            local_basis=lambda y, i, m: local_bspline_basis(
                y, i, m, order=order, regularity=order - 1, out_dim=1
            ),
            spline_regularity=order - 1,
        )
        scheme = EllipticFEscheme(
            model, VariablesFE(basis=basis, nb_variables=1, dof_map=BSplineDofMap)
        )
        scheme = EllipticFEscheme.solve(scheme, max_iter=1)
        print(f"{sigma:>10.3g}{l2_error(scheme):>16.3e}")


def _plot(results):
    h = 1.0 / np.array(MESHES)
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.5), sharey=True)
    colours = {
        "dirichlet": "tab:red",
        "mixed": "tab:blue",
        "robin": "tab:green",
        "nitsche": "tab:purple",
    }

    for ax, order in zip(axes, (2, 3)):
        for kind, colour in colours.items():
            errors = results[(order, kind)]
            ax.loglog(h, errors, "o-", color=colour, label=kind)
        # Slopes to read against, anchored on the finest Robin point.
        anchor = results[(order, "robin")][-1]
        for slope, style in ((order + 1, "-"), (2, "--")):
            ax.loglog(
                h,
                anchor * (h / h[-1]) ** slope,
                style,
                color="0.5",
                linewidth=1,
                label=f"order {slope}",
            )
        ax.set_title(f"p = {order}, r = {order - 1}")
        ax.set_xlabel("h")
        ax.grid(True, which="both", alpha=0.3)
    axes[0].set_ylabel("relative L2 error")
    axes[0].legend(fontsize=8)
    fig.suptitle(
        "Same Dirichlet data, two impositions: strong (red) caps at order 2, "
        "Nitsche (purple) holds p+1"
    )
    fig.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()