"""The JET poloidal cross-section as a nested hierarchy of cell graphs.

Shared by the graph-network examples of this directory. Gmsh meshes the
coarsest level of the real JET wall (read from its EQDSK file) once, at
``mesh_size``; every finer level is an exact 1-to-4 split, so the parent map
between levels is exact and the finest level is a genuine refinement of the
coarsest -- which is what lets a network trained on one level be read on the
next. The meshes are cached next to this file; Gmsh does not repeat itself.
"""

from __future__ import annotations

import re
from pathlib import Path

import numpy as np

from scimba_jax.domains.tokamak import read_eqdsk_wall
from scimba_jax.linear_approximation.meshes.unstructured_mesh import UnstructuredMesh
from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized
from scimba_jax.mapping.mesh_hierarchy import mesh_hierarchy_from_points
from scimba_jax.neural_operator.data_for_no.hierarchic_mesh_data import (
    HierarchicMeshData,
    _radii_per_level,
)
from scimba_jax.neural_operator.data_for_no.mesh_data import MeshData

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


def jet_wall(max_seg: float = 0.03) -> np.ndarray:
    """The JET wall as an ordered polygon, deduplicated and densified.

    Args:
        max_seg: the longest segment kept along the wall.

    Returns:
        ``(n_points, 2)``.
    """
    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
    radial, vertical = read_eqdsk_wall(str(_DATA / "eqdsk_jet_compare.dat"))
    return namespace["_densify_polygon"](
        namespace["_dedupe_polygon"](np.stack([radial[:-1], vertical[:-1]], axis=1)),
        max_seg=max_seg,
    )


def jet_hierarchy(
    n_levels: int = 2,
    mesh_size: float = 0.13,
    ball_radii=(),
    ball_distance: str = "euclidean",
    cross_radii=(),
    quad_order: int = 2,
) -> HierarchicMeshData:
    """Nested cell graphs of JET, coarsest first, with their balls precomputed.

    Args:
        n_levels: how many nested levels.
        mesh_size: Gmsh's characteristic length at the COARSEST level.
        ball_radii: physical radii to precompute -- one flat sequence for
            every level, or one sequence per level, coarsest first.
        ball_distance: ``"euclidean"`` or ``"geodesic"``.
        cross_radii: radii of the balls between consecutive levels, for a
            neural-operator pooling.
        quad_order: the cell quadrature order.

    Returns:
        The hierarchy.
    """
    cache = _HERE / f"jet_hierarchy_h{mesh_size}_L{n_levels}.npz"
    if not cache.exists():
        hierarchy = mesh_hierarchy_from_points(
            jet_wall(), n_levels=n_levels, order=1, mesh_size=mesh_size, smooth=True
        )
        arrays = {}
        for k, macro in enumerate(hierarchy.levels):
            arrays[f"nodes_{k}"] = np.asarray(macro.nodes)
            arrays[f"cells_{k}"] = np.asarray(macro.cells)
        for k in range(1, n_levels):
            arrays[f"parents_{k - 1}"] = np.asarray(hierarchy.parent_of(k))
        np.savez_compressed(cache, **arrays)
    stored = np.load(cache)
    quad = UnitSquareTensorized(dim=2, order=quad_order)
    per_level = _radii_per_level(ball_radii, n_levels)
    levels = [
        MeshData(
            dim=2,
            mesh=UnstructuredMesh(
                nodes=stored[f"nodes_{k}"],
                cells=stored[f"cells_{k}"],
                ref_quad=quad,
                order=1,
            ),
            entity="cells",
            ball_radii=per_level[k],
            ball_distance=ball_distance,
        )
        for k in range(n_levels)
    ]
    parents = [stored[f"parents_{k}"] for k in range(n_levels - 1)]
    return HierarchicMeshData(levels, parents, cross_radii)