"""The lopsided domain as a nested hierarchy of cell graphs.

Shared by the graph-network examples of this directory. Gmsh meshes the
coarsest level of the lopsided D of ``mesh/irregular_domains.py`` 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 sys
from pathlib import Path

import numpy as np

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
sys.path.insert(0, str(_HERE.parents[1] / "mesh"))
from irregular_domains import lopsided_wall  # noqa: E402, F401  (re-exported)


def lopsided_arrays(n_levels: int = 2, mesh_size: float = 0.13) -> tuple:
    """The cached nodes, cells and parent maps of the lopsided hierarchy.

    Gmsh meshes the coarsest level once; the arrays are cached next to this
    file.

    Args:
        n_levels: how many nested levels.
        mesh_size: Gmsh's characteristic length at the COARSEST level.

    Returns:
        ``(nodes, cells, parents)``: one array per level, coarsest first, and
        one fine-to-coarse map per transition.
    """
    cache = _HERE / f"lopsided_hierarchy_h{mesh_size}_L{n_levels}.npz"
    if not cache.exists():
        hierarchy = mesh_hierarchy_from_points(
            lopsided_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)
    return (
        [stored[f"nodes_{k}"] for k in range(n_levels)],
        [stored[f"cells_{k}"] for k in range(n_levels)],
        [stored[f"parents_{k}"] for k in range(n_levels - 1)],
    )


def hierarchy_from_arrays(
    nodes,
    cells,
    parents,
    ball_radii=(),
    ball_distance: str = "euclidean",
    cross_radii=(),
    quad_order: int = 2,
    capacities=None,
    cross_capacity=None,
) -> HierarchicMeshData:
    """A hierarchy of cell graphs from nodes, cells and parent maps.

    Args:
        nodes: one ``(n_nodes, 2)`` array per level, coarsest first.
        cells: one ``(n_cells, 4)`` array per level.
        parents: one fine-to-coarse map per transition.
        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.
        quad_order: the cell quadrature order.
        capacities: one ``MeshData`` capacity per level, or ``None``.
        cross_capacity: the cross balls' declared sizes, or ``None``.

    Returns:
        The hierarchy.
    """
    n_levels = len(nodes)
    quad = UnitSquareTensorized(dim=2, order=quad_order)
    per_level = _radii_per_level(ball_radii, n_levels)
    capacities = capacities or [None] * n_levels
    levels = [
        MeshData(
            dim=2,
            mesh=UnstructuredMesh(
                nodes=nodes[k], cells=cells[k], ref_quad=quad, order=1
            ),
            entity="cells",
            ball_radii=per_level[k],
            ball_distance=ball_distance,
            capacity=capacities[k],
        )
        for k in range(n_levels)
    ]
    return HierarchicMeshData(levels, parents, cross_radii, cross_capacity)


def lopsided_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 the lopsided domain, 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.
    """
    nodes, cells, parents = lopsided_arrays(n_levels, mesh_size)
    return hierarchy_from_arrays(
        nodes, cells, parents, ball_radii, ball_distance, cross_radii, quad_order
    )