"""This package implements the layers for different Graph Neural Networks."""

from __future__ import annotations

from matgl.config import BACKEND
from matgl.layers._activations import ActivationFunction
from matgl.layers._bond import BondExpansion
from matgl.layers._core import (
    MLP,
    GatedEquivariantBlock,
    GatedMLP,
    build_gated_equivariant_mlp,
)
from matgl.layers._norm import GraphNorm

if BACKEND == "DGL":
    from matgl.layers._atom_ref_dgl import AtomRef
    from matgl.layers._basis import FourierExpansion, RadialBesselFunction, SphericalBesselWithHarmonics
    from matgl.layers._core_dgl import EdgeSet2Set, GatedMLPNorm, MLPNorm
    from matgl.layers._embedding_dgl import EmbeddingBlock, NeighborEmbedding, TensorEmbedding
    from matgl.layers._graph_convolution_dgl import (
        CHGNetAtomGraphBlock,
        CHGNetBondGraphBlock,
        CHGNetGraphConv,
        CHGNetLineGraphConv,
        M3GNetBlock,
        M3GNetGraphConv,
        MEGNetBlock,
        MEGNetGraphConv,
        TensorNetInteraction,
    )
    from matgl.layers._readout_dgl import (
        AttentiveFPReadout,
        GlobalPool,
        ReduceReadOut,
        Set2SetReadOut,
        WeightedAtomReadOut,
        WeightedReadOut,
        WeightedReadOutPair,
    )
    from matgl.layers._three_body import ThreeBodyInteractions
    from matgl.layers._zbl_dgl import NuclearRepulsion
else:
    from matgl.layers._atom_ref_pyg import AtomRefPyG  # type: ignore[assignment]
    from matgl.layers._embedding_pyg import TensorEmbedding  # type: ignore[assignment]
    from matgl.layers._graph_convolution_pyg import TensorNetInteraction  # type: ignore[assignment]
    from matgl.layers._readout_pyg import (  # type: ignore[assignment]
        ReduceReadOut,
        WeightedAtomReadOut,
        WeightedReadOut,
    )
    from matgl.layers._zbl_pyg import NuclearRepulsionPyG  # type: ignore[assignment]
