"""-Delta u = 1 on unstructured quadrilateral meshes read from GMSH.

Three domains, none of them a box: the unit disk, and the poloidal
cross-sections of JET and MAST read from their EQDSK wall files. The point is
that nothing in the discretization knows about it. ``UnstructuredMesh`` gives
the bases the same two things a Cartesian mesh gives them -- reference square
to cell and back -- and ``UnstructuredLagrangeDofMap`` reads the gluing from
the file instead of deriving it from a grid size. The Lagrange basis, the
scheme, the solve are unchanged.

The elements are **isoparametric**: the geometry of a cell is interpolated
with the very same shape functions as the unknowns on it, so a degree-3 mesh
follows a curved wall with degree-3 accuracy and costs nothing extra to
describe.

On the disk the exact solution is ``(1 - r^2)/4``, which gives the error table
below; on the tokamaks there is none, so the check is convergence under
refinement plus the two things that must hold whatever the domain: ``u = 0``
on the boundary exactly, and ``u > 0`` inside by the maximum principle.

Needs the optional ``mesh`` extra (pygmsh/meshio/gmsh). Saves
``laplacian_unstructured.png`` next to this file and shows it: the meshes
on the top row, the solutions below.
"""

import re
import time
from pathlib import Path

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

from scimba_jax.domains.tokamak import read_eqdsk_wall
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.galerkin.dg.elliptic_dg_scheme import (
    EllipticDGscheme,
)
from scimba_jax.linear_approximation.galerkin.dg.flux import SIPGFlux
from scimba_jax.linear_approximation.galerkin.fem.elliptic_fe_scheme import (
    EllipticFEscheme,
)
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_dg import VariablesDG
from scimba_jax.linear_approximation.variables.variables_fe import VariablesFE
from scimba_jax.mapping.macro_mesh import macro_mesh_from_points, macro_mesh_ogrid_disk
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.plots.plots_galerkin import sample_solution

_HERE = Path(__file__).resolve().parent
_DATA = _HERE.parents[5] / "src/scimba_jax/domains/tokamak/data"
_TOKAMAK_HELPERS = _HERE.parents[3] / "mesh/meshes_gmesh/example_macro_mesh_tokamak.py"


def _wall_helpers():
    """``_dedupe_polygon`` and ``_densify_polygon`` from the meshing example.

    Lifted rather than re-derived: the EQDSK wall needs both before GMSH can
    mesh it (near-duplicate points break the 1-D mesher, sparse segments make
    the spline bulge away from the true wall), and that reasoning belongs to
    the example that established it.
    """
    source = _TOKAMAK_HELPERS.read_text()
    namespace = {"np": np}
    for name in ("_dedupe_polygon", "_densify_polygon"):
        match = re.search(rf"def {name}\(.*?(?=\ndef |\n# %%)", source, re.S)
        exec(match.group(0), namespace)  # noqa: S102
    return namespace["_dedupe_polygon"], namespace["_densify_polygon"]


def cached_tokamak_mesh(name, path, mesh_size, order):
    """Mesh once, save, and reuse -- because GMSH does not repeat itself.

    Calling ``macro_mesh_from_points`` twice on the same wall gives the same
    number of cells and *different coordinates*: measured over four calls, four
    distinct meshes. Fixing ``Mesh.RandomSeed`` is not enough, and the pattern
    says why -- the first call in a fresh process is reproducible, later ones
    are not, so it is GMSH state carried between meshes rather than an unseeded
    draw.

    That is not a cosmetic annoyance here. Most draws give the same answer, but
    some are distorted enough to move it: this example once printed 0.384 for
    JET and 2.711 the next run, same code, same cell count -- and DG, far more
    sensitive to cell shape, reached 430. An example whose number changes each
    time is not an example, and a refinement study over three unrelated meshes
    is not a refinement study.

    Args:
        name: Machine name, used for the cache file.
        path: EQDSK file.
        mesh_size: Target characteristic length.
        order: Geometric degree.

    Returns:
        ``(nodes, cells)``, from the cache if it exists.
    """
    cache = _HERE / f"mesh_{name.lower()}_p{order}_h{mesh_size}.npz"
    if cache.exists():
        stored = np.load(cache)
        return stored["nodes"], stored["cells"]

    dedupe, densify = _wall_helpers()
    radial, vertical = read_eqdsk_wall(str(path))
    wall = densify(dedupe(np.stack([radial[:-1], vertical[:-1]], axis=1)), max_seg=0.03)
    macro = macro_mesh_from_points(wall, order=order, mesh_size=mesh_size, smooth=True)
    np.savez_compressed(cache, nodes=macro.nodes, cells=macro.cells)
    print(f"  meshed {name} and saved {cache.name}")
    return macro.nodes, macro.cells


def solve_laplacian(mesh, order):
    """Solve ``-Delta u = 1`` with ``u = 0`` on the whole boundary."""
    basis = AnalyticBasis(
        nb_basis=(order + 1) ** 2,
        out_dim=1,
        mesh=mesh,
        local_basis_by_logical=lambda y, i, m: local_lagrange_basis_by_logical(
            y, i, m, order=order, out_dim=1
        ),
        local_basis=lambda y, i, m: local_lagrange_basis(
            y, i, m, order=order, out_dim=1
        ),
        basis_type="scalar",
    )
    model = AbstractPhysicalWeakModel.from_weak_form(
        LaplacianWeakForm(dim=2, f=lambda x: jnp.ones(())),
        dirichlet=lambda x: jnp.zeros(1),
    )
    variables = VariablesFE(
        basis=basis, nb_variables=1, dof_map=UnstructuredLagrangeDofMap
    )
    return EllipticFEscheme.solve(EllipticFEscheme(model, variables), max_iter=1)


def solve_laplacian_dg(mesh, order):
    """The same problem, discontinuous Galerkin with an interior penalty.

    ``h=None`` is the load-bearing argument. The penalty needs a length scale,
    and on an unstructured mesh a single global one is wrong at both ends at
    once: too weak on the large cells, where coercivity is lost, and too strong
    on the small ones, where it only wrecks the conditioning. Left to None the
    flux takes each face's own size, computed from the cell's quadrature
    weights -- which already carry the mapping's Jacobian, so a curved cell
    needs nothing extra.
    """
    basis = AnalyticBasis(
        nb_basis=(order + 1) ** 2,
        out_dim=1,
        mesh=mesh,
        local_basis_by_logical=lambda y, i, m: local_lagrange_basis_by_logical(
            y, i, m, order=order, out_dim=1
        ),
        local_basis=lambda y, i, m: local_lagrange_basis(
            y, i, m, order=order, out_dim=1
        ),
        basis_type="scalar",
    )
    model = AbstractPhysicalWeakModel.from_weak_form(
        LaplacianWeakForm(dim=2, f=lambda x: jnp.ones(())),
        dirichlet=lambda x: jnp.zeros(1),
    )
    scheme = EllipticDGscheme(
        model,
        VariablesDG(basis=basis, nb_variables=1),
        SIPGFlux(sigma=order * (order + 1), h=None),
    )
    return EllipticDGscheme.solve(scheme, max_iter=1)


# ── The disk, where the answer is known ───────────────────────────────────────

ORDER_DISK = 2
print("-Delta u = 1 on the unit disk,  exact u = (1 - r^2)/4\n")
print(
    f"{'cells':>7}{'FEM dofs':>10}{'FEM error':>13}{'ratio':>8}"
    f"{'DG dofs':>10}{'DG error':>13}{'ratio':>8}{'FEM [s]':>9}{'DG [s]':>9}"
)
print("-" * 86)

previous = None
disk_scheme = None
for refinement in (1, 2, 4, 8):
    start = time.perf_counter()
    macro = macro_mesh_ogrid_disk(radius=1.0, n=refinement, order=ORDER_DISK)
    mesh = UnstructuredMesh.from_macro_mesh(
        macro, UnitSquareTensorized(dim=2, order=2 * ORDER_DISK + 2)
    )
    meshing_time = time.perf_counter() - start

    start = time.perf_counter()
    disk_scheme = solve_laplacian(mesh, ORDER_DISK)
    jax.block_until_ready(disk_scheme.variables.dofsl)
    fem_time = time.perf_counter() - start

    start = time.perf_counter()
    disk_dg = solve_laplacian_dg(mesh, ORDER_DISK)
    jax.block_until_ready(disk_dg.variables.dofsl)
    dg_time = time.perf_counter() - start

    used = np.unique(np.asarray(mesh.cells))
    coordinates = np.asarray(mesh.nodes)[used]
    exact = (1.0 - (coordinates**2).sum(1)) / 4.0
    error = float(
        np.sqrt(np.mean((np.asarray(disk_scheme.variables.dofsl)[:, 0] - exact) ** 2))
    )
    # DG has no nodal values to compare: sampled on the cells instead.
    dg_points, _, dg_values = sample_solution(disk_dg, n_side=6)
    dg_error = float(
        np.sqrt(np.mean((dg_values - (1.0 - (dg_points**2).sum(1)) / 4.0) ** 2))
    )
    ratio = "" if previous is None else f"{previous[0] / error:8.1f}"
    dg_ratio = "" if previous is None else f"{previous[1] / dg_error:8.1f}"
    previous = (error, dg_error)
    print(
        f"{mesh.n_cells_total:>7}{disk_scheme.variables.n_nodes_total:>10}"
        f"{error:>13.4e}{ratio:>8}{disk_dg.variables.dofsl.size:>10}"
        f"{dg_error:>13.4e}{dg_ratio:>8}{fem_time:>9.2f}{dg_time:>9.2f}"
    )

print(
    "\nRatios near 16 = order 4, which is what an isoparametric Q2 gives here."
    "\nDG carries every cell's own DOFs, so it has ~4x more of them for the"
    "\nsame mesh, and its error is sampled inside the cells rather than at"
    "\nnodes -- the two columns are not comparable one to one, their rates are."
    "\nThe solve time is dominated by tracing and compiling, not by the linear"
    "\nalgebra at these sizes -- it is a per-mesh cost, not a per-DOF one."
)

# ── The tokamaks, where it is not ─────────────────────────────────────────────

ORDER_TOKAMAK = 3
tokamaks = {
    "JET": (_DATA / "eqdsk_jet_compare.dat", 0.13),
    "MAST": (_DATA / "g_p49320_t0.60000", 0.13),
}

print("\n-Delta u = 1 on real tokamak cross-sections (EQDSK wall)\n")
print(
    f"{'machine':>8}{'cells':>7}{'area [m2]':>12}{'FEM max u':>12}{'u on bnd':>11}"
    f"{'DG max u':>11}{'FEM [s]':>9}{'DG [s]':>9}"
)
print("-" * 80)

solved, solved_dg = {}, {}
for name, (path, mesh_size) in tokamaks.items():
    start = time.perf_counter()
    nodes, cells = cached_tokamak_mesh(name, path, mesh_size, ORDER_TOKAMAK)
    mesh = UnstructuredMesh(
        nodes=nodes,
        cells=cells,
        ref_quad=UnitSquareTensorized(dim=2, order=2 * ORDER_TOKAMAK + 2),
        order=ORDER_TOKAMAK,
    )
    meshing_time = time.perf_counter() - start

    start = time.perf_counter()
    scheme = solve_laplacian(mesh, ORDER_TOKAMAK)
    jax.block_until_ready(scheme.variables.dofsl)
    solve_time = time.perf_counter() - start
    solved[name] = scheme

    start = time.perf_counter()
    dg_scheme = solve_laplacian_dg(mesh, ORDER_TOKAMAK)
    jax.block_until_ready(dg_scheme.variables.dofsl)
    dg_time = time.perf_counter() - start
    solved_dg[name] = dg_scheme

    values = np.asarray(scheme.variables.dofsl)[:, 0]
    boundary = scheme.variables.boundary_dofs()
    area = sum(
        float(jnp.sum(mesh._local_weights_points(c)[0]))
        for c in range(mesh.n_cells_total)
    )
    _, _, dg_values = sample_solution(dg_scheme, n_side=6)
    print(
        f"{name:>8}{mesh.n_cells_total:>7}{area:>12.4f}{values.max():>12.5f}"
        f"{np.abs(values[boundary]).max():>11.1e}{dg_values.max():>11.5f}"
        f"{solve_time:>9.2f}{dg_time:>9.2f}"
    )

# ── Draw ──────────────────────────────────────────────────────────────────────


def cell_outlines(mesh, n_samples=25):
    """The four edges of every cell, in physical space.

    Sampled through the cell map rather than drawn corner to corner: on a
    degree-3 mesh the edges are cubics, and joining the corners with straight
    lines would draw a different mesh from the one being solved on.
    """
    line = jnp.linspace(0.0, 1.0, n_samples)
    zeros, ones = jnp.zeros_like(line), jnp.ones_like(line)
    edges = [
        jnp.stack([line, zeros], -1),
        jnp.stack([ones, line], -1),
        jnp.stack([line, ones], -1),
        jnp.stack([zeros, line], -1),
    ]

    def one_cell(cell):
        return jnp.stack([mesh._unit_hypercube_to_cell(cell, e) for e in edges])

    return np.asarray(jax.vmap(one_cell)(jnp.arange(mesh.n_cells_total)))


figure, axes = plt.subplots(3, 3, figsize=(16, 15))
panels = [("unit disk", disk_scheme), *solved.items()]
dg_panels = [("unit disk", disk_dg), *solved_dg.items()]

for column, (title, scheme) in enumerate(panels):
    mesh = scheme.variables.mesh
    length_unit = "x" if title == "unit disk" else "R [m]"
    height_unit = "y" if title == "unit disk" else "Z [m]"

    # Top: the mesh actually solved on
    ax = axes[0, column]
    for cell_edges in cell_outlines(mesh):
        for edge in cell_edges:
            ax.plot(edge[:, 0], edge[:, 1], "-", color="#4C78A8", lw=0.6)
    ax.set_title(f"{title} — {mesh.n_cells_total} cells, degree {mesh.order}")
    ax.set_aspect("equal")
    ax.set_xlabel(length_unit)
    ax.set_ylabel(height_unit)

    # Bottom: the solution
    ax = axes[1, column]
    points, triangles, values = sample_solution(scheme)
    drawing = ax.tricontourf(
        points[:, 0], points[:, 1], triangles, values, levels=40, cmap="turbo"
    )
    figure.colorbar(drawing, ax=ax, fraction=0.046)
    ax.set_title(f"{title} — CG-FEM, max u = {values.max():.4f}")
    ax.set_aspect("equal")
    ax.set_xlabel(length_unit)
    ax.set_ylabel(height_unit)

for column, (title, scheme) in enumerate(dg_panels):
    ax = axes[2, column]
    points, triangles, values = sample_solution(scheme)
    drawing = ax.tricontourf(
        points[:, 0], points[:, 1], triangles, values, levels=40, cmap="turbo"
    )
    figure.colorbar(drawing, ax=ax, fraction=0.046)
    ax.set_title(f"{title} — DG SIPG, max u = {values.max():.4f}")
    ax.set_aspect("equal")
    ax.set_xlabel("x" if title == "unit disk" else "R [m]")
    ax.set_ylabel("y" if title == "unit disk" else "Z [m]")

figure.suptitle("-Δu = 1 on unstructured quadrilaterals: mesh, CG-FEM, DG-SIPG")
figure.tight_layout()
figure.savefig(_HERE / "laplacian_unstructured.png", dpi=130)
print(f"\nSaved {_HERE / 'laplacian_unstructured.png'}")
plt.show()