r"""A graph U-Net trained on ONE mesh of the lopsided domain, read on a batch of others in one program.

    -Delta u = f    in the lopsided domain,    u = 0 on the wall

The case of a graph operator used in practice: the mesh is FIXED while it
trains -- a constant of the problem for the gradient, but an argument of the
compiled program, never baked into it -- and at inference it is read on OTHER
meshes, all of them in a single ``vmap``.

The other meshes are the training mesh with its interior nodes jittered, by
``JITTER`` times the cell size of each level; the wall does not move. Same
topology, so the same parent maps; different geometry, so different cell
measures, centres and ball pairs -- every shape a graph would otherwise carry
per mesh. ``MeshData(capacity=...)`` pads them all to one declared size: the
padded nodes have a zero measure and the padded edges and pairs are loops on a
padded node, so the real nodes see exactly their own graph (checked below on
the training mesh, padded against unpadded).

1. the hierarchy of the training mesh (levels 0-3 of the nested hierarchy of the lopsided domain
   of ``gnn_unet_lopsided_laplacian.py``) and ``N_PERTURBED`` jittered copies, all
   padded to the same capacities;
2. the gno U-Net of ``gnn_unet_lopsided_laplacian.py``, trained on the UNPADDED
   training hierarchy -- the padding plays no part in the training;
3. ONE jitted ``vmap`` over the hierarchies and the test sources; the truth is a
   P1 finite-element solve on each jittered mesh.

Run:  python gno_unet_lopsided_perturbed_meshes.py
"""

import math
import sys
import time
from pathlib import Path

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
from scipy.spatial import cKDTree

sys.path.insert(0, str(Path(__file__).resolve().parent))
from lopsided_hierarchy import (  # noqa: E402
    hierarchy_from_arrays,
    lopsided_arrays,
    lopsided_wall,
)
from lopsided_laplacian_data import (  # noqa: E402
    FiniteElementSolver,
    draw_sources,
    relative_error,
)

from scimba_jax.neural_operator.layers.graph_based import (  # noqa: E402
    GraphOperator,
    GraphUNet,
)
from scimba_jax.nonlinear_approximation.numerical_solvers.no_projectors import (  # noqa: E402
    NOProjector,
)
from scimba_jax.nonlinear_approximation.optimizers.schedules import (  # noqa: E402
    CosineDecay,
)

jax.config.update("jax_enable_x64", True)

# ── The network and its training: gnn_unet_lopsided_laplacian.py's, gno mode ─────
N_TRAIN, N_TEST = 128, 32
N_EPOCHS, BATCH_SIZE, LEARNING_RATE = 400, 16, 3e-3
N_STEPS = N_EPOCHS * math.ceil(N_TRAIN / BATCH_SIZE)
HIDDEN, BLOCKS_PER_LEVEL = 16, 2
RADIUS, RADIUS_GROWTH = 0.1, 2.0
MESH_SIZE, N_LEVELS = 0.4, 4

# ── The other meshes ──────────────────────────────────────────────────────────
N_PERTURBED = 4
JITTER = 0.15  # interior nodes moved by up to this fraction of the cell size

# ══════════════════════════════════════════════════════════════════════════════
# 1. The training mesh and its jittered copies, padded to one capacity
# ══════════════════════════════════════════════════════════════════════════════
nodes, cells, parents = lopsided_arrays(N_LEVELS + 1, MESH_SIZE)
nodes, cells, parents = nodes[:N_LEVELS], cells[:N_LEVELS], parents[: N_LEVELS - 1]
radii = [RADIUS * RADIUS_GROWTH**k for k in range(N_LEVELS)]  # finest ... coarsest
ball_radii = [(radii[N_LEVELS - 1 - k],) for k in range(N_LEVELS)]
cross_radii = tuple(radii[: N_LEVELS - 1])


def wall_nodes(cell_nodes: np.ndarray, points: np.ndarray) -> np.ndarray:
    """Nodes on the domain's boundary: on an edge that only one cell owns.

    Args:
        cell_nodes: ``(n_cells, 4)``.
        points: ``(n_nodes, 2)``.

    Returns:
        The boundary node indices.
    """
    # the four nodes of a cell, in cyclic order around its centre
    centre = points[cell_nodes].mean(axis=1, keepdims=True)
    gap = points[cell_nodes] - centre
    order = np.argsort(np.arctan2(gap[..., 1], gap[..., 0]), axis=1)
    cyclic = np.take_along_axis(cell_nodes, order, axis=1)
    edges = np.sort(
        np.stack([cyclic, np.roll(cyclic, -1, axis=1)], axis=-1).reshape(-1, 2), axis=1
    )
    unique, count = np.unique(edges, axis=0, return_counts=True)
    return np.unique(unique[count == 1])


def jittered(k: int, seed: int) -> np.ndarray:
    """Level ``k``'s nodes, the interior ones moved, the wall fixed.

    Args:
        k: the level.
        seed: the random state.

    Returns:
        ``(n_nodes, 2)``.
    """
    rng = np.random.default_rng(seed * 100 + k)
    moved = nodes[k] + JITTER * (MESH_SIZE / 2**k) * rng.uniform(-1, 1, nodes[k].shape)
    fixed = wall_nodes(cells[k], nodes[k])
    moved[fixed] = nodes[k][fixed]
    return moved


meshes = [nodes] + [
    [jittered(k, seed) for k in range(N_LEVELS)] for seed in range(1, N_PERTURBED + 1)
]
exact = [
    hierarchy_from_arrays(m, cells, parents, ball_radii, cross_radii=cross_radii)
    for m in meshes
]
# Every size the hierarchies carry, and the largest of each across them.
capacities = [
    {
        "nodes": max(h.levels[k].n_nodes for h in exact) + 1,
        "edges": max(h.levels[k].n_edges for h in exact) + 1,
        "balls": max(int(h.levels[k].ball(ball_radii[k][0])[0].shape[0]) for h in exact)
        + 1,
    }
    for k in range(N_LEVELS)
]
cross_capacity = {
    r: max(
        int(h.cross_ball(t, r)[0].shape[0]) for h in exact for t in range(N_LEVELS - 1)
    )
    + 1
    for r in cross_radii
}
padded = [
    hierarchy_from_arrays(
        m,
        cells,
        parents,
        ball_radii,
        cross_radii=cross_radii,
        capacities=capacities,
        cross_capacity=cross_capacity,
    )
    for m in meshes
]
for k in range(N_LEVELS):
    pairs = [int(h.levels[k].ball(ball_radii[k][0])[0].shape[0]) for h in exact]
    print(
        f"level {k}: {exact[0].levels[k].n_nodes} cells; ball pairs per mesh "
        f"{pairs} -> padded to {capacities[k]['balls']}"
    )

# ══════════════════════════════════════════════════════════════════════════════
# 2. Data on every mesh, training on the first (unpadded)
# ══════════════════════════════════════════════════════════════════════════════
key = jax.random.key(0)
wall_tree = cKDTree(lopsided_wall())
finest = N_LEVELS - 1
candidates = exact[0].finest.point_cloud
draws_train = draw_sources(jax.random.fold_in(key, 1), N_TRAIN, candidates, wall_tree)
draws_test = draw_sources(jax.random.fold_in(key, 2), N_TEST, candidates, wall_tree)
started = time.perf_counter()
solvers = [
    FiniteElementSolver(m[finest], cells[finest], h.finest.point_cloud)
    for m, h in zip(meshes, exact, strict=True)
]
f_train, u_train = solvers[0](*draws_train)
tests = [solver(*draws_test) for solver in solvers]
print(
    f"finite-element data on {len(meshes)} meshes in {time.perf_counter() - started:.0f} s"
)
F_SCALE, U_SCALE = float(jnp.std(f_train)), float(jnp.std(u_train))

network = GraphUNet(
    2,
    1,
    1,
    HIDDEN,
    n_levels=N_LEVELS,
    mode="gno",
    radius=RADIUS,
    radius_growth=RADIUS_GROWTH,
    blocks_per_level=BLOCKS_PER_LEVEL,
    message_kind="kernel",
    update_kind="mlp",
    message_kwargs={"hidden_sizes": [32]},
    key=jax.random.fold_in(key, 3),
)
operator = GraphOperator(network, exact[0])
print(f"\ngno GraphUNet: {network.ndof()} parameters")
projector = NOProjector(
    operator,
    (f_train / F_SCALE, u_train / U_SCALE),
    learning_rate=LEARNING_RATE,
    schedule=CosineDecay(N_STEPS, final_ratio=0.05),
)
started = time.perf_counter()
_, projector = projector.project(
    jax.random.fold_in(key, 4),
    operator,
    N_EPOCHS,
    batch_size=BATCH_SIZE,
    tqdm_desc="  gno",
)
trained = projector.operator.network
print(f"  trained in {time.perf_counter() - started:.0f} s")

# ══════════════════════════════════════════════════════════════════════════════
# 3. One program for every mesh and every test source
# ══════════════════════════════════════════════════════════════════════════════


@jax.jit
def predict_all(network, inputs, hierarchies):
    """The network on a batch of hierarchies, a batch of sources on each.

    Every argument is an ARGUMENT, nothing closed over: the hierarchies are
    operands of the program (see ``predict`` in ``gnn_unet_lopsided_laplacian.py``).

    Args:
        network: the trained U-Net.
        inputs: ``(n_meshes, n_sources, capacity, 1)``.
        hierarchies: the padded hierarchies, stacked.

    Returns:
        ``(n_meshes, n_sources, capacity, 1)``.
    """
    return jax.vmap(
        lambda f, h: jax.vmap(lambda one: network(one, h))(f), in_axes=(0, 0)
    )(inputs, hierarchies)


n_real = [h.finest.n_nodes for h in exact]
capacity = capacities[finest]["nodes"]


def pad(values, n):
    return jnp.concatenate(
        [values, jnp.zeros(values.shape[:-2] + (capacity - n, values.shape[-1]))],
        axis=-2,
    )


inputs = jnp.stack(
    [pad(f / F_SCALE, n) for (f, _), n in zip(tests, n_real, strict=True)]
)
stacked = jax.tree_util.tree_map(lambda *x: jnp.stack(x), *padded)
started = time.perf_counter()
outputs = jax.block_until_ready(predict_all(trained, inputs, stacked))
first = time.perf_counter() - started
started = time.perf_counter()
jax.block_until_ready(predict_all(trained, inputs, stacked))
second = time.perf_counter() - started
print(
    f"\none program, {len(meshes)} meshes x {N_TEST} sources: "
    f"{first:.1f} s with its compilation, {second:.2f} s after"
)

# the padding is invisible: the training mesh, padded, against unpadded
unpadded = jax.jit(lambda n, f, h: jax.vmap(lambda one: n(one, h))(f))(
    trained, tests[0][0] / F_SCALE, exact[0]
)  # arguments, not a closure: a closed-over hierarchy is folded as constants
gap = float(jnp.max(jnp.abs(outputs[0, :, : n_real[0]] - unpadded)))
print(f"training mesh, padded vs unpadded: max difference {gap:.1e}")

errors = []
for k, ((_, u), n) in enumerate(zip(tests, n_real, strict=True)):
    prediction = outputs[k, :, :n] * U_SCALE
    errors.append(float(jnp.mean(relative_error(prediction, u))))
    label = "training mesh" if k == 0 else f"jittered mesh {k}"
    print(f"  {label:<16} relative L2 {errors[-1]:.3e}")

# ══════════════════════════════════════════════════════════════════════════════
# 4. One source on the training mesh and on a jittered one
# ══════════════════════════════════════════════════════════════════════════════
sample, other = 0, 1
u_max = float(max(jnp.max(jnp.abs(tests[m][1][sample])) for m in (0, other)))
figure, axes = plt.subplots(2, 3, figsize=(15, 9))
wall = lopsided_wall()
for row, m in enumerate((0, other)):
    points = np.asarray(exact[m].finest.point_cloud)
    truth = tests[m][1][sample]
    prediction = outputs[m, sample, : n_real[m]] * U_SCALE
    panels = [
        ("u FEM", truth, u_max, "viridis"),
        ("u gno U-Net", prediction, u_max, "viridis"),
        ("|error|", jnp.abs(prediction - truth), None, "magma"),
    ]
    for axis, (title, values, limit, cmap) in zip(axes[row], panels):
        picture = axis.scatter(
            points[:, 0],
            points[:, 1],
            c=np.asarray(values)[:, 0],
            s=4,
            cmap=cmap,
            vmin=0.0 if limit else None,
            vmax=limit,
        )
        axis.plot(*np.vstack([wall, wall[:1]]).T, "k-", linewidth=0.8)
        label = "training mesh" if m == 0 else f"jittered mesh {m}"
        axis.set_title(f"{title}, {label} (rel. L2 {errors[m]:.2f})")
        axis.set_aspect("equal")
        figure.colorbar(picture, ax=axis, shrink=0.8)
figure.suptitle(
    f"gno U-Net trained on one mesh of the lopsided domain, read on {N_PERTURBED} jittered ones "
    f"in one program (jitter {JITTER} h)"
)
figure.tight_layout()
figure.savefig(Path(__file__).with_suffix(".png"), dpi=110)
plt.show()