r"""The current hole: an internal kink in reduced resistive MHD, CG-FEM on the disk.

    d_t psi   = [psi, phi] + eta (J - J_c)
    d_t omega = [omega, phi] + [psi, J] + nu Delta omega
    J = Delta psi,   omega = Delta phi,      all four zero on the circle

on the unit disk, with the driven current of Deriaz, Despres, Faccanoni,
Gostaf, Imbert-Gerard, Sadaka and Sart (*Magnetic equations with FreeFem++:
the Grad-Shafranov equation and the current hole*, ESAIM Proc. 32, 2011,
sec. 3, after Huysmans and Czarny):

    J_c = j_1 (1 - r^4) - j_2 (1 - r^2)^8,   j_1 = 0.2,  j_2 = 0.2666,
    eta = 1e-5,  nu = 1e-6,   psi(0) = Delta^-1 J_c,  J(0) = J_c,  phi = omega = 0.

The current density is NEGATIVE near the axis -- the "current hole" -- and
the equilibrium is unstable to an internal kink: the flux surfaces stay
axisymmetric through a long linear stage, then the hole is expelled outward
and the current profile reconnects. No perturbation is added: as in the
paper, the asymmetry of the unstructured mesh seeds the mode, which is why
the disk is meshed by Gmsh rather than as a structured O-grid.

What the discretisation is, and the two things it needed:

* one CG-FEM ``"vec"`` space of four Q2 components on the same mesh, the
  weak form of :class:`ReducedMHDWeakForm` (the brackets as
  ``grad b . grad^perp a``, the two Laplacians as weak constraints);
* Crank-Nicolson with Newton-BiCGSTAB at every step (the linearised system
  is not symmetric: brackets), ``dt = 1`` as in the paper's figure 6;
* the two constraint rows carry NO time derivative -- this is a DAE, not an
  ODE -- which ``TimeDiscreteFEscheme``'s ``mass_weights = (1, 0, 0, 1)``
  expresses; an explicit stage cannot advance it and is refused;
* a physics-based preconditioner, the block LU of Philip, Chacón and
  Pernice (JCP 2008) written as :class:`BlockSchurPreconditioner`: the two
  constraints are inverted with V-cycles on the (constant) Laplacian and
  mass hierarchies, and ``psi`` through the parabolised Schur complement
  :class:`ReducedMHDSchurWeakForm`, relinearised at every step. Without it
  the constraints ``omega = Delta phi``, ``J = Delta psi`` give the
  Jacobian the conditioning of a bi-Laplacian: at Q2 the unpreconditioned
  BiCGSTAB no longer converges at all (NaN within 110 steps), with it a
  step takes 2-3 Krylov iterations;
* a KEPT iteration matrix (``jacobian_steps``/``jacobian_tol``): the stage
  Jacobian moves by 3e-5 in norm over 100 steps here, so reassembling it at
  every step -- 52 ms of the 116 a step then costs -- buys nothing. It is
  looked at every 10 steps and rebuilt only when one matrix-free product
  says it has drifted past 2e-4, which takes a step to ~35 ms without
  moving the energies in nine digits. The mesh is therefore a NESTED ladder
  (a coarse Gmsh disk refined 1-to-4, ``mesh_hierarchy``), which is what
  the multigrid needs.

Diagnostics: the magnetic and kinetic energies ``1/2 int |grad psi|^2`` and
``1/2 int |grad phi|^2`` in time -- the kinetic one grows exponentially
during the linear stage, its slope is the growth rate -- and snapshots of
``J`` through the kink (the paper's figure 7).

The sequence is the paper's figure 7 -- the profile axisymmetric for a long
while, the hole displaced off axis, the negative-current island pushed off
centre with current sheets around it, then a wider flatter hole -- but WHEN
it happens is not a property of the equations here: with no perturbation
added, the mode grows from the asymmetry of the mesh and the round-off of
the solver, so the onset is ``ln(A_sat / A_0) / gamma`` and moves with the
seed. Measured at Q2 on this ladder: the kinetic energy passes 2.6e-07 at
t = 2000, still rising, with a growth rate ``gamma = 2.5e-03``; the earlier
Q1 run (Gmsh at h = 0.05, unpreconditioned, a noisier seed) reached its
saturation by t = 1500 with ``gamma = 2.9e-03``. Hence ``T_FINAL = 3000``
and snapshots spanning 2000-3000. The magnetic energy loses about 1e-3 of
itself over the run, the resistive dissipation at ``eta = 1e-5``.

⚠ What the mesh leaves in ``J`` (a second derivative, so 1/h^2 times
anything the mapping does): the irregular vertices of the recombined quad
mesh show up as grid-scale ripple, 4.4e-03 rms against a 0.22 range, with
four hot spots where the coarse Gmsh mesh has valence-3 corners. Measured:
doubling the quadrature changes it by 3 %, so it is the mesh and not
aliasing -- and it is also precisely what seeds the kink.

Needs the optional ``mesh`` extra (gmsh).
"""

import time

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_lagrange_basis,
    local_lagrange_basis_by_logical,
)
from scimba_jax.linear_approximation.basis.dof_map import UnstructuredLagrangeDofMap
from scimba_jax.linear_approximation.basis.general_bases import AnalyticBasis
from scimba_jax.linear_approximation.error_analysis import h1_seminorm_error
from scimba_jax.linear_approximation.galerkin.fem.elliptic_fe_scheme import (
    EllipticFEscheme,
)
from scimba_jax.linear_approximation.galerkin.fem.time_discrete_fe_scheme import (
    TimeDiscreteFEscheme,
    stage_scheme,
)
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.multigrid import MG
from scimba_jax.linear_approximation.solvers.newton import NewtonKrylov
from scimba_jax.linear_approximation.solvers.schur import (
    BlockSchurPreconditioner,
    MultigridInverse,
)
from scimba_jax.linear_approximation.solvers.smoothers import DampedJacobiSmoother
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 macro_mesh_disk
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.laplacian_weak_form import (
    LaplacianWeakForm,
)
from scimba_jax.physical_models.classical_weakform.reduced_mhd_weak_form import (
    MASS_WEIGHTS,
    ReducedMHDSchurWeakForm,
    ReducedMHDWeakForm,
)
from scimba_jax.plots.plots_galerkin import sample_solution
from scimba_jax.time_discrete.butcher_tableau import build_crank_nicolson_tableau

# The paper's data, unchanged.
J1, J2 = 0.2, 0.2666
ETA, NU = 1e-5, 1e-6
T_FINAL = 3000.0
DT = 1.0
#: Gmsh size of the COARSEST level; the finest is ``MESH_SIZE / 2**(N_LEVELS-1)``
#: (~0.05 here, ~1 700 Q2 cells: the paper's figure 6 uses ~0.05 at Q1).
MESH_SIZE = 0.21
N_LEVELS = 3
ORDER = 2
#: Components, for the preconditioner's sweep.
PSI, PHI, J, OMEGA = 0, 1, 2, 3
#: Times at which J is drawn (the paper's figure 7 spans the kink).
SNAPSHOTS = (0.0, 1000.0, 2000.0, 2400.0, 2700.0, 3000.0)
CHUNK = 100  # steps per compiled scan; the energies are read between chunks


def driven_current(x):
    r2 = x[0] * x[0] + x[1] * x[1]
    return J1 * (1.0 - r2 * r2) - J2 * (1.0 - r2) ** 8


def make_mesh_hierarchy(mesh_size=MESH_SIZE, n_levels=N_LEVELS, order=ORDER):
    """Nested disks: a coarse Gmsh mesh (unstructured, hence not symmetric)
    split 1-to-4 ``n_levels - 1`` times, as the multigrid wants; the last
    level is the computational mesh. Isoparametric: the geometry has the
    degree of the unknowns."""
    return mesh_hierarchy(
        macro_mesh_disk(radius=1.0, order=order, mesh_size=mesh_size),
        n_levels=n_levels,
    )


def make_mesh(macro, order=ORDER):
    # Gauss with `order + 1` points per direction: exact for the mass and
    # the Laplacians of Q_order, the brackets slightly under-integrated --
    # the usual variational crime. Measured on a stage's element matrices
    # at Q1: 7.8 ms at 4 points, 19.5 ms at the 16 of `2 * order + 2`, for
    # the same solution.
    return UnstructuredMesh.from_macro_mesh(
        macro, UnitSquareTensorized(dim=2, order=2 * order)
    )


def make_variables(mesh, n_components, order=ORDER):
    basis = AnalyticBasis(
        nb_basis=(order + 1) ** 2,
        out_dim=n_components,
        mesh=mesh,
        local_basis=lambda y, i, m: local_lagrange_basis(
            y, i, m, order=order, out_dim=n_components
        ),
        local_basis_by_logical=lambda y, i, m: local_lagrange_basis_by_logical(
            y, i, m, order=order, out_dim=n_components
        ),
        basis_type="scalar" if n_components == 1 else "vec",
    )
    return VariablesFE(
        basis=basis, nb_variables=n_components, dof_map=UnstructuredLagrangeDofMap
    )


def initial_flux(mesh):
    """``psi_0 = Delta^-1 J_c``: the Grad-Shafranov equilibrium of the profile.

    ``-Delta psi_0 = -J_c`` with ``psi_0 = 0`` on the circle, solved once by
    the stationary scheme; its nodal DOFs are the first component's.
    """
    scalar = make_variables(mesh, 1)
    model = AbstractPhysicalWeakModel.from_weak_form(
        LaplacianWeakForm(dim=2, f=lambda x: -driven_current(x)),
        dirichlet=lambda x: jnp.zeros(1),
    )
    scheme = EllipticFEscheme.solve(EllipticFEscheme(model, scalar), max_iter=1)
    return scheme, scheme.variables.dofsl


def _zero(x):
    """The reference of the seminorm: a NAMED function, so that the jitted
    norm keeps one compiled program -- a fresh ``lambda`` at every call is a
    new key, and was measured as 2 recompilations per call (0.8 s)."""
    return 0.0


def energies(scalar_scheme, dofsl):
    """``1/2 int |grad psi|^2`` and ``1/2 int |grad phi|^2`` from the DOFs.

    The H1 seminorm of each component, read through a scalar scheme on the
    same mesh and numbering (the one that solved for ``psi_0``).
    """
    out = []
    for k in (0, 1):
        scalar_scheme.variables.dofsl = dofsl[:, k : k + 1]
        out.append(0.5 * float(h1_seminorm_error(scalar_scheme, _zero)) ** 2)
    return out


def _scalar_scheme(form, variables):
    model = AbstractPhysicalWeakModel.from_weak_form(
        form, dirichlet=lambda x: jnp.zeros(1)
    )
    return EllipticFEscheme(model, variables)


def make_preconditioner(hierarchy, variables, a, eta=ETA, order=ORDER):
    """The block LU of the stage Jacobian, each block inverted by one V-cycle.

    Sweep ``J, omega, phi, psi, J, omega, phi``: the predictor eliminates the
    two constraints and the vorticity at ``delta psi = 0``, ``psi`` is solved
    through the parabolised Schur complement, and the corrector redoes the
    rest with ``delta psi`` known; the couplings come from the stage's own
    element matrices.

    ⚠ **The two CONSTRAINT rows are built ONCE for the whole run.** ``J =
    Delta psi`` and ``omega = Delta phi`` are linear, so their stage blocks
    (``a M`` and ``a K``) never move; the assembly cannot know that, the
    physicist can, and :meth:`MultigridInverse.frozen_component` says it --
    no Galerkin coarsening, no coarse factorisation, no diagonal probe for
    them at any step. Only ``omega`` (whose block carries ``[., phi]``) and
    ``psi`` are relinearised.

    ⚠ **Every diagonal block is READ OFF those same element matrices**
    (``component=``), not rebuilt from a weak form. It is the true block --
    ``a M`` on the ``J`` row, ``a K`` on the ``phi`` row, ``M + a nu K -
    a [., phi]`` on ``omega`` -- and in particular it carries the stage's
    factor ``a = a_ii dt``, which a hand-written mass or Laplacian does not:
    inverting the unscaled ones cost 7 BiCGSTAB iterations per step against
    3. Only what eliminating ``J`` and ``phi`` ADDS to the ``psi`` row
    (``a eta K`` and the Alfven term ``a^2 (B.grad)(B.grad)``,
    :class:`ReducedMHDSchurWeakForm` with ``diagonal=False``) is assembled,
    once per Newton step.

    The hierarchies -- meshes, transfers, coarse schemes -- are built ONCE
    here; a Laplacian is what fills them, but only their structure is used,
    every level's operator being the stage's block and its Galerkin
    coarsenings.

    Args:
        hierarchy: The nested macro-meshes.
        variables: The four-component space on the finest level.
        a: ``a_ii * dt`` of the implicit stage (``dt / 2`` for Crank-Nicolson).
        eta: The resistivity.
        order: The degree of the unknowns.

    Returns:
        The :class:`BlockSchurPreconditioner`.
    """
    cache = {}

    def scalar_variables(macro):
        if id(macro) not in cache:
            cache[id(macro)] = make_variables(make_mesh(macro, order), 1, order)
        return cache[id(macro)]

    def laplace(macro):
        return _scalar_scheme(
            LaplacianWeakForm(dim=2, f=lambda x: 0.0), scalar_variables(macro)
        )

    finest = hierarchy[hierarchy.n_levels - 1]

    def schur_complement(dofsl_lin):
        return _scalar_scheme(
            ReducedMHDSchurWeakForm(variables, dofsl_lin, a, eta, diagonal=False),
            scalar_variables(finest),
        )

    smoother = DampedJacobiSmoother(2.0 / 3.0)

    def cycle():
        return MG(build_hierarchy_nested(laplace, hierarchy), smoother)

    # The stage operator itself, as a probe: `a_ii * dt = a` is all that
    # enters the bilinear form, and the linear form is not read (see
    # `stage_scheme`). The two constraint rows are read off it once.
    probe = stage_scheme(
        ReducedMHDWeakForm(eta=eta, nu=NU, j_c=driven_current),
        variables,
        a,
        1.0,
        dirichlet=lambda x: jnp.zeros(4),
        mass_weights=MASS_WEIGHTS,
    )
    zeros = jnp.zeros_like(variables.dofsl)
    return BlockSchurPreconditioner(
        sweep=[J, OMEGA, PHI, PSI, J, OMEGA, PHI],
        inverses={
            J: MultigridInverse.frozen_component(cycle(), probe, zeros, J),
            PHI: MultigridInverse.frozen_component(cycle(), probe, zeros, PHI),
            OMEGA: MultigridInverse(cycle(), component=OMEGA),
            PSI: MultigridInverse(
                cycle(), component=PSI, scheme_factory=schur_complement
            ),
        },
    )


def make_time_scheme(variables, preconditioner, eta=ETA, nu=NU, dt=DT):
    return TimeDiscreteFEscheme(
        spatial_weak_form_factory=ReducedMHDWeakForm(
            eta=eta, nu=nu, j_c=driven_current
        ),
        variables=variables,
        butcher_tableau=build_crank_nicolson_tableau(),
        dt=dt,
        dirichlet=lambda x: jnp.zeros(4),
        mass_weights=MASS_WEIGHTS,
        # The brackets make the linearised system non-symmetric: BiCGSTAB,
        # not the default CG. Newton converges in ONE step at dt = 1, so the
        # whole cost of a step is the Krylov loop. Measured at Q2 on the
        # 3-level ladder (6 841 nodes, 27 364 DOFs): without preconditioner
        # a stage costs thousands of BiCGSTAB iterations and the time
        # integration gives NaN within 110 steps; with it, 1 Newton step and
        # 2 to 3 BiCGSTAB iterations, ~120 ms per step.
        # 1e-7 keeps the energies and the growth rate
        # (measured at Q1: 897 iterations at 1e-9, 291 at 1e-7 for a CN step
        # whose own error is far above either).
        # ⚠ ``solver=`` and not ``preconditioner=``: the scheme forwards the
        # unknown kwargs as one dict, which a pytree-valued preconditioner
        # would make dynamic together with the integer caps next to it (see
        # ``2d_anisotropic_diffusion_disk_preconditioners.make_solver``).
        solver=NewtonKrylov(
            cg_solver="bicgstab",
            max_iter=8,
            max_iter_linear=200,
            tol=1e-7,
            preconditioner=preconditioner,
        ),
        # The kept iteration matrix: the stage operator is looked at every
        # 10 steps and reassembled only if one matrix-free product says it
        # has drifted by more than 2e-4. Measured here: 116 -> ~35 ms per
        # step, the same 2-3 BiCGSTAB iterations, the same energies to nine
        # digits. The probe (5.7 ms against 52 for an assembly) reads 3e-5
        # through the linear stage and 5e-4 once the kink develops, so the
        # operator is rebuilt a handful of times over the 2 000 steps --
        # early on because the equilibrium settles, then when the mode grows.
        jacobian_steps=10,
        jacobian_tol=2e-4,
    )


def run(
    t_final=T_FINAL,
    dt=DT,
    mesh_size=MESH_SIZE,
    n_levels=N_LEVELS,
    eta=ETA,
    nu=NU,
    snapshots=SNAPSHOTS,
):
    hierarchy = make_mesh_hierarchy(mesh_size, n_levels)
    mesh = make_mesh(hierarchy[n_levels - 1])
    scalar_scheme, psi0 = initial_flux(mesh)
    variables = make_variables(mesh, 4)
    start = time.time()
    preconditioner = make_preconditioner(hierarchy, variables, 0.5 * dt, eta)
    print(f"preconditioner hierarchies: {time.time() - start:.1f} s")
    scheme = make_time_scheme(variables, preconditioner, eta, nu, dt)
    # (psi, phi, J, omega) at t = 0: the equilibrium flux, no flow, J = J_c.
    dofsl = scheme.initialize(lambda x: jnp.array([0.0, 0.0, driven_current(x), 0.0]))
    dofsl = dofsl.at[:, 0].set(psi0[:, 0])
    print(
        f"current hole: {mesh.n_cells_total} Gmsh Q{ORDER} cells on {n_levels} "
        f"nested levels, {dofsl.size} DOFs, dt = {dt}, T = {t_final}, "
        f"eta = {eta}, nu = {nu}"
    )

    n_steps = int(round(t_final / dt))
    e_mag, e_kin = energies(scalar_scheme, dofsl)
    times, magnetic, kinetic = [0.0], [e_mag], [e_kin]
    frames = {0.0: np.asarray(dofsl)}
    wanted = sorted(t for t in snapshots if t > 0)
    start = time.time()
    t = 0.0
    done = 0
    while done < n_steps:
        n = min(CHUNK, n_steps - done)
        dofsl, _ = scheme.solve(dofsl, t0=t, nt=n, keep_history=False)
        jax.block_until_ready(dofsl)
        done += n
        t = done * dt
        e_mag, e_kin = energies(scalar_scheme, dofsl)
        times.append(t)
        magnetic.append(e_mag)
        kinetic.append(e_kin)
        if wanted and t >= wanted[0] - 1e-9:
            frames[wanted.pop(0)] = np.asarray(dofsl)
        if done % (CHUNK * 4) == 0 or done == n_steps:
            print(
                f"  t = {t:7.1f}  E_mag = {e_mag:.6f}  E_kin = {e_kin:.3e}"
                f"  ({time.time() - start:.0f} s)"
            )
    return scheme, dofsl, np.array(times), np.array(magnetic), np.array(kinetic), frames


if __name__ == "__main__":
    scheme, dofsl, times, magnetic, kinetic, frames = run()

    # Growth rate of the kink: the slope of log E_kin over the linear stage,
    # taken where the kinetic energy is well above round-off and well below
    # its saturation.
    mask = (kinetic > 1e-10 * kinetic.max()) & (kinetic < 1e-2 * kinetic.max())
    if mask.sum() > 3:
        slope = np.polyfit(times[mask], np.log(kinetic[mask]), 1)[0]
        print(
            f"\nlinear stage: E_kin ~ exp({slope:.3e} t), growth rate {slope / 2:.3e}"
        )

    # ── Figures: J through the kink, and the energies ─────────────────────
    scalar_scheme = EllipticFEscheme.solve(
        EllipticFEscheme(
            AbstractPhysicalWeakModel.from_weak_form(
                LaplacianWeakForm(dim=2, f=lambda x: -driven_current(x)),
                dirichlet=lambda x: jnp.zeros(1),
            ),
            make_variables(scheme.variables.mesh, 1),
        ),
        max_iter=1,
    )
    points, triangles, _ = sample_solution(scalar_scheme, n_side=4)

    def sampled(nodal):
        scalar_scheme.variables.dofsl = jnp.asarray(nodal)[:, None]
        return sample_solution(scalar_scheme, n_side=4)[2]

    snaps = sorted(frames.items())
    # The four unknowns, one row each, at the same instants, on one colour
    # map: psi and J on the range of their initial state (the flux barely
    # moves, the current reorganises inside it), phi and omega symmetric
    # about zero on the largest value they reach (they start at zero).
    names = ("psi", "phi", "J", "omega")
    all_frames = np.stack([nodal for _, nodal in snaps])
    ranges = []
    for column in range(4):
        if column in (0, 2):
            low, high = all_frames[0, :, column].min(), all_frames[0, :, column].max()
        else:
            high = np.abs(all_frames[:, :, column]).max()
            low = -high
        margin = 0.05 * (high - low) if high > low else 1e-12
        ranges.append(np.linspace(low - margin, high + margin, 41))
    figure, axes = plt.subplots(
        4, len(snaps), figsize=(2.6 * len(snaps), 10.4), constrained_layout=True
    )
    for column, (name, levels, row) in enumerate(zip(names, ranges, axes)):
        for axis, (t_snap, nodal) in zip(row, snaps):
            values = sampled(nodal[:, column])
            filled = axis.tricontourf(
                points[:, 0],
                points[:, 1],
                triangles,
                values,
                levels=levels,
                cmap="jet",
                extend="both",
            )
            # Few lines on purpose: J is a second derivative, so the mesh's
            # irregular vertices put ~2 % of grid-scale ripple in it, and a
            # dense set of levels turns that into festoons that read as
            # structure. One line in eight (five per panel) shows the shape.
            axis.tricontour(
                points[:, 0],
                points[:, 1],
                triangles,
                values,
                levels=levels[::8],
                colors="k",
                linewidths=0.3,
            )
            axis.set(aspect="equal", xticks=[], yticks=[])
            if column == 0:
                axis.set_title(f"t = {t_snap:.0f}")
        row[0].set_ylabel(name, fontsize=12)
        figure.colorbar(filled, ax=row.tolist(), shrink=0.9)
    figure.suptitle("current hole: psi, phi, J, omega through the internal kink")

    figure, axis = plt.subplots(figsize=(8, 4.5))
    axis.semilogy(times, kinetic, label="kinetic  1/2 int |grad phi|^2")
    axis.semilogy(times, magnetic, label="magnetic 1/2 int |grad psi|^2")
    axis.set(
        xlabel="t", ylabel="energy", title="energies: the kink grows, then saturates"
    )
    axis.grid(alpha=0.25)
    axis.legend()
    plt.show()