"""Grad-div on Raviart-Thomas and curl-curl on Nédélec, on an unstructured curved disk.

    H(div):   u - grad(div u) = f,     u.n = 0 on the circle   (RT_k)
    H(curl):  u + curl(curl u) = f,    u.t = 0 on the circle   (Nédélec_k)

Both are the natural coercive problems of their space: one field, a
symmetric positive definite system, solved by CG. What makes the disk the
test worth running, where the unit square is not:

* the cells are CURVED (an O-grid of degree-2 quadrilaterals read as an
  unstructured mesh), so the contravariant (RT) and covariant (Nédélec) Piola
  transforms are genuinely at work -- on a straight axis-aligned cell they are
  identities and a wrong one goes unnoticed;
* the DOFs are glued through the STORED face list of the unstructured mesh
  (``UnstructuredFaceDofMap``), with its orientation signs, not through the
  index arithmetic of a Cartesian grid.

The manufactured solution ``u = rho w`` with ``rho = 1 - r^2`` vanishes on the
circle, so BOTH its normal and its tangential trace do: one field serves the
two problems without a lifting (the spaces only impose a homogeneous datum,
since their DOFs are moments). It is not polynomial, so no space contains it
and the rates are real.

Expected: order ``k + 1`` in L2 for RT_k and Nédélec_k. The geometry is of
degree 2 (a parabola through three points of each boundary edge), whose own
error is O(h^3): it does not cap these orders.

Measured (CPU, O-grid n = 2, 4, 8, i.e. 20, 80, 320 cells; relative L2 error,
[observed order]):

    RT_0 grad-div         2.17e-01          1.02e-01 [1.09]    5.00e-02 [1.03]
    RT_1 grad-div         4.07e-02          1.09e-02 [1.90]    2.84e-03 [1.94]
    Nédélec_0 curl-curl   1.85e-01          7.83e-02 [1.24]    3.67e-02 [1.09]
    Nédélec_1 curl-curl   3.27e-02          8.59e-03 [1.93]    2.25e-03 [1.93]

(44 to 2 592 DOFs; 52 s for the whole table, each solve being mostly its own
compilation -- every refinement is a new mesh, hence new array shapes.) The
same elements on the Cartesian square are in the ``fem_convergence``
benchmark; what this adds is the curved, unstructured mesh.

Usage: ``python solve_hdiv_hcurl_unstructured_disk_2d.py [--no-plot]``.
"""

import functools
import sys
import time

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

from scimba_jax.linear_approximation.basis.analytic_bases import (
    local_nedelec_basis,
    local_nedelec_basis_by_logical,
    local_rt_basis,
    local_rt_basis_by_logical,
    nedelec_dof_layout,
    rt_dof_layout,
)
from scimba_jax.linear_approximation.basis.dof_map import UnstructuredFaceDofMap
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.unstructured_mesh import (
    UnstructuredMesh,
    unit_cell_quad_points,
)
from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized
from scimba_jax.linear_approximation.variables.variables_vector_element import (
    VariablesVectorElement,
)
from scimba_jax.mapping.macro_mesh import macro_mesh_ogrid_disk
from scimba_jax.physical_models.abstract_physical_weak_model import (
    AbstractPhysicalWeakModel,
)
from scimba_jax.physical_models.classical_weakform.hcurl_weak_form import HcurlWeakForm
from scimba_jax.physical_models.classical_weakform.hdiv_weak_form import HdivWeakForm

DIM = 2
GEOMETRY_ORDER = 2
REFINEMENTS = (2, 4, 8)  # O-grid parameter n: 5 n^2 cells
ORDERS = (0, 1)


# ── The manufactured solution: both traces vanish on the circle ──────────────


def rho(x):
    """``1 - r^2``: zero on the unit circle, smooth across the centre."""
    return 1.0 - x[0] ** 2 - x[1] ** 2


def u_exact(x):
    """``rho w``: the whole field vanishes on the circle, so both traces do."""
    return rho(x) * jnp.array([jnp.cos(x[1]), jnp.sin(x[0])])


def div_u(x):
    return jnp.trace(jax.jacfwd(u_exact)(x))


def curl_u(x):
    jacobian = jax.jacfwd(u_exact)(x)
    return jacobian[1, 0] - jacobian[0, 1]


# Sources differentiated rather than written out: a hand-derived source is one
# more thing to get wrong, and the point of the run is the discretisation.
def f_hdiv(x):
    """``u - grad(div u)``."""
    return u_exact(x) - jax.grad(div_u)(x)


def f_hcurl(x):
    """``u + curl(curl u)``, with ``curl(s) = (ds/dy, -ds/dx)`` for a scalar."""
    g = jax.grad(curl_u)(x)
    return u_exact(x) + jnp.array([g[1], -g[0]])


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


# ── Mesh and spaces ──────────────────────────────────────────────────────────


def disk(n, quad_order):
    """The O-grid disk, ``5 n^2`` curved cells, as an unstructured mesh."""
    macro = macro_mesh_ogrid_disk(radius=1.0, n=n, order=GEOMETRY_ORDER)
    return UnstructuredMesh.from_macro_mesh(
        macro, UnitSquareTensorized(dim=DIM, order=quad_order)
    )


# ⚠ One callable per (family, order), cached: a basis keeps its function in the
# pytree's aux_data, compared by IDENTITY, and a lambda rebuilt per mesh would
# make every refinement a new treedef -- a recompilation of everything.
@functools.lru_cache(maxsize=None)
def basis_callables(family, order):
    """``(local_basis, local_basis_by_logical, layout)`` of RT_k or Nédélec_k."""
    physical, logical, layout_of = {
        "rt": (local_rt_basis, local_rt_basis_by_logical, rt_dof_layout),
        "nedelec": (
            local_nedelec_basis,
            local_nedelec_basis_by_logical,
            nedelec_dof_layout,
        ),
    }[family]
    return (
        lambda y, i, m: physical(y, i, m, order=order, out_dim=DIM),
        lambda y, i, m: logical(y, i, m, order=order, out_dim=DIM),
        layout_of(order),
    )


def scheme(family, order, n):
    """Grad-div on RT_k or curl-curl on Nédélec_k, on the disk of parameter n."""
    mesh = disk(n, 2 * order + 4)
    local, by_logical, layout = basis_callables(family, order)
    basis = AnalyticBasis(
        nb_basis=layout.nb_local_dofs(DIM),
        out_dim=DIM,
        mesh=mesh,
        basis_type="field",
        local_basis=local,
        local_basis_by_logical=by_logical,
        dof_layout=layout,
    )
    form = (
        HdivWeakForm(dim=DIM, f=f_hdiv)
        if family == "rt"
        else HcurlWeakForm(dim=DIM, f=f_hcurl)
    )
    model = AbstractPhysicalWeakModel.from_weak_form(form, dirichlet=zero)
    return EllipticFEscheme(
        model,
        VariablesVectorElement(
            basis=basis, nb_variables=DIM, dof_map=UnstructuredFaceDofMap
        ),
    )


# ── Evaluation on the cells' own quadrature ──────────────────────────────────


def cell_values(solved):
    """``u_h`` and ``u`` at every quadrature point, cell by cell.

    On the mesh's own quadrature, cell known: no point location, and no point
    outside the disk. ``vmap`` over the cells, not a Python loop, which would
    compile once per cell.

    Returns:
        ``(weights, points, u_h, u)``, each flattened over cells and points.
    """
    variables = solved.variables
    mesh = variables.mesh
    x_hat = unit_cell_quad_points(mesh)

    def per_cell(cell):
        weights, points = mesh._local_weights_points(cell)

        def at(point, hat):
            return variables.local_expansion_pure(
                variables, variables.dofsl, cell, point, hat
            )

        return weights, points, jax.vmap(at)(points, x_hat), jax.vmap(u_exact)(points)

    weights, points, got, want = jax.vmap(per_cell)(jnp.arange(mesh.n_cells_total))
    return (
        np.asarray(weights).ravel(),
        np.asarray(points).reshape(-1, DIM),
        np.asarray(got).reshape(-1, DIM),
        np.asarray(want).reshape(-1, DIM),
    )


def l2_error(solved):
    """``||u_h - u|| / ||u||`` on the disk."""
    weights, _, got, want = cell_values(solved)
    error = np.sum(weights * np.sum((got - want) ** 2, axis=-1))
    return float(np.sqrt(error / np.sum(weights * np.sum(want**2, axis=-1))))


# ── The runs ─────────────────────────────────────────────────────────────────


LABELS = {"rt": "RT_{k} grad-div", "nedelec": "Nédélec_{k} curl-curl"}


def main():
    print(__doc__.split("\n\n")[0])
    print(
        f"\nO-grid disk, geometry of degree {GEOMETRY_ORDER}; relative L2 error "
        "[observed order], DOFs, time (first solve, compilation included)\n"
    )
    finest = {}
    for family in ("rt", "nedelec"):
        for order in ORDERS:
            row, previous = "", None
            for n in REFINEMENTS:
                started = time.perf_counter()
                solved = EllipticFEscheme.solve(scheme(family, order, n), max_iter=1)
                jax.block_until_ready(solved.variables.dofsl)
                seconds = time.perf_counter() - started
                error = l2_error(solved)
                # h ~ 1/n on the O-grid, so the order reads against n, never
                # as a log2 of consecutive columns.
                rate = (
                    ""
                    if previous is None
                    else f" [{np.log(previous[0] / error) / np.log(n / previous[1]):.2f}]"
                )
                previous = (error, n)
                n_dofs = int(solved.variables.dofsl.size)
                row += (
                    f"{f'{error:.2e}{rate}':>17}{f'{n_dofs}d':>8}{f'{seconds:.1f}s':>7}"
                )
            label = LABELS[family].format(k=order)
            print(f"{label:24}expect {order + 1}{row}", flush=True)
            finest[(family, order)] = solved
    if "--no-plot" not in sys.argv:
        plot(finest)


def plot(finest):
    """The highest order of each problem: ``|u_h|`` and its arrows, the error."""
    order = max(ORDERS)
    figure, axes = plt.subplots(2, 2, figsize=(10, 9))
    for row, family in enumerate(("rt", "nedelec")):
        _, points, got, want = cell_values(finest[(family, order)])
        triangles = mtri.Triangulation(points[:, 0], points[:, 1])
        magnitude = np.linalg.norm(got, axis=-1)
        error = np.linalg.norm(got - want, axis=-1)
        left, right = axes[row]
        shown = left.tricontourf(triangles, magnitude, levels=30, cmap="viridis")
        figure.colorbar(shown, ax=left, fraction=0.046)
        every = max(1, len(points) // 400)
        left.quiver(
            points[::every, 0],
            points[::every, 1],
            got[::every, 0],
            got[::every, 1],
            color="white",
            width=0.003,
        )
        left.set_title(f"{LABELS[family].format(k=order)}: |u_h| and u_h")
        shown = right.tricontourf(triangles, error, levels=30, cmap="magma")
        figure.colorbar(shown, ax=right, fraction=0.046, format="%.0e")
        right.set_title(
            f"|u_h - u|, {int(finest[(family, order)].variables.dofsl.size)} DOFs"
        )
        for axis in (left, right):
            axis.add_patch(plt.Circle((0, 0), 1.0, fill=False, color="0.3", lw=0.8))
            axis.set_aspect("equal")
            axis.set_xticks([])
            axis.set_yticks([])
    figure.suptitle(
        f"Unstructured curved disk, {REFINEMENTS[-1]}: {5 * REFINEMENTS[-1] ** 2} cells"
    )
    figure.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()