nels=out_dim_edges,
            out_channels=out_dim,
            num_layers=num_output_layers,
            activation=activation,
        )

    def forward(self, batch: Union[Data, Batch]) -> Union[Data, Batch]:
        r"""
        forward function of the layer
        Parameters:
            batch: pyg Batch graphs to pass through the layer
        Returns:
            batch: pyg Batch graphs
        """
        assert (
            "radius_edge_index" in batch
        ), "radius_edge_index not in batch, make sure to use 3D encoder firstly"
        # (j, i) = edge_index
        i, j, idx_i, idx_j, idx_k, idx_kj, idx_ji = triplets(
            batch.radius_edge_index, num_nodes=batch.feat.size(0)
        )
        x, P = batch.edge_feat, batch.feat
        rbf, sbf = batch.edge_rbf, batch.triplet_sbf

        # apply MLP to node embeddings
        P = self.node_model(P)  # [num_nodes, out_dim]

        # rbf and sbf should be computed during pos encoder
        x = self.interaction_block(x, rbf, sbf, idx_kj, idx_ji)
        P = P + self.output_block(x, rbf, i, num_nodes=batch.feat.size(0))  # [num_nodes, out_dim]

        batch.edge_feat = x  # updated edge features
        batch.feat = P  # updated node features

        return batch

    ############################################################################################################
    @classproperty
    def layer_supports_edges(cls) -> bool:
        r"""
        Return a boolean specifying if the layer type supports edges or not.

        Returns:

            supports_edges: bool
                Always ``True`` for the current class
        """
        return True

    @property
    def layer_inputs_edges(self) -> bool:
        r"""
        Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:

            bool:
                Always ``True`` for the current class
        """
        return True

    @property
    def layer_outputs_edges(self) -> bool:
        r"""
        Abstract method. Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:

            bool:
                Always ``True`` for the current class
        """
        return True

    @property
    def out_dim_factor(self) -> int:
        r"""
        Get the factor by which the output dimension is multiplied for
        the next layer.

        For standard layers, this will return ``1``.

        But for others, such as ``GatLayer``, the output is the concatenation
        of the outputs from each head, so the out_dim gets multiplied by
        the number of heads, and this function should return the number
        of heads.

        Returns:

            int:
                Always ``1`` for the current class
        """
        return 1


# graphium/nn/pyg_layers/gated_gcn_pyg.py

"""
Unit tests for the different layers of graphium/nn/pyg_layers/...

The layers are not thoroughly tested due to the difficulty of testing them

adapated from https://github.com/rampasek/GraphGPS/blob/main/graphgps/layer/gps_layer.py
"""

from typing import Union, Callable
from functools import partial

import torch
import torch.nn as nn

from torch_geometric.nn.conv import MessagePassing
from torch_scatter import scatter
from torch_geometric.data import Data, Batch

from graphium.nn.base_graph_layer import BaseGraphStructure, check_intpus_allow_int
from graphium.nn.base_layers import FCLayer
from graphium.utils.decorators import classproperty


class GatedGCNPyg(MessagePassing, BaseGraphStructure):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        in_dim_edges: int,
        out_dim_edges: int = None,
        activation: Union[Callable, str] = "relu",
        dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
        **kwargs,
    ):
        r"""
        ResGatedGCN: Residual Gated Graph ConvNets
        An Experimental Study of Neural Networks for Variable Graphs (Xavier Bresson and Thomas Laurent, ICLR 2018)
        https://arxiv.org/pdf/1711.07553v2.pdf

        Parameters:

            in_dim:
                Input feature dimensions of the layer

            out_dim:
                Output feature dimensions of the layer, and for the edges

            in_dim_edges:
                Input edge-feature dimensions of the layer

            activation:
                activation function to use in the layer

            dropout:
                The ratio of units to dropout. Must be between 0 and 1

            normalization:
                Normalization to use. Choices:

                - "none" or `None`: No normalization
                - "batch_norm": Batch normalization
                - "layer_norm": Layer normalization
                - `Callable`: Any callable function

        """
        MessagePassing.__init__(self, aggr="add", flow="source_to_target", node_dim=-2)
        BaseGraphStructure.__init__(
            self,
            in_dim=in_dim,
            out_dim=out_dim,
            activation=activation,
            dropout=dropout,
            normalization=normalization,
            **kwargs,
        )

        self._initialize_activation_dropout_norm()

        # Allow int32 in the edge_index
        self.__check_input__ = partial(check_intpus_allow_int, self)
        if out_dim_edges is None:
            out_dim_edges = in_dim_edges

        # Initialize the layers for the gating
        self.A = FCLayer(in_dim, out_dim, activation=None, bias=True)
        self.B = FCLayer(in_dim, out_dim, activation=None, bias=True)
        self.C = FCLayer(in_dim_edges, out_dim, activation=None, bias=True)
        self.D = FCLayer(in_dim, out_dim, activation=None, bias=True)
        self.E = FCLayer(in_dim, out_dim, activation=None, bias=True)

        self.edge_out = FCLayer(
            in_dim=out_dim, out_dim=out_dim_edges, activation=None, dropout=dropout, bias=True
        )

    def forward(
        self,
        batch: Union[Data, Batch],
    ) -> Union[Data, Batch]:
        r"""
        Forward pass the Gated GCN layer
        extract the following from the batch:
        x, node features with dim [n_nodes, in_dim]
        e, edge features with dim [n_edges, in_dim]
        edge_index with dim [2, n_edges]

        Parameters:
            batch: pyg Batch graph to pass through the layer
        Returns:
            batch: pyg Batch graph
        """

        x, e, edge_index = batch.feat, batch.edge_feat, batch.edge_index

        # Apply the linear layers
        Ax = self.A(x)
        Bx = self.B(x)
        Ce = self.C(e)
        Dx = self.D(x)
        Ex = self.E(x)

        # Propagate, and apply norm, activation, dropout
        x, e = self.propagate(edge_index, Bx=Bx, Dx=Dx, Ex=Ex, Ce=Ce, e=e, Ax=Ax)
        x = self.apply_norm_activation_dropout(x, batch_idx=batch.batch)
        e = self.edge_out(e)

        # Output
        batch.feat = x
        batch.edge_feat = e

        return batch

    def message(self, Dx_i: torch.Tensor, Ex_j: torch.Tensor, Ce: torch.Tensor) -> torch.Tensor:
        """
        message function
        Parameters:
            Dx_i: tensor with dimension [n_edges, out_dim]
            Ex_j: tensor with dimension [n_edges, out_dim]
            Ce: tensor with dimension [n_edges, out_dim]
        Returns:
            sigma_ij: tensor with dimension [n_edges, out_dim]
        """
        e_ij = Dx_i + Ex_j + Ce
        sigma_ij = torch.sigmoid(e_ij)

        self.e = e_ij
        return sigma_ij

    def aggregate(
        self,
        sigma_ij: torch.Tensor,
        index: torch.Tensor,
        Bx_j: torch.Tensor,
        Bx: torch.Tensor,
    ) -> torch.Tensor:
        r"""
        aggregation function of the layer
        Parameters:
            sigma_ij: the output from message() function with dim [n_edges, out_dim]
            index: dim [n_edges]
            Bx_j: dim [n_edges, out_dim]
            Bx: dim [n_nodes, out_dim]
        Returns:
            out: dim [n_nodes, out_dim]
        """
        dim_size = Bx.shape[0]  # or None ??   <--- Double check this

        sum_sigma_x = sigma_ij * Bx_j
        numerator_eta_xj = scatter(sum_sigma_x, index, 0, None, dim_size, reduce="sum")

        sum_sigma = sigma_ij
        denominator_eta_xj = scatter(sum_sigma, index, 0, None, dim_size, reduce="sum")

        out = numerator_eta_xj / (denominator_eta_xj + 1e-6)
        return out

    def update(self, aggr_out: torch.Tensor, Ax: torch.Tensor):
        r"""
        update function of the layer
        Parameters:
            aggr_out: the output from aggregate() function with dim [n_nodes, out_dim]
            Ax: tensor with dim [n_nodes, out_dim]
        Returns:
            x: dim [n_nodes, out_dim]
            e_out: dim [n_edges, out_dim_edges]
        """
        x = Ax + aggr_out
        e_out = self.e
        del self.e
        return x, e_out

    @classproperty
    def layer_supports_edges(cls) -> bool:
        r"""
        Return a boolean specifying if the layer type supports edges or not.

        Returns:
            bool:
                Always ``True`` for the current class
        """
        return True

    @property
    def layer_inputs_edges(self) -> bool:
        r"""
        Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:
            bool:
                Always ``True`` for the current class
        """
        return True

    @property
    def layer_outputs_edges(self) -> bool:
        r"""
        Abstract method. Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:
            bool:
                Always ``True`` for the current class
        """
        return True

    @property
    def out_dim_factor(self) -> int:
        r"""
        Get the factor by which the output dimension is multiplied for
        the next layer.

        For standard layers, this will return ``1``.

        But for others, such as ``GatLayer``, the output is the concatenation
        of the outputs from each head, so the out_dim gets multiplied by
        the number of heads, and this function should return the number
        of heads.

        Returns:

            int:
                Always ``1`` for the current class
        """
        return 1


# graphium/nn/pyg_layers/gcn_pyg.py

from typing import Callable, Union
from functools import partial

import torch_geometric.nn as pyg_nn
from torch_geometric.data import Data, Batch

from graphium.nn.base_graph_layer import BaseGraphModule, check_intpus_allow_int
from graphium.utils.decorators import classproperty


class GCNConvPyg(BaseGraphModule):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        activation: Union[Callable, str] = "relu",
        dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
        **kwargs,
    ):
        super().__init__(
            in_dim=in_dim,
            out_dim=out_dim,
            activation=activation,
            dropout=dropout,
            normalization=normalization,
            **kwargs,
        )

        self.model = pyg_nn.GCNConv(
            in_channels=self.in_dim, out_channels=out_dim, add_self_loops=False, normalize=False
        )
        self.model.__check_input__ = partial(check_intpus_allow_int, self)

    def forward(
        self,
        batch: Union[Data, Batch],
    ) -> Union[Data, Batch]:
        r"""
        forward function of the layer
        Parameters:
            batch: pyg Batch graphs to pass through the layer
        Returns:
            batch: pyg Batch graphs
        """
        batch.feat = self.model(batch.feat, batch.edge_index)
        batch.feat = self.apply_norm_activation_dropout(batch.feat, batch_idx=batch.batch)

        return batch

    @classproperty
    def layer_supports_edges(cls) -> bool:
        r"""
        Return a boolean specifying if the layer type supports edges or not.

        Returns:

            supports_edges: bool
                Always ``False`` for the current class
        """
        return False

    @property
    def layer_inputs_edges(self) -> bool:
        r"""
        Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:

            bool:
                Always ``False`` for the current class
        """
        return False

    @property
    def layer_outputs_edges(self) -> bool:
        r"""
        Abstract method. Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:

            bool:
                Always ``False`` for the current class
        """
        return False

    @property
    def out_dim_factor(self) -> int:
        r"""
        Get the factor by which the output dimension is multiplied for
        the next layer.

        For standard layers, this will return ``1``.

        But for others, such as ``GatLayer``, the output is the concatenation
        of the outputs from each head, so the out_dim gets multiplied by
        the number of heads, and this function should return the number
        of heads.

        Returns:

            int:
                Always ``1`` for the current class
        """
        return 1


# graphium/nn/pyg_layers/gin_pyg.py

from typing import Callable, Union, Optional
from functools import partial

import torch_geometric.nn as pyg_nn
from torch_geometric.data import Data, Batch

from graphium.nn.base_graph_layer import BaseGraphModule, check_intpus_allow_int
from graphium.nn.base_layers import MLP
from graphium.utils.decorators import classproperty


class GINConvPyg(BaseGraphModule):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        activation: Union[Callable, str] = "relu",
        dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
        **kwargs,
    ):
        r"""
        GIN: Graph Isomorphism Networks
        HOW POWERFUL ARE GRAPH NEURAL NETWORKS? (Keyulu Xu, Weihua Hu, Jure Leskovec and Stefanie Jegelka, ICLR 2019)
        https://arxiv.org/pdf/1810.00826.pdf

        [!] code uses the pytorch-geometric implementation of GINConv

        Parameters:

            in_dim:
                Input feature dimensions of the layer

            out_dim:
                Output feature dimensions of the layer

            activation:
                activation function to use in the layer

            dropout:
                The ratio of units to dropout. Must be between 0 and 1

            normalization:
                Normalization to use. Choices:

                - "none" or `None`: No normalization
                - "batch_norm": Batch normalization
                - "layer_norm": Layer normalization
                - `Callable`: Any callable function

            init_eps :
                Initial :math:`\epsilon` value, default: ``0``.

            learn_eps :
                If True, :math:`\epsilon` will be a learnable parameter.

        """

        super().__init__(
            in_dim=in_dim,
            out_dim=out_dim,
            activation=activation,
            dropout=dropout,
            normalization=normalization,
            **kwargs,
        )

        gin_nn = MLP(
            in_dim=self.in_dim,
            hidden_dims=self.in_dim,
            out_dim=self.out_dim,
            depth=2,
            activation=self.activation_layer,
            last_activation="none",
            normalization=self.normalization,
            last_normalization="none",
        )

        self.model = pyg_nn.GINConv(gin_nn)
        self.model.__check_input__ = partial(check_intpus_allow_int, self)

    def forward(
        self,
        batch: Union[Data, Batch],
    ) -> Union[Data, Batch]:
        r"""
        forward function of the layer
        Parameters:
            batch: pyg Batch graphs to pass through the layer
        Returns:
            batch: pyg Batch graphs
        """
        batch.feat = self.model(batch.feat, batch.edge_index)
        batch.feat = self.apply_norm_activation_dropout(batch.feat, batch_idx=batch.batch)

        return batch

    @classproperty
    def layer_supports_edges(cls) -> bool:
        r"""
        Return a boolean specifying if the layer type supports edges or not.

        Returns:

            supports_edges: bool
                Always ``False`` for the current class
        """
        return False

    @property
    def layer_inputs_edges(self) -> bool:
        r"""
        Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:

            bool:
                Always ``False`` for the current class
        """
        return False

    @property
    def layer_outputs_edges(self) -> bool:
        r"""
        Abstract method. Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:

            bool:
                Always ``False`` for the current class
        """
        return False

    @property
    def out_dim_factor(self) -> int:
        r"""
        Get the factor by which the output dimension is multiplied for
        the next layer.

        For standard layers, this will return ``1``.

        But for others, such as ``GatLayer``, the output is the concatenation
        of the outputs from each head, so the out_dim gets multiplied by
        the number of heads, and this function should return the number
        of heads.

        Returns:

            int:
                Always ``1`` for the current class
        """
        return 1


class GINEConvPyg(BaseGraphModule):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        in_dim_edges: Optional[int] = None,
        activation: Union[Callable, str] = "relu",
        dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
        **kwargs,
    ):
        r"""
        GINE: Graph Isomorphism Networks with Edges
        Strategies for Pre-training Graph Neural Networks
        Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, Jure Leskovec
        https://arxiv.org/abs/1905.12265

        [!] code uses the pytorch-geometric implementation of GINEConv

        Parameters:

            in_dim:
                Input feature dimensions of the layer

            out_dim:
                Output feature dimensions of the layer

            activation:
                activation function to use in the layer

            dropout:
                The ratio of units to dropout. Must be between 0 and 1

            normalization:
                Normalization to use. Choices:

                - "none" or `None`: No normalization
                - "batch_norm": Batch normalization
                - "layer_norm": Layer normalization
                - `Callable`: Any callable function

            init_eps :
                Initial :math:`\epsilon` value, default: ``0``.

            learn_eps :
                If True, :math:`\epsilon` will be a learnable parameter.

        """

        super().__init__(
            in_dim=in_dim,
            out_dim=out_dim,
            activation=activation,
            dropout=dropout,
            normalization=normalization,
        )

        gin_nn = MLP(
            in_dim=self.in_dim,
            hidden_dims=self.in_dim,
            out_dim=self.out_dim,
            depth=2,
            activation=self.activation_layer,
            last_activation="none",
            normalization=self.normalization,
            last_normalization="none",
        )

        self.model = pyg_nn.GINEConv(gin_nn, edge_dim=in_dim_edges)  # , node_dim=-1)
        self.model.__check_input__ = partial(check_intpus_allow_int, self)

    def forward(
        self,
        batch: Union[Data, Batch],
    ) -> Union[Data, Batch]:
        r"""
        forward function of the layer
        Parameters:
            batch: pyg Batch graphs to pass through the layer
        Returns:
            batch: pyg Batch graphs
        """
        batch.feat = self.model(batch.feat, batch.edge_index, batch.edge_feat)
        batch.feat = self.apply_norm_activation_dropout(batch.feat, batch_idx=batch.batch)

        return batch

    @classproperty
    def layer_supports_edges(cls) -> bool:
        r"""
        Return a boolean specifying if the layer type supports edges or not.

        Returns:

            supports_edges: bool
                Always ``True`` for the current class
        """
        return True

    @property
    def layer_inputs_edges(self) -> bool:
        r"""
        Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:

            bool:
                Always ``True`` for the current class
        """
        return True

    @property
    def layer_outputs_edges(self) -> bool:
        r"""
        Abstract method. Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:

            bool:
                Always ``False`` for the current class
        """
        return False

    @property
    def out_dim_factor(self) -> int:
        r"""
        Get the factor by which the output dimension is multiplied for
        the next layer.

        For standard layers, this will return ``1``.

        But for others, such as ``GatLayer``, the output is the concatenation
        of the outputs from each head, so the out_dim gets multiplied by
        the number of heads, and this function should return the number
        of heads.

        Returns:

            int:
                Always ``1`` for the current class
        """
        return 1


# graphium/nn/pyg_layers/gps_pyg.py

import torch
from copy import deepcopy
from typing import Callable, Union, Optional, Dict, Any
from torch.nn import Module
from torch import Tensor
from torch_geometric.data import Batch
from graphium.nn.base_graph_layer import BaseGraphModule
from graphium.nn.base_layers import FCLayer, MultiheadAttentionMup, MLP
from graphium.nn.pyg_layers import (
    GatedGCNPyg,
    GINConvPyg,
    GINEConvPyg,
    PNAMessagePassingPyg,
    MPNNPlusPyg,
)
from graphium.data.utils import get_keys
from graphium.utils.decorators import classproperty
from graphium.ipu.to_dense_batch import (
    to_dense_batch,
    to_sparse_batch,
    to_packed_dense_batch,
    to_sparse_batch_from_packed,
)
from graphium.ipu.ipu_utils import is_running_on_ipu

PYG_LAYERS_DICT = {
    "pyg:gin": GINConvPyg,
    "pyg:gine": GINEConvPyg,
    "pyg:gated-gcn": GatedGCNPyg,
    "pyg:pna-msgpass": PNAMessagePassingPyg,
    "pyg:mpnnplus": MPNNPlusPyg,
}

ATTENTION_LAYERS_DICT = {
    "full-attention": MultiheadAttentionMup,
    "none": None,
}


class GPSLayerPyg(BaseGraphModule):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        in_dim_edges: Optional[int] = None,
        out_dim_edges: Optional[int] = None,
        activation: Union[Callable, str] = "relu",
        dropout: float = 0.0,
        node_residual: Optional[bool] = True,
        normalization: Union[str, Callable] = "none",
        mpnn_type: str = "pyg:gine",
        mpnn_kwargs=None,
        attn_type: str = "full-attention",
        precision: str = "32",
        biased_attention_key: Optional[str] = None,
        attn_kwargs=None,
        droppath_rate_attn: float = 0.0,
        droppath_rate_ffn: float = 0.0,
        hidden_dim_scaling: float = 4.0,
        **kwargs,
    ):
        r"""
        GPS layer implementation in pyg
        adapated from https://github.com/rampasek/GraphGPS/blob/main/graphgps/layer/gps_layer.py
        GPS: Recipe for a General, Powerful, Scalable Graph Transformer
        Ladislav Rampášek, Mikhail Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, Dominique Beaini
        https://arxiv.org/abs/2205.12454

        GPS++: An Optimised Hybrid MPNN/Transformer for Molecular Property Prediction
        Dominic Masters, Josef Dean, Kerstin Klaser, Zhiyi Li, Sam Maddrell-Mander, Adam Sanders, Hatem Helal, Deniz Beker, Ladislav Rampášek, Dominique Beaini
        https://arxiv.org/abs/2212.02229

        Parameters:

            in_dim:
                Input node feature dimensions of the layer

            out_dim:
                Output node feature dimensions of the layer

            in_dim:
                Input edge feature dimensions of the layer

            out_dim:
                Output edge feature dimensions of the layer

            in_dim_edges:
                input edge-feature dimensions of the layer

            out_dim_edges:
                output edge-feature dimensions of the layer

            activation:
                activation function to use in the layer

            dropout:
                The ratio of units to dropout. Must be between 0 and 1

            node_residual:
                If node residual is used after on the gnn layer output

            normalization:
                Normalization to use. Choices:

                - "none" or `None`: No normalization
                - "batch_norm": Batch normalization
                - "layer_norm": Layer normalization
                - `Callable`: Any callable function

            mpnn_type:
                type of mpnn used, choose from "pyg:gin", "pyg:gine", "pyg:gated-gcn", "pyg:pna-msgpass" and "pyg:mpnnplus"

            mpnn_kwargs:
                kwargs for mpnn layer

            attn_type:
                type of attention used, choose from "full-attention" and "none"

            attn_kwargs:
                kwargs for attention layer

            droppath_rate_attn:
                stochastic depth drop rate for attention layer https://arxiv.org/abs/1603.09382

            droppath_rate_ffn:
                stochastic depth drop rate for ffn layer https://arxiv.org/abs/1603.09382

            mpnn_type:
                Type of MPNN layer to use. Choices specified in PYG_LAYERS_DICT

            mpnn_kwargs:
                Keyword arguments to pass to the MPNN layer

            attn_type:
                Type of attention layer to use. Choices specified in ATTENTION_LAYERS_DICT

            biased_attention_key:
                indicates if biased attention is used by specifying a key corresponding to the pyg attribute in the batch (processed by the gaussian kernel encoder)
                default: None means biased attention is not used

            attn_kwargs:
                Keyword arguments to pass to the attention layer

        """

        super().__init__(
            in_dim=in_dim,
            out_dim=out_dim,
            activation=activation,
            dropout=dropout,
            normalization=normalization,
            droppath_rate=droppath_rate_attn,
            **kwargs,
        )
        # Set the other attributes
        self.in_dim_edges = in_dim_edges
        self.out_dim_edges = out_dim_edges

        # Dropout layers
        self.dropout_local = self.dropout_layer
        self.dropout_attn = self._parse_dropout(dropout=self.dropout)

        # DropPath layers
        self.droppath_ffn = self._parse_droppath(droppath_rate_ffn)

        # Residual connections
        self.node_residual = node_residual

        self.precision = precision

        # MLP applied at the end of the GPS layer
        self.mlp = MLP(
            in_dim=in_dim,
            hidden_dims=int(hidden_dim_scaling * in_dim),
            out_dim=in_dim,
            depth=2,
            activation=activation,
            dropout=self.dropout,
            last_dropout=self.dropout,
        )
        self.f_out = FCLayer(in_dim, out_dim, normalization=normalization)

        # Normalization layers
        self.norm_layer_local = self._parse_norm(normalization=self.normalization, dim=in_dim)
        self.norm_layer_attn = self._parse_norm(normalization=self.normalization, dim=in_dim)
        self.norm_layer_ff = self._parse_norm(self.normalization)

        self.biased_attention_key = biased_attention_key
        # Initialize the MPNN and Attention layers
        self.mpnn = self._parse_mpnn_layer(mpnn_type, mpnn_kwargs)
        self.attn_layer = self._parse_attn_layer(attn_type, self.biased_attention_key, attn_kwargs)

    def forward(self, batch: Batch) -> Batch:
        r"""
        forward function of the layer
        Parameters:
            batch: pyg Batch graphs to pass through the layer
        Returns:
            batch: pyg Batch graphs
        """
        # pe, feat, edge_index, edge_feat = batch.pos_enc_feats_sign_flip, batch.feat, batch.edge_index, batch.edge_feat
        feat = batch.feat

        feat_in = feat  # for first residual connection

        # Local MPNN with edge attributes.
        batch_out = self.mpnn(batch.clone())
        h_local = batch_out.feat
        if self.dropout_local is not None:
            h_local = self.dropout_local(h_local)
        if self.node_residual:
            h_local = feat_in + h_local  # Residual connection for nodes, not used in gps++.
        if self.norm_layer_local is not None:
            h_local = self.norm_layer_local(h_local)

        # Multi-head attention.
        if self.attn_layer is not None:
            h_attn = self._self_attention_block(feat, feat_in, batch)
            # Combine local and global outputs.
            feat = h_local + h_attn
        else:
            feat = h_local

        # MLP block, with skip connection
        feat_mlp = self.mlp(feat)
        # Add the droppath to the output of the MLP
        batch_size = None if feat.device.type != "ipu" else batch.graph_is_true.shape[0]
        if self.droppath_ffn is not None:
            feat_mlp = self.droppath_ffn(feat_mlp, batch.batch, batch_size)
        feat = feat + feat_mlp

        feat = self.f_out(feat)

        batch_out.feat = feat

        return batch_out

    def _parse_mpnn_layer(self, mpnn_type, mpnn_kwargs: Dict[str, Any]) -> Optional[Module]:
        """Parse the MPNN layer."""

        mpnn_kwargs = deepcopy(mpnn_kwargs)
        if mpnn_kwargs is None:
            mpnn_kwargs = {}

        # Set the default values
        mpnn_kwargs = deepcopy(mpnn_kwargs)
        mpnn_kwargs.setdefault("in_dim", self.in_dim)
        mpnn_kwargs.setdefault("out_dim", self.in_dim)
        mpnn_kwargs.setdefault("in_dim_edges", self.in_dim_edges)
        mpnn_kwargs.setdefault("out_dim_edges", self.out_dim_edges)
        # TODO: The rest of default values
        self.mpnn_kwargs = mpnn_kwargs

        # Initialize the MPNN layer
        mpnn_class = PYG_LAYERS_DICT[mpnn_type]
        mpnn_layer = mpnn_class(**mpnn_kwargs, layer_depth=self.layer_depth, layer_idx=self.layer_idx)

        return mpnn_layer

    def _parse_attn_layer(
        self, attn_type, biased_attention_key: str, attn_kwargs: Dict[str, Any]
    ) -> Optional[Module]:
        """
        parse the input attention layer and check if it is valid
        Parameters:
            attn_type: type of the attention layer
            biased_attention_key: key for the attenion bias
        Returns:
            attn_layer: the attention layer
        """

        # Set the default values for the Attention layer
        if attn_kwargs is None:
            attn_kwargs = {}
        attn_kwargs.setdefault("embed_dim", self.in_dim)
        attn_kwargs.setdefault("num_heads", 1)
        attn_kwargs.setdefault("dropout", self.dropout)
        attn_kwargs.setdefault("batch_first", True)
        self.attn_kwargs = attn_kwargs

        # Initialize the Attention layer
        attn_layer, attn_class = None, None
        if attn_type is not None:
            attn_class = ATTENTION_LAYERS_DICT[attn_type]
        if attn_class is not None:
            attn_layer = attn_class(biased_attention_key, **attn_kwargs)
        return attn_layer

    def _use_packing(self, batch: Batch) -> bool:
        """
        Check if we should use packing for the batch of graphs.
        """
        batch_keys = get_keys(batch)
        return "pack_from_node_idx" in batch_keys and "pack_attn_mask" in batch_keys

    def _to_dense_batch(
        self,
        h: Tensor,
        batch: Batch,
        batch_size: Optional[int] = None,
        max_num_nodes_per_graph: Optional[int] = None,
        on_ipu: bool = False,
    ) -> Tensor:
        """
        Convert the batch of graphs to a dense batch.
        """

        if self._use_packing(batch):
            attn_mask = batch.pack_attn_mask
            key_padding_mask = None
            idx = batch.pack_from_node_idx
            h_dense = to_packed_dense_batch(
                h,
                pack_from_node_idx=idx,
                pack_attn_mask=attn_mask,
                max_num_nodes_per_pack=100,  # TODO: This should be a parameter
            )
        else:
            attn_mask = None
            h_dense, key_padding_mask, idx = to_dense_batch(
                h,
                batch=batch.batch,  # The batch index as a vector that indicates for nodes of which graph it belongs to
                batch_size=batch_size,
                max_num_nodes_per_graph=max_num_nodes_per_graph,
                drop_nodes_last_graph=on_ipu,
            )
            key_padding_mask = ~key_padding_mask
        return h_dense, attn_mask, key_padding_mask, idx

    def _to_sparse_batch(self, batch: Batch, h_dense: Tensor, idx: Tensor) -> Tensor:
        """
        Convert the dense batch back to a sparse batch.
        """
        if self._use_packing(batch):
            h = to_sparse_batch_from_packed(
                h_dense,
                pack_from_node_idx=idx,
            )
        else:
            h = to_sparse_batch(
                h_dense,
                mask_idx=idx,
            )
        return h

    def _self_attention_block(self, feat: Tensor, feat_in: Tensor, batch: Batch) -> Tensor:
        """
        Applying the multi-head self-attention to the batch of graphs.
        First the batch is converted from [num_nodes, hidden_dim] to [num_graphs, max_num_nodes, hidden_dim]
        Then the self-attention is applied on each graph
        Then the batch is converted again to [num_nodes, hidden_dim]
        """

        # Multi-head attention.
        on_ipu = is_running_on_ipu()
        max_num_nodes_per_graph = None
        if on_ipu:
            max_num_nodes_per_graph = self.max_num_nodes_per_graph

        # Convert the tensor to a dense batch, then back to a sparse batch
        batch_size = None if feat.device.type != "ipu" else batch.graph_is_true.shape[0]

        # h[num_nodes, hidden_dim] -> h_dense[num_graphs, max_num_nodes, hidden_dim]
        feat_dense, attn_mask, key_padding_mask, idx = self._to_dense_batch(
            feat,
            batch=batch,  # The batch index as a vector that indicates for nodes of which graph it belongs to
            batch_size=batch_size,
            max_num_nodes_per_graph=max_num_nodes_per_graph,
            on_ipu=on_ipu,
        )

        attn_bias = None
        if self.biased_attention_key is not None:
            attn_bias = batch[self.biased_attention_key]

        # h_dense[num_graphs, max_num_nodes, hidden_dim] -> feat_attn[num_graphs, max_num_nodes, hidden_dim]
        feat_attn = self._sa_block(
            feat_dense, attn_bias=attn_bias, attn_mask=attn_mask, key_padding_mask=key_padding_mask
        )

        # feat_attn[num_graphs, max_num_nodes, hidden_dim] -> feat_attn[num_nodes, hidden_dim]
        feat_attn = self._to_sparse_batch(batch, feat_attn, idx)

        # Dropout, residual, norm
        if self.dropout_attn is not None:
            feat_attn = self.dropout_attn(feat_attn)
        feat_attn = feat_in + feat_attn
        if self.norm_layer_attn is not None:
            feat_attn = self.norm_layer_attn(feat_attn)
        if self.droppath_layer is not None:
            self.droppath_layer(feat_attn, batch.batch, batch_size=batch_size)

        # Combine local and global outputs.
        return feat + feat_attn

    def _sa_block(
        self, x: torch.Tensor, attn_bias: torch.Tensor, attn_mask=None, key_padding_mask=None
    ) -> torch.Tensor:
        """
        Self-attention block.
        Parameters:
            x: input tensor
            attn_bias: attention bias tensor
            attn_mask: None
            key_padding_mask: None
        Returns:
            x: output tensor
        """
        x = self.attn_layer(
            x,
            x,
            x,
            attn_bias=attn_bias,
            precision=self.precision,
            attn_mask=attn_mask,
            key_padding_mask=key_padding_mask,
            need_weights=False,
        )[0]
        return x

    @classproperty
    def layer_supports_edges(cls) -> bool:
        r"""
        Return a boolean specifying if the layer type supports edges or not.

        Returns:

            supports_edges: bool
                Always ``True`` for the current class
        """
        return True

    @property
    def layer_inputs_edges(self) -> bool:
        r"""
        Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:

            bool:
                Always ``True`` for the current class
        """
        return True

    @property
    def layer_outputs_edges(self) -> bool:
        r"""
        Abstract method. Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:

            bool:
                Always ``False`` for the current class
        """
        return self.mpnn.layer_outputs_edges

    @property
    def out_dim_factor(self) -> int:
        r"""
        Get the factor by which the output dimension is multiplied for
        the next layer.

        For standard layers, this will return ``1``.

        But for others, such as ``GatLayer``, the output is the concatenation
        of the outputs from each head, so the out_dim gets multiplied by
        the number of heads, and this function should return the number
        of heads.

        Returns:

            int:
                Always ``1`` for the current class
        """
        return 1


# graphium/nn/pyg_layers/mpnn_pyg.py

from typing import Callable, Optional, Union, Tuple, List

import torch
from torch import Tensor, IntTensor, LongTensor
from graphium.nn.base_graph_layer import BaseGraphModule
from graphium.nn.base_layers import MLP
from graphium.utils.decorators import classproperty
from torch_geometric.nn.aggr import MultiAggregation, Aggregation
from torch_geometric.data import Batch


class MPNNPlusPyg(BaseGraphModule):
    def __init__(
        self,
        in_dim: int = 64,
        out_dim: int = 64,
        activation: Union[str, Callable] = "gelu",
        dropout: float = 0.3,
        normalization: Union[str, Callable] = "layer_norm",
        gather_from: str = "both",
        scatter_to: str = "both",
        node_combine_method: str = "concat",
        num_node_mlp: int = 2,
        mlp_expansion_ratio: int = 4,
        use_edges: bool = True,
        in_dim_edges: Optional[int] = 32,
        out_dim_edges: Optional[int] = 32,
        aggregation_method: Optional[List[Union[str, Aggregation]]] = ["sum"],
        num_edge_mlp: Optional[int] = 2,
        use_globals: bool = True,
        edge_dropout_rate: Optional[float] = 0.0035,
        **kwargs,
    ):
        r"""
            MPNNPlusPyg: InteractionNetwork layer witg edges and global feature, GPS++ type of GNN layer
            GPS++: An Optimised Hybrid MPNN/Transformer for Molecular Property Prediction
            Dominic Masters, Josef Dean, Kerstin Klaser, Zhiyi Li, Sam Maddrell-Mander, Adam Sanders,
            Hatem Helal, Deniz Beker, Ladislav Rampášek, Dominique Beaini
            https://arxiv.org/abs/2212.02229

        Parameters:

            in_dim:
                Input feature dimensions of the nodes

            out_dim:
                Output feature dimensions of the nodes

            activation:
                Activation function to use in the edge and node model

            dropout:
                The ratio of units to dropout at the end within apply_norm_activation_dropout.

            normalization:
                Normalization to use with the edge, node models and at the end
                within apply_norm_activation_dropout. Choices:

                - "none" or `None`: No normalization
                - "batch_norm": Batch normalization
                - "layer_norm": Layer normalization
                - `Callable`: Any callable function

            gather_from:
                The method to gather features from. Could choose from:
                "senders", "receivers" and "both".

            scatter_to:
                The method to scatter features to. Could choose from:
                "senders", "receivers" and "both".

            node_combine_method:
                The method to combine the node features, Could choose from:
                "sum" and "concat".

            aggregation_method:
                Methods for aggregating (scatter) messages built from node and edge features.
                Provide a list of `Aggregation` or strings.
                supported strings are:

                - "sum" / "add" (Default)
                - "mean"
                - "max"
                - "min"
                - "softmax"
                - "median"
                - "std"
                - "var"

            num_node_mlp:
                Number of mlp layer used for node model

            mlp_expansion_ratio:
                Expansion ratio for node and edge mlp

            use_edges:
                If edge features are used

            in_dim_edges:
                Input feature dimensions of the edges

            out_dim_edges:
                Output feature dimensions of the edges

            num_edge_mlp:
                Number of mlp layer used for edge model

            edge_dropout_rate:
                dropout rate for the edges
        """

        super().__init__(
            in_dim=in_dim,
            out_dim=out_dim,
            activation=activation,
            dropout=dropout,
            normalization=normalization,
            **kwargs,
        )

        self.gather_from = gather_from
        self.scatter_to = scatter_to
        self.node_combine_method = node_combine_method
        self.num_node_mlp = num_node_mlp
        self.mlp_expansion_ratio = mlp_expansion_ratio

        self.use_edges = use_edges
        self.in_dim_edges = in_dim_edges
        self.out_dim_edges = out_dim_edges
        self.num_edge_mlp = num_edge_mlp
        self.edge_dropout_rate = edge_dropout_rate

        self.aggregator = MultiAggregation(aggregation_method)

        # node_model:
        edge_dim = self.out_dim_edges if use_edges else self.in_dim_edges
        if self.node_combine_method == "concat":
            node_model_in_dim = 3 * self.in_dim + 2 * edge_dim
        elif self.node_combine_method == "sum":
            node_model_in_dim = 2 * self.in_dim + edge_dim
        else:
            raise ValueError(f"node_combine_method {self.node_combine_method} not recognised.")
        node_model_hidden_dim = self.mlp_expansion_ratio * self.in_dim
        self.node_model = MLP(
            in_dim=node_model_in_dim,
            hidden_dims=node_model_hidden_dim,
            out_dim=self.out_dim,
            depth=self.num_node_mlp,
            activation=self.activation_layer,
            last_dropout=self.dropout,
            normalization=self.normalization,
        )

        # edge_model:
        if self.node_combine_method == "concat":
            edge_model_in_dim = 2 * self.in_dim + self.in_dim_edges
        elif self.node_combine_method == "sum":
            edge_model_in_dim = self.in_dim + self.in_dim_edges
        else:
            raise ValueError(f"node_combine_method {self.node_combine_method} not recognised.")
        edge_model_hidden_dim = self.mlp_expansion_ratio * self.in_dim_edges
        self.edge_model = MLP(
            in_dim=edge_model_in_dim,
            hidden_dims=edge_model_hidden_dim,
            out_dim=self.out_dim_edges,
            depth=self.num_edge_mlp,
            activation=self.activation_layer,
            last_dropout=self.edge_dropout_rate,
            normalization=self.normalization,
        )

        self.use_globals = use_globals

    def gather_features(
        self,
        input_features: Tensor,
        senders: Union[IntTensor, LongTensor],
        receivers: Union[IntTensor, LongTensor],
    ) -> Tuple[Tensor, Tensor, Tensor]:
        r"""
        Function to gather node features based on the senders and receivers of the edge indices.

        Parameters:

            input_features:
                Node features of the batch

            senders:
                Senders of the edge_index of the batch

            receivers:
                Receivers of the edge_index of the batch

        Output:
            Gathered node features (sender and receiver) summed up or concatenated
            Gathered sender features
            Gathered receiver features
        """

        out = []

        receiver_features = input_features[receivers]
        sender_features = input_features[senders]

        if self.gather_from == "receivers":
            out.append(receiver_features)

        if self.gather_from == "senders":
            out.append(sender_features)

        if self.gather_from == "both":
            if self.node_combine_method == "sum":
                out.append(receiver_features + sender_features)
            elif self.node_combine_method == "concat":
                out.append(torch.cat([receiver_features, sender_features], dim=-1))
            else:
                raise ValueError(f"node_combine_method {self.node_combine_method} not recognised.")

        return out, sender_features, receiver_features

    def aggregate_features(
        self,
        input_features: Tensor,
        senders: Union[IntTensor, LongTensor],
        receivers: Union[IntTensor, LongTensor],
        sender_features: Tensor,
        receiver_features: Tensor,
        size: int,
    ) -> Tensor:
        r"""
        Function to aggregate (scatter) messages built from node and edge features.

        Parameters:

            input_features:
                Edge features of the batch

            senders:
                Senders of the edge_index of the batch

            receivers:
                Receivers of the edge_index of the batch

            sender_features:
                Senders features gathered from the gather_features function

            receiver_features:
                Receiver features gathered from the gather_features function

            size:
                size of the aggregation, equals to the total number of nodes

        Returns:
            Tensor:
                Aggregated node features

        """

        out = []
        aggregated_features = []

        if self.scatter_to in ["receivers", "both"]:
            # using direct_neighbour_aggregation to generate the message
            message = torch.cat([input_features, sender_features], dim=-1)
            # sum method is used with aggregators
            aggregated_features.append(self.aggregator(message, receivers, dim_size=size))

        if self.scatter_to in ["senders", "both"]:
            # using direct_neighbour_aggregation to generate the message
            message = torch.cat([input_features, receiver_features], dim=-1)
            # sum method is used with aggregators
            aggregated_features.append(self.aggregator(message, senders, dim_size=size))

        if self.node_combine_method == "sum" and self.scatter_to == "both":
            out.append(aggregated_features[0] + aggregated_features[1])
        elif self.scatter_to == "both":
            out.append(torch.cat([aggregated_features[0], aggregated_features[1]], dim=-1))
        else:
            out.extend(aggregated_features)

        return out

    def forward(self, batch: Batch) -> Batch:
        r"""
        Forward function of the MPNN Plus layer
        Parameters:
            batch:
                pyg Batch graph to pass through the layer
        Returns:
            batch:
                pyg Batch graph with updated node and edge features
        """
        senders = batch.edge_index[0]
        receivers = batch.edge_index[1]
        # ---------------EDGE step---------------
        edge_model_input, sender_nodes, receiver_nodes = self.gather_features(batch.feat, senders, receivers)

        if self.use_edges:
            edge_model_input.append(batch.edge_feat)
            edge_model_input = torch.cat([edge_model_input[0], edge_model_input[1]], dim=-1)
            # edge dropout included in the edge_model
            batch.edge_feat = self.edge_model(edge_model_input)
        else:
            batch.edge_feat = edge_model_input

        # ---------------NODE step---------------
        # message + aggregate
        node_count_per_pack = batch.feat.shape[-2]
        node_model_input = self.aggregate_features(
            batch.edge_feat, senders, receivers, sender_nodes, receiver_nodes, node_count_per_pack
        )
        node_model_input.append(batch.feat)
        batch.feat = torch.cat([node_model_input[0], node_model_input[1]], dim=-1)
        batch.feat = self.node_model(batch.feat)

        # ---------------Apply norm activation and dropout---------------
        # use dropout value of the layer (default 0.3)
        batch.feat = self.apply_norm_activation_dropout(
            batch.feat,
            normalization=False,
            activation=False,
            batch_idx=batch.batch,
            batch_size=batch.num_graphs,
        )

        return batch

    @classproperty
    def layer_supports_edges(cls) -> bool:
        r"""
        Return a boolean specifying if the layer type supports edges or not.

        Returns:

            supports_edges: bool
                Always ``True`` for the current class
        """
        return True

    @property
    def layer_inputs_edges(self) -> bool:
        r"""
        Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:

            bool:
                Always ``True`` for the current class
        """
        return True

    @property
    def layer_outputs_edges(self) -> bool:
        r"""
        Abstract method. Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:

            bool:
                Always ``True`` for the current class
        """
        return True

    @property
    def out_dim_factor(self) -> int:
        r"""
        Get the factor by which the output dimension is multiplied for
        the next layer.

        For standard layers, this will return ``1``.

        But for others, such as ``GatLayer``, the output is the concatenation
        of the outputs from each head, so the out_dim gets multiplied by
        the number of heads, and this function should return the number
        of heads.

        Returns:

            int:
                Always ``1`` for the current class
        """
        return 1


# graphium/nn/pyg_layers/pna_pyg.py

from typing import Dict, List, Optional, Union, Callable
from functools import partial

import torch
from torch import Tensor
from torch_scatter import scatter

from torch_geometric.nn.conv import MessagePassing
from torch_geometric.typing import OptTensor
from torch_geometric.utils import degree
from torch_geometric.data import Data, Batch

from graphium.utils.decorators import classproperty
from graphium.nn.base_layers import MLP, FCLayer, get_activation
from graphium.nn.base_graph_layer import BaseGraphStructure, check_intpus_allow_int


class PNAMessagePassingPyg(MessagePassing, BaseGraphStructure):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        aggregators: List[str],
        scalers: List[str],
        activation: Union[Callable, str] = "relu",
        dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
        avg_d: Dict[str, float] = {"log": 1.0, "lin": 1.0},
        last_activation: Union[Callable, str] = "none",
        posttrans_layers: int = 1,
        pretrans_layers: int = 1,
        in_dim_edges: int = 0,
        **kwargs,
    ):
        r"""
        Implementation of the message passing architecture of the PNA message passing layer,
        previously known as `PNALayerComplex`. This layer applies an MLP as
        pretransformation to the concatenation of $[h_u, h_v, e_{uv}]$ to generate
        the messages, with $h_u$ the node feature, $h_v$ the neighbour node features,
        and $e_{uv}$ the edge feature between the nodes $u$ and $v$.

        After the pre-transformation, it aggregates the messages
        multiple aggregators and scalers,
        concatenates their results, then applies an MLP on the concatenated
        features.

        PNA: Principal Neighbourhood Aggregation
        Gabriele Corso, Luca Cavalleri, Dominique Beaini, Pietro Lio, Petar Velickovic
        https://arxiv.org/abs/2004.05718

        [!] code adapted from pytorch-geometric implementation of PNAConv

        Parameters:

            in_dim:
                Input feature dimensions of the layer

            out_dim:
                Output feature dimensions of the layer

            aggregators:
                Set of aggregation function identifiers,
                e.g. "mean", "max", "min", "std", "sum", "var", "moment3".
                The results from all aggregators will be concatenated.

            scalers:
                Set of scaling functions identifiers
                e.g. "identidy", "amplification", "attenuation"
                The results from all scalers will be concatenated

            activation:
                activation function to use in the layer

            dropout:
                The ratio of units to dropout. Must be between 0 and 1

            normalization:
                Normalization to use. Choices:

                - "none" or `None`: No normalization
                - "batch_norm": Batch normalization
                - "layer_norm": Layer normalization
                - `Callable`: Any callable function

            avg_d:
                Average degree of nodes in the training set, used by scalers to normalize

            last_activation:
                activation function to use in the last layer of the internal MLP

            posttrans_layers:
                number of layers in the MLP transformation after the aggregation

            pretrans_layers:
                number of layers in the transformation before the aggregation

            in_dim_edges:
                size of the edge features. If 0, edges are ignored

        """

        MessagePassing.__init__(self, node_dim=0)
        BaseGraphStructure.__init__(
            self,
            in_dim=in_dim,
            out_dim=out_dim,
            activation=activation,
            dropout=dropout,
            normalization=normalization,
            **kwargs,
        )

        # Allow int32 as edge index
        self.__check_input__ = partial(check_intpus_allow_int, self)

        self.aggregators = aggregators
        self.scalers = scalers

        # Edge dimensions
        self.in_dim_edges = in_dim_edges
        self.edge_encoder = None
        if self.in_dim_edges > 0:
            self.edge_encoder = FCLayer(self.in_dim_edges, self.in_dim_edges, activation=None)

        # Initializing basic attributes
        self.avg_d = avg_d
        self.last_activation = get_activation(last_activation)

        # MLP used on each pair of nodes with their edge MLP(h_u, h_v, e_uv)
        self.pretrans = MLP(
            in_dim=2 * in_dim + in_dim_edges,
            hidden_dims=in_dim,
            out_dim=in_dim,
            depth=pretrans_layers,
            activation=self.activation,
            last_activation=self.last_activation,
            dropout=dropout,
            normalization=normalization,
            last_normalization=normalization,
        )

        # MLP used on the aggregated messages of the neighbours
        self.posttrans = MLP(
            in_dim=(len(aggregators) * len(scalers)) * self.in_dim,
            hidden_dims=self.out_dim,
            out_dim=self.out_dim,
            depth=posttrans_layers,
            activation=self.activation,
            last_activation=self.last_activation,
            dropout=dropout,
            normalization=normalization,
            last_normalization=normalization,
        )

    def forward(self, batch: Union[Data, Batch]) -> Union[Data, Batch]:
        r"""
        forward function of the layer
        Parameters:
            batch: pyg Batch graphs
        Returns:
            batch: pyg Batch graphs
        """
        feat, edge_index, edge_feat = batch.feat, batch.edge_index, batch.edge_feat

        out = self.propagate(edge_index, x=feat, edge_feat=edge_feat, size=None)
        out = self.posttrans(out)  # No more towers and concat with x
        batch.feat = out
        return batch

    def message(self, x_i: Tensor, x_j: Tensor, edge_feat: OptTensor) -> Tensor:
        r"""
        message function

        Parameters:
            x_i: node features
            x_j: neighbour node features
            edge_feat: edge features
        Returns:
            feat: the message
        """
        feat: Tensor = x_i  # Dummy.
        if (edge_feat is not None) and (self.edge_encoder is not None):
            edge_feat = self.edge_encoder(edge_feat)
            feat = torch.cat([x_i, x_j, edge_feat], dim=-1)
        else:
            feat = torch.cat([x_i, x_j], dim=-1)

        return self.pretrans(feat)  # No more towers

    def aggregate(
        self,
        inputs: Tensor,
        index: Tensor,
        edge_index: Tensor,
        dim_size: Optional[int] = None,
    ) -> Tensor:
        r"""
        aggregate function

        Parameters:
            inputs: input features
            index: index of the nodes
            edge_index: edge index
            dim_size: dimension size
        Returns:
            out: aggregated features
        """
        outs = []

        for aggregator in self.aggregators:
            if aggregator == "sum":
                out = scatter(inputs, index, 0, None, dim_size, reduce="sum")
            elif aggregator == "mean":
                out = scatter(inputs, index, 0, None, dim_size, reduce="mean")
            elif aggregator == "min":
                out = scatter(inputs, index, 0, None, dim_size, reduce="min")
            elif aggregator == "max":
                out = scatter(inputs, index, 0, None, dim_size, reduce="max")
            elif aggregator in ["var", "std"]:
                mean = scatter(inputs, index, 0, None, dim_size, reduce="mean")
                mean_squares = scatter(inputs * inputs, index, 0, None, dim_size, reduce="mean")
                out = mean_squares - mean * mean
                if aggregator == "std":
                    out = torch.sqrt(torch.relu(out) + 1e-5)
            else:
                raise ValueError(f'Unknown aggregator "{aggregator}".')
            outs.append(out)
        out = torch.cat(outs, dim=-1)

        deg = degree(index, dim_size, dtype=inputs.dtype)
        deg = deg.clamp_(1).view(-1, 1)

        outs = []
        for scaler in self.scalers:
            if scaler == "identity":
                out_scaler = out
            elif scaler == "amplification":
                out_scaler = out * (torch.log(deg + 1) / self.avg_d["log"])
            elif scaler == "attenuation":
                out_scaler = out * (self.avg_d["log"] / torch.log(deg + 1))
            elif scaler == "linear":
                out_scaler = out * (deg / self.avg_d["lin"])
            elif scaler == "inverse_linear":
                out_scaler = out * (self.avg_d["lin"] / deg)
            else:
                raise ValueError(f'Unknown scaler "{scaler}".')
            outs.append(out_scaler)
        return torch.cat(outs, dim=-1)

    @property
    def layer_outputs_edges(self) -> bool:
        r"""
        Abstract method. Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:

            bool:
                Always ``False`` for the current class
        """
        return False

    @property
    def out_dim_factor(self) -> int:
        r"""
        Get the factor by which the output dimension is multiplied for
        the next layer.

        For standard layers, this will return ``1``.

        But for others, such as ``GatLayer``, the output is the concatenation
        of the outputs from each head, so the out_dim gets multiplied by
        the number of heads, and this function should return the number
        of heads.

        Returns:

            int:
                Always ``1`` for the current class
        """
        return 1

    @property
    def layer_inputs_edges(self) -> bool:
        r"""
        Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:

            bool:
                Returns ``self.edge_features``
        """
        return self.edge_features

    @classproperty
    def layer_supports_edges(cls) -> bool:
        r"""
        Return a boolean specifying if the layer type supports edges or not.

        Returns:

            bool:
                Always ``True`` for the current class
        """
        return True


# graphium/nn/pyg_layers/pooling_pyg.py

import torch
import torch.nn as nn
from torch import Tensor, LongTensor
from typing import List, Union, Callable, Tuple, Optional, Dict
from copy import deepcopy

from torch_scatter import scatter
from torch_geometric.data import Data, Batch

from graphium.nn.base_layers import MLP, FCLayer
from graphium.utils.tensor import ModuleListConcat, ModuleWrap
from graphium.nn.base_layers import MuReadoutGraphium

EPS = 1e-6


def scatter_logsum_pool(x: Tensor, batch: LongTensor, dim: int = 0, dim_size: Optional[int] = None) -> Tensor:
    r"""
    Apply pooling over the nodes in the graph using a mean aggregation,
    but scaled by the log of the number of nodes. This gives the same
    expressive power as the sum, but helps deal with graphs that are
    significantly larger than others by using a logarithmic scale.

    $$r^{(i)} = \frac{\log N_i}{N_i}\sum_{k=1}^{N_i} x^{(i)}_k$$

    Parameters:
        x (Tensor): Node feature matrix
            :math:`\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}`.
        batch (LongTensor): Batch vector :math:`\mathbf{b} \in {\{ 0, \ldots,
            B-1\}}^N`, which assigns each node to a specific example.
        size (int, optional): Batch-size :math:`B`.
            Automatically calculated if not given. (default: :obj:`None`)
    Returns:
        the pooled features tensor
    """
    dim_size = int(batch.max().detach() + 1) if dim_size is None else dim_size
    mean_pool = scatter(x, batch, dim=dim, dim_size=dim_size, reduce="mean")
    num_nodes = scatter(
        torch.ones(x.shape[:-1], dtype=x.dtype, device=x.device),
        batch,
        dim=dim,
        dim_size=dim_size,
        reduce="sum",
    )
    lognum = torch.log(num_nodes)
    return mean_pool * lognum.unsqueeze(-1)


def scatter_std_pool(x: Tensor, batch: LongTensor, dim: int = 0, dim_size: Optional[int] = None):
    r"""Returns batch-wise graph-level-outputs by taking the channel-wise
    minimum across the node dimension, so that for a single graph
    :math:`\mathcal{G}_i` its output is computed by

    .. math::
        \mathbf{r}_i = \mathrm{max}_{n=1}^{N_i} \, \mathbf{x}_n

    Parameters:
        x (Tensor): Node feature matrix
            :math:`\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}`.
        batch (LongTensor): Batch vector :math:`\mathbf{b} \in {\{ 0, \ldots,
            B-1\}}^N`, which assigns each node to a specific example.
        size (int, optional): Batch-size :math:`B`.
            Automatically calculated if not given. (default: :obj:`None`)

    Returns:
        the pooled features tensor
    """
    dim_size = int(batch.max().detach() + 1) if dim_size is None else dim_size
    mean = scatter(x, batch, dim=dim, out=None, dim_size=dim_size, reduce="mean")
    mean_squares = scatter(x * x, batch, dim=dim, out=None, dim_size=dim_size, reduce="mean")
    out = mean_squares - mean * mean
    return torch.sqrt(torch.relu(out) + 1e-5)


class PoolingWrapperPyg(ModuleWrap):
    def __init__(self, func, feat_type, *args, **kwargs) -> None:
        super().__init__(func, *args, **kwargs)
        self.feat_type = feat_type

    def forward(self, g: Batch, feature: Tensor, *args, **kwargs):
        """
        forward function
        Parameters:
            g: the pyg batch graph
            feature: the node features
        Returns:
            the pooled features
        """
        dim_size = g.num_graphs
        if self.feat_type == "node":
            index = g.batch
        elif self.feat_type == "edge":
            index = g.batch[g.edge_index][0]
        else:
            index = g.batch
        return self.func(feature, index, dim_size=dim_size, *args, **kwargs, **self.kwargs)


def parse_pooling_layer_pyg(in_dim: int, pooling: Union[str, List[str]], feat_type: str = "node", **kwargs):
    r"""
    Select the pooling layers from a list of strings, and put them
    in a Module that concatenates their outputs.

    Parameters:

        in_dim:
            The dimension at the input layer of the pooling

        pooling:
            The list of pooling layers to use. The accepted strings are:

            - "none": No pooling
            - "sum": Sum all the nodes for each graph
            - "mean": Mean all the nodes for each graph
            - "logsum": Mean all the nodes then multiply by log(num_nodes) for each graph
            - "max": Max all the nodes for each graph
            - "min": Min all the nodes for each graph
            - "std": Standard deviation of all the nodes for each graph

    """

    # Create the pooling layer
    pool_layer = ModuleListConcat()
    out_pool_dim = 0
    if isinstance(pooling, str):
        pooling = [pooling]
    assert feat_type in ["node", "edge", "global"]
    for this_pool in pooling:
        this_pool = None if this_pool is None else this_pool.lower()
        out_pool_dim += in_dim
        if this_pool == "sum":
            pool_layer.append(PoolingWrapperPyg(scatter, dim=0, reduce="add", feat_type=feat_type, **kwargs))
        elif this_pool == "mean":
            pool_layer.append(PoolingWrapperPyg(scatter, dim=0, reduce="mean", feat_type=feat_type, **kwargs))
        elif this_pool == "logsum":
            pool_layer.append(PoolingWrapperPyg(scatter_logsum_pool, dim=0, feat_type=feat_type, **kwargs))
        elif this_pool == "max":
            pool_layer.append(PoolingWrapperPyg(scatter, dim=0, reduce="max", feat_type=feat_type, **kwargs))
        elif this_pool == "min":
            pool_layer.append(PoolingWrapperPyg(scatter, dim=0, reduce="min", feat_type=feat_type, **kwargs))
        elif this_pool == "std":
            pool_layer.append(PoolingWrapperPyg(scatter_std_pool, dim=0, feat_type=feat_type, **kwargs))
        elif (this_pool == "none") or (this_pool is None):
            pass
        else:
            raise NotImplementedError(f"Undefined pooling `{this_pool}`")

    return pool_layer, out_pool_dim


class VirtualNodePyg(nn.Module):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        in_dim_edges: Optional[int],
        out_dim_edges: Optional[int],
        vn_type: Union[type(None), str] = "sum",
        activation: Union[str, Callable] = "relu",
        dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
        bias: bool = True,
        residual: bool = True,
        use_edges: bool = False,
        **kwargs,
    ):
        r"""
        The VirtualNode is a layer that pool the features of the graph,
        applies a neural network layer on the pooled features,
        then add the result back to the node features of every node.

        Parameters:

            in_dim:
                Input feature dimensions of the virtual node layer.

            out_dim:
                Output feature dimensions of the virtual node layer.

            in_dim_edges:
                Input feature dimensions of the virtual node layer for the edges.

            out_dim_edges:
                Output feature dimensions of the virtual node layer for the edges.

            vn_type:
                The type of the virtual node. Choices are:

                - "none": No pooling
                - "sum": Sum all the nodes for each graph
                - "mean": Mean all the nodes for each graph
                - "logsum": Mean all the nodes then multiply by log(num_nodes) for each graph
                - "max": Max all the nodes for each graph
                - "min": Min all the nodes for each graph
                - "std": Standard deviation of all the nodes for each graph

            activation:
                activation function to use in the neural network layer.

            dropout:
                The ratio of units to dropout. Must be between 0 and 1

            normalization:
                Normalization to use. Choices:

                - "none" or `None`: No normalization
                - "batch_norm": Batch normalization
                - "layer_norm": Layer normalization
                - `Callable`: Any callable function

            bias:
                Whether to add a bias to the neural network

            residual:
                Whether all virtual nodes should be connected together
                via a residual connection

            use_edges:
                Boolean flag to select if edges are used in the global node
                aggregation and update of features

        """
        super().__init__()
        if (vn_type is None) or (vn_type.lower() == "none"):
            self.vn_type = None
            self.fc_layer = None
            self.residual = None
            return
        # Layer sizes stored here
        self.in_dim_edges = in_dim_edges
        self.out_dim_edges = out_dim_edges
        self.in_dim_nodes = in_dim
        self.out_dim_nodes = out_dim

        # Use edge features in pooling
        self.use_edges = use_edges
        has_edges = (
            (in_dim_edges is not None)
            and (out_dim_edges is not None)
            and (in_dim_edges > 0)
            and (out_dim_edges > 0)
        )
        if self.use_edges and not has_edges:
            raise ValueError("The edge features are not defined but `use_edges` is True")
        self.vn_type = vn_type.lower()
        self.layer, out_pool_dim = parse_pooling_layer_pyg(
            in_dim=self.in_dim_nodes, pooling=self.vn_type, feat_type="node"
        )
        self.residual = residual

        if self.use_edges:
            self.edge_layer, out_edge_pool_dim = parse_pooling_layer_pyg(
                in_dim=self.in_dim_edges, pooling=self.vn_type, feat_type="edge"
            )

            out_pool_dim = out_pool_dim + out_edge_pool_dim

        self.fc_layer = FCLayer(
            in_dim=out_pool_dim,
            out_dim=out_pool_dim,
            activation=activation,
            dropout=dropout,
            normalization=normalization,
            bias=bias,
        )

        # Projection layers from the pooling layer to node and edge feature sizes
        self.node_projection = MuReadoutGraphium(out_pool_dim, self.out_dim_nodes)
        self.edge_projection = None
        if self.use_edges:
            self.edge_projection = MuReadoutGraphium(out_pool_dim, self.out_dim_edges)

    def forward(
        self, g: Union[Data, Batch], feat: Tensor, vn_feat: LongTensor, edge_feat: Tensor
    ) -> Tuple[Tensor, Tensor, Tensor]:
        r"""
        Apply the virtual node layer.

        Parameters:

            g:
                PyG Graphs or Batched graphs.

            feat (torch.Tensor[..., N, Din]):
                Node feature tensor, before convolution.
                `N` is the number of nodes, `Din` is the input features

            vn_feat (torch.Tensor[..., M, Din]):
                Graph feature of the previous virtual node, or `None`
                `M` is the number of graphs, `Din` is the input features.
                It is added to the result after the MLP, as a residual connection

            edge_feat (torch.Tensor[..., E, Din]):
                Edge feature tensor, before convolution.
                `E` is the number of edges, `Din` is the input features

        Returns:

            `feat = torch.Tensor[..., N, Dout]`:
                Node feature tensor, after convolution and residual.
                `N` is the number of nodes, `Dout` is the output features of the layer and residual

            `vn_feat = torch.Tensor[..., M, Dout]`:
                Graph feature tensor to be used at the next virtual node, or `None`
                `M` is the number of graphs, `Dout` is the output features

            `edge_feat = torch.Tensor[..., N, Dout]`:
                Edge feature tensor, after convolution and residual - if edges are used, otherwise returned unchanged.
                `N` is the number of edges, `Dout` is the output features of the layer and residual

        """

        # Pool the features
        if self.vn_type is None:
            return feat, vn_feat, edge_feat
        pool = self.layer(g, feat)
        if self.use_edges:
            edge_pool = self.edge_layer(g, edge_feat)
            pool = torch.cat((pool, edge_pool), -1)

        vn_h_temp = self.fc_layer.forward(vn_feat + pool)

        if self.residual:
            vn_feat = vn_feat + vn_h_temp
        else:
            vn_feat = vn_h_temp

        # Add the virtual node projections to the node features
        feat = feat + self.node_projection(vn_feat[g.batch])

        if self.use_edges:
            # Add the virtual node projections to the edge features
            edge_feat = edge_feat + self.edge_projection(vn_feat[g.batch[g.edge_index][0]])
        return feat, vn_feat, edge_feat


# graphium/nn/pyg_layers/utils.py

import math
import torch
import torch.nn as nn
from torch_geometric.data import Batch
from typing import Tuple
from torch import Tensor

from torch_geometric.typing import SparseTensor

from graphium.nn.base_layers import MLP, get_norm
from graphium.ipu.to_dense_batch import to_dense_batch, to_sparse_batch


class PreprocessPositions(nn.Module):
    """
    Compute 3D attention bias and 3D node features according to the 3D position information.
    """

    def __init__(
        self,
        num_heads,
        embed_dim,
        num_kernel,
        in_dim=3,
        num_layers=2,
        activation="gelu",
        first_normalization="none",
    ):
        r"""
        Parameters:
            num_heads:
                Number of attention heads used in self-attention.
            embed_dim:
                Hidden dimension of node features.
            num_kernel:
                Number of gaussian kernels.
            num_layers: The number of layers in the MLP.
            activation: The activation function used in the MLP.
            first_normalization: The normalization function used before the gaussian kernel.

        """
        super().__init__()
        self.num_heads = num_heads
        self.num_kernel = num_kernel
        self.embed_dim = embed_dim
        self.first_normalization = get_norm(first_normalization, dim=in_dim)

        self.gaussian = GaussianLayer(self.num_kernel, in_dim=in_dim)
        self.gaussian_proj = MLP(
            in_dim=self.num_kernel,
            hidden_dims=self.num_kernel,
            out_dim=self.num_heads,
            depth=num_layers,
            activation=activation,
            last_layer_is_readout=True,  # Since the output is not proportional to the hidden dim, but to the number of heads
        )

        # make sure the 3D node feature has the same dimension as the embedding size
        # so that it can be added to the original node features
        self.node_proj = nn.Linear(self.num_kernel, self.embed_dim)

    def forward(
        self, batch: Batch, max_num_nodes_per_graph: int, on_ipu: bool, positions_3d_key: str
    ) -> Tuple[Tensor, Tensor]:
        r"""
        Inputs:
            batch:
                Batch object.
            max_num_nodes_per_graph:
                Maximum number of nodes per graph.
            on_ipu:
                If model rus on IPU.
            positions_3d_key:
                The key of the pyg graph object that contains the 3D positions.

        """

        pos = batch[positions_3d_key]
        if self.first_normalization is not None:
            pos = self.first_normalization(pos)
        batch_size = None if pos.device.type != "ipu" else batch.graph_is_true.shape[0]
        # batch_size = None if batch.feat.device.type != "ipu" else batch.graph_is_true.shape[0] #[Andy] batch.feat is only available after passing through layers, not a good attribute to check
        # pos: [batch, nodes, 3]
        # padding_mask: [batch, nodes]
        # idx: [totoal_nodes]
        pos, mask, idx = to_dense_batch(
            pos,
            batch=batch.batch,
            batch_size=batch_size,
            max_num_nodes_per_graph=max_num_nodes_per_graph,
            drop_nodes_last_graph=on_ipu,
        )
        # check nan with the pos from to_dense_batch,
        # and generate mask. 1 for nan, 0 for other values.
        # pos consists of real nodes and padding nodes
        # for real nodes, if 3d position does not exit, it is nans. For padding nodes, 3d positions will be 0
        # if the first node of a molecule has 3d position as nan, the whole molecule will be masked out.
        # [batch]
        nan_mask = torch.isnan(pos)[:, 0, 0]
        # apply nan_mask on pos so that it does not give nan gradient
        # when applying gaussian kernels
        pos.masked_fill_(nan_mask.unsqueeze(1).unsqueeze(2), 0.0)
        # we need the opposite of mask output
        padding_mask = ~mask
        # [batch, nodes]
        batch, n_node, _ = pos.shape
        # [batch, nodes, nodes, 3]
        delta_pos = pos.unsqueeze(1) - pos.unsqueeze(2)
        # [batch, nodes, nodes]
        distance = delta_pos.norm(dim=-1).view(-1, n_node, n_node)
        # [batch, nodes, nodes, num_kernel]
        distance_feature = self.gaussian(distance)
        # [batch, nodes, nodes, num_heads]
        attn_bias = self.gaussian_proj(distance_feature)
        # [batch, num_heads, nodes, nodes]
        attn_bias = attn_bias.permute(0, 3, 1, 2).contiguous()
        # apply padding_mask on attn_bias
        # unsqueezed mask size: [batch, 1, 1, nodes] apply on tensor [batch, num_heads, nodes, nodes]
        attn_bias.masked_fill_(
            padding_mask.unsqueeze(1).unsqueeze(2),
            float("-1000"),
        )
        # apply nan_mask on attn_bias
        # unsqueezed mask size: [batch, 1, 1, 1] apply on tensor [batch, num_heads, nodes, nodes]
        attn_bias.masked_fill_(
            nan_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1),
            0.0,
        )
        # apply padding_mask on distance_feature
        # unsqueezed mask size: [batch, 1, nodes, 1] apply on tensor [batch, nodes, nodes, num_kernel]
        distance_feature.masked_fill_(padding_mask.unsqueeze(1).unsqueeze(-1).to(torch.bool), 0.0)
        # [batch, nodes, num_kernel]
        distance_feature_sum = distance_feature.sum(dim=-2)
        # Output of GaussianLayer is FP32, cast to dtype of self.node_proj here
        distance_feature_sum = distance_feature_sum.to(self.node_proj.weight.dtype)
        # [batch, nodes, embed_dim]
        node_feature = self.node_proj(distance_feature_sum)
        # apply nan_mask on node_feature
        # unsqueezed mask size: [batch, 1, 1] apply on tensor [batch, nodes, embed_dim]
        node_feature.masked_fill_(nan_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), 0.0)
        # [total_nodes, embed_dim]
        node_feature = to_sparse_batch(node_feature, idx)

        return attn_bias, node_feature


class GaussianLayer(nn.Module):
    def __init__(self, num_kernels=128, in_dim=3):
        r"""
            Gaussian kernel function that applied on the all-to-all 3D distances.
        Parameters:
            num_kernels:
                Number of gaussian kernel used.
        """
        super().__init__()
        self.num_kernels = num_kernels
        self.means = nn.Embedding(1, num_kernels)
        self.stds = nn.Embedding(1, num_kernels)
        nn.init.uniform_(self.means.weight, 0, in_dim)
        nn.init.uniform_(self.stds.weight, 0, in_dim)

    def forward(self, input: Tensor) -> Tensor:
        # [batch, nodes, nodes, 1]
        input = input.unsqueeze(-1)
        # [batch, nodes, nodes, num_kernels]
        expanded_input = input.expand(-1, -1, -1, self.num_kernels)
        # [num_kernels]
        mean = self.means.weight.float().view(-1)
        # [num_kernels]
        std = self.stds.weight.float().view(-1).abs() + 0.01  # epsilon is 0.01 that matches gps++ value
        pre_exp_factor = (2 * math.pi) ** 0.5
        # [batch, nodes, nodes, num_kernels]
        tensor_with_kernel = torch.exp(-0.5 * (((expanded_input - mean) / std) ** 2)) / (pre_exp_factor * std)
        return tensor_with_kernel


def triplets(
    edge_index: Tensor,
    num_nodes: int,
) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
    r"""Generates triplets from the given edge indices.
        A triplet is defined as a path of length two,
        such that if node A is connected to node B,
        and node B is connected to node C, then there is a triplet (A, B, C).

    Parameters:
        edge_index (LongTensor): The edge indices.
        num_nodes (int): The number of nodes.

    Returns:
        col: The sink node indices of edges from the edge indices.
        row: The source node indices of edges from the edge indices.
        idx_i: The sink node indices of the triplets.
        idx_j: The middle node indices of the triplets.
        idx_k: The source node indices of the triplets.
        idx_kj: The indices of edges those from the source node to the middle node.
        idx_ji: The indices of edges those from the middle node to the sink node.
    """
    row, col = edge_index  # j->i

    value = torch.arange(row.size(0), device=row.device)
    adj_t = SparseTensor(row=col, col=row, value=value, sparse_sizes=(num_nodes, num_nodes))
    adj_t_row = adj_t[row]
    num_triplets = adj_t_row.set_value(None).sum(dim=1).to(torch.long)

    # Node indices (k->j->i) for triplets.
    idx_i = col.repeat_interleave(num_triplets)
    idx_j = row.repeat_interleave(num_triplets)
    idx_k = adj_t_row.storage.col()

    # Remove self-loop triplets d->b->d
    mask = idx_i != idx_k  # Remove i == k triplets.
    idx_i, idx_j, idx_k = idx_i[mask], idx_j[mask], idx_k[mask]

    # Edge indices (k->j, j->i) for triplets.
    idx_kj = adj_t_row.storage.value()[mask]
    idx_ji = adj_t_row.storage.row()[mask]
    return col, row, idx_i, idx_j, idx_k, idx_kj, idx_ji


# expts/dataset_benchmark.py

import os
from os.path import dirname, abspath
import yaml
from omegaconf import DictConfig
from datetime import datetime
from copy import deepcopy

import graphium
from graphium.config._loader import load_datamodule, load_accelerator
import time
from typing import Optional, List, Sequence
import wandb
import statistics
import tqdm
import torch
import numpy as np

# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(graphium.__file__)))
os.chdir(MAIN_DIR)
# CONFIG_FILE = "expts/neurips2023_configs/debug/config_large_gcn_debug.yaml"
CONFIG_FILE = "expts/neurips2023_configs/config_large_gcn.yaml"
# CONFIG_FILE = "expts/configs/config_pcqmv2_mpnn.yaml"
# CONFIG_FILE = "expts/configs/config_ipu_qm9.yaml"


def benchmark(fn, *args, message="", log2wandb=False, **kwargs):
    start = time.time()
    value = fn(*args, **kwargs)
    duration = time.time() - start
    print(f"{message} {duration:.3f} secs")
    if log2wandb:
        wandb.log({message: duration})
    return value


def benchmark_dataloader(dataloader, name, n_epochs=5, log2wandb=False):
    print(f"length of {name} dataloader: {len(dataloader)}")
    epoch_times = [0] * n_epochs
    tputs = [0] * n_epochs
    n_batches = [0] * n_epochs
    n_graphs = [0] * n_epochs
    n_nodes = [0] * n_epochs
    n_edges = [0] * n_epochs
    for i in range(n_epochs):
        start = time.time()
        for data in tqdm.tqdm(dataloader):
            n_batches[i] += 1
            n_graphs[i] += torch.sum(torch.max(data["features"]["batch"], dim=-1).values).item()
            n_nodes[i] += np.prod(data["features"]["batch"].shape)
            n_edges[i] += np.prod(data["features"]["edge_weight"].shape)
        epoch_times[i] = time.time() - start
        tputs[i] = n_graphs[i] / epoch_times[i]

    average_tput = statistics.mean(tputs)
    print(f"{name} dataloader average tput {average_tput}")
    print(f"{name} dataloader tputs per epoch {tputs}")
    print(f"{name} dataloader epoch times {epoch_times}")
    print(f"{name} dataloader total batches per epoch {n_batches}")
    print(f"{name} dataloader total graphs per epoch {n_graphs}")
    print(f"{name} dataloader total nodes per epoch {n_nodes}")
    print(f"{name} dataloader total edges per epoch {n_edges}")

    if log2wandb:
        wandb.log({"average tput": average_tput})

        for i in range(n_epochs):
            d = {
                "epoch": i,
                "tput per epoch": tputs[i],
                "epoch times ": epoch_times[i],
                "batches per epoch": n_batches[i],
                "graphs per epoch": n_graphs[i],
                "nodes per epoch": n_nodes[i],
                "edges per epoch": n_edges[i],
            }
            print(d)
            wandb.log(d)


def main(
    cfg: DictConfig,
    stages: Optional[Sequence[str]] = None,
    run_name: str = "dataset_benchmark",
    add_date_time: bool = True,
    log2wandb: bool = False,
) -> None:
    if add_date_time:
        run_name += "_" + datetime.now().strftime("%d.%m.%Y_%H.%M.%S")

    if log2wandb:
        wandb.init(project="multitask-gnn", name=run_name, config=cfg)

    cfg = deepcopy(cfg)
    cfg, accelerator_type = load_accelerator(cfg)
    # Load and initialize the dataset
    datamodule = benchmark(
        load_datamodule, cfg, accelerator_type, message="Load duration", log2wandb=log2wandb
    )

    benchmark(datamodule.prepare_data, message="Prepare duration", log2wandb=log2wandb)

    if False:  # stages is not None:
        for stage in stages:
            benchmark(datamodule.setup, stage, message=f"Setup {stage} duration", log2wandb=log2wandb)
    else:
        benchmark(datamodule.setup, message=f"Setup duration", log2wandb=log2wandb)

    if stages is None or {"train", "fit"}.intersection(stages):
        dataloader = datamodule.train_dataloader()
        benchmark_dataloader(
            dataloader, name="train", n_epochs=cfg["trainer"]["trainer"]["max_epochs"], log2wandb=log2wandb
        )
    if stages is None or {"val", "valid", "validation"}.intersection(stages):
        dataloader = datamodule.val_dataloader()
        benchmark_dataloader(
            dataloader,
            name="validation",
            n_epochs=cfg["trainer"]["trainer"]["max_epochs"],
            log2wandb=log2wandb,
        )
    if stages is None or {"test", "testing"}.intersection(stages):
        dataloader = datamodule.test_dataloader()
        benchmark_dataloader(
            dataloader, name="testing", n_epochs=cfg["trainer"]["trainer"]["max_epochs"], log2wandb=log2wandb
        )


if __name__ == "__main__":
    with open(os.path.join(MAIN_DIR, CONFIG_FILE), "r") as f:
        cfg = yaml.safe_load(f)
    main(cfg, stages=["train"], log2wandb=True)


# expts/debug_yaml.py

# test_yaml.py


import yaml

CONFIG_FILE = "neurips2023_configs/config_small_mpnn.yaml"

import re
from collections import defaultdict

# def get_anchors_and_aliases(filepath):
#     anchors = defaultdict(list)
#     current_level = {}

#     with open(filepath, "r") as file:
#         for line in file:
#             indent = len(line) - len(line.lstrip(' '))
#             key_match = re.search(r'(\w+):', line)
#             anchor_match = re.search(r'&(\w+)', line)
#             alias_match = re.search(r'\*(\w+)', line)
#             if key_match:
#                 key = key_match.group(1)
#                 # Compute the full path of the current key.
#                 full_path = '.'.join([current_level[i] for i in sorted(current_level.keys()) if i < indent] + [key])
#                 current_level[indent] = key
#                 # Remove any keys that are indented more than the current line.
#                 keys_to_remove = [i for i in current_level if i > indent]
#                 for i in keys_to_remove:
#                     del current_level[i]
#             else:
#                 full_path = '.'.join([current_level[i] for i in sorted(current_level.keys())])
#             if anchor_match:
#                 anchor = anchor_match.group(1)
#                 anchors[anchor].append(full_path)
#             if alias_match:
#                 alias = alias_match.group(1)
#                 anchors[alias].append(full_path)
#     return {k: v for k, v in anchors.items() if len(v) > 1}


def get_anchors_and_aliases(filepath):
    anchors = defaultdict(list)
    current_level = {}
    anchor_to_path = {}

    with open(filepath, "r") as file:
        for line in file:
            indent = len(line) - len(line.lstrip(" "))
            key_match = re.search(r"(\w+):", line)
            anchor_match = re.search(r"&(\w+)", line)
            alias_match = re.search(r"\*(\w+)", line)
            if key_match:
                key = key_match.group(1)
                # Compute the full path of the current key.
                full_path = ".".join(
                    [current_level[i] for i in sorted(current_level.keys()) if i < indent] + [key]
                )
                current_level[indent] = key
                # Remove any keys that are indented more than the current line.
                keys_to_remove = [i for i in current_level if i > indent]
                for i in keys_to_remove:
                    del current_level[i]
            else:
                full_path = ".".join([current_level[i] for i in sorted(current_level.keys())])
            if anchor_match:
                anchor = anchor_match.group(1)
                anchor_to_path[anchor] = full_path
            if alias_match:
                alias = alias_match.group(1)
                if alias in anchor_to_path:
                    anchors[anchor_to_path[alias]].append(full_path)
    return anchors


refs = get_anchors_and_aliases(CONFIG_FILE)
import ipdb

ipdb.set_trace()
pass


# expts/main_run_finetuning.py

# General imports
import os
from os.path import dirname, abspath
from omegaconf import DictConfig, OmegaConf
import timeit
from loguru import logger
from datetime import datetime
from lightning.pytorch.utilities.model_summary import ModelSummary

from graphium.utils.mup import set_base_shapes
from graphium.finetuning import modify_cfg_for_finetuning, MolecularFinetuning

# Current project imports
import graphium
from graphium.config._loader import (
    load_datamodule,
    load_metrics,
    load_architecture,
    load_predictor,
    load_trainer,
    save_params_to_wandb,
    load_accelerator,
)
from graphium.utils.safe_run import SafeRun

import hydra

# WandB
import wandb

# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(graphium.__file__)))
os.chdir(MAIN_DIR)


@hydra.main(version_base=None, config_path="hydra-configs", config_name="finetune")
def main(cfg: DictConfig) -> None:
    cfg = OmegaConf.to_container(cfg, resolve=True)

    cfg = modify_cfg_for_finetuning(cfg)

    run_name: str = "main"
    add_date_time: bool = True

    st = timeit.default_timer()

    date_time_suffix = ""
    if add_date_time:
        date_time_suffix = datetime.now().strftime("%d.%m.%Y_%H.%M.%S")

    wandb.init(entity=cfg["constants"]["entity"], project=cfg["constants"]["name"], config=cfg)

    # Initialize the accelerator
    cfg, accelerator_type = load_accelerator(cfg)

    # Load and initialize the dataset
    datamodule = load_datamodule(cfg, accelerator_type)

    # Initialize the network
    model_class, model_kwargs = load_architecture(
        cfg,
        in_dims=datamodule.in_dims,
    )

    datamodule.prepare_data()

    metrics = load_metrics(cfg)
    logger.info(metrics)

    predictor = load_predictor(
        cfg, model_class, model_kwargs, metrics, accelerator_type, datamodule.task_norms
    )

    ########################
    # Changes for finetuning

    # Load pretrained & replace in predictor
    pretrained_model = predictor.load_pretrained_models('dummy-pretrained-model').model     # make pretrained model part of config
    
    # Adapt pretrained model to new task
    # We need to overwrite shared weights with pretrained

    predictor.model.overwrite_with_pretrained(cfg, pretrained_model)

    # (Un)freezing will be handled by finetuning callback added to trainer

    predictor.model = set_base_shapes(predictor.model, base=None)       # how do we deal with muP for finetuning; it
    ########################

    logger.info(predictor.model)
    logger.info(ModelSummary(predictor, max_depth=4))

    trainer = load_trainer(cfg, run_name, accelerator_type, date_time_suffix)
    
    ########################
    # Changes for finetuning

    # Add the pl.BaseFinetuning callback to trainer
    
    trainer.callbacks.append(
        MolecularFinetuning(cfg)
    )
    ########################

    save_params_to_wandb(trainer.logger, cfg, predictor, datamodule)

    # Determine the max num nodes and edges in training and validation
    logger.info("About to set the max nodes etc.")
    predictor.set_max_nodes_edges_per_graph(datamodule, stages=["train", "val"])

    # Run the model training
    with SafeRun(name="TRAINING", raise_error=cfg["constants"]["raise_train_error"], verbose=True):
        trainer.fit(model=predictor, datamodule=datamodule)

    # Determine the max num nodes and edges in testing
    predictor.set_max_nodes_edges_per_graph(datamodule, stages=["test"])

    # Run the model testing
    with SafeRun(name="TESTING", raise_error=cfg["constants"]["raise_train_error"], verbose=True):
        trainer.test(model=predictor, datamodule=datamodule)  # , ckpt_path=ckpt_path)

    logger.info("--------------------------------------------")
    logger.info("total computation used", timeit.default_timer() - st)
    logger.info("--------------------------------------------")
    wandb.finish()

    return trainer.callback_metrics


if __name__ == "__main__":
    main()


# expts/main_run_get_fingerprints.py

# General imports
import os
from os.path import dirname, abspath
import yaml
import numpy as np
import pandas as pd
import torch
import fsspec
from lightning.pytorch.utilities.model_summary import ModelSummary

# Current project imports
import graphium
from graphium.config._loader import load_datamodule, load_trainer
from graphium.utils.fs import mkdir
from graphium.trainer.predictor import PredictorModule


# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(graphium.__file__)))
os.chdir(MAIN_DIR)


# MODEL_FILE = "models_checkpoints/micro_ZINC/model.ckpt"
# CONFIG_FILE = "expts/config_micro_ZINC.yaml"


def main() -> None:
    LIST_CONCAT_LAST_LAYERS = [1, 0, [1, 2], [0, 1, 2]]
    DATA_NAME_ALL = ["molbace"]  # , "mollipo", "moltox21", "molHIV"]
    MODEL_PATH = "gs://graphium-public/pretrained-models"
    MODEL_NAME = "graphium-zinc-micro-dummy-test"
    MODEL_FILE = f"{MODEL_PATH}/{MODEL_NAME}/model.ckpt"
    MODEL_CONFIG = f"{MODEL_PATH}/{MODEL_NAME}/configs.yaml"

    predictor = PredictorModule.load_from_checkpoint(MODEL_FILE)

    print(predictor.model)
    print(ModelSummary(predictor, max_depth=4))

    for data_name in DATA_NAME_ALL:
        DATA_CONFIG = f"{MAIN_DIR}/expts/config_{data_name}_pretrained.yaml"

        with fsspec.open(DATA_CONFIG, "r") as f:
            data_cfg = yaml.safe_load(f)
        with fsspec.open(os.path.join(MODEL_CONFIG), "r") as f:
            model_cfg = yaml.safe_load(f)

        # Load and initialize the dataset
        data_cfg["datamodule"]["args"]["featurization"] = model_cfg["datamodule"]["args"]["featurization"]
        datamodule = load_datamodule(data_cfg)
        print("\ndatamodule:\n", datamodule, "\n")

        for concat_last_layers in LIST_CONCAT_LAST_LAYERS:
            export_dir = f"predictions/fingerprints-model-{MODEL_NAME}"
            export_df_path = f"{export_dir}/{data_name}-concatlayers-{concat_last_layers}.csv.gz"

            predictor.model.concat_last_layers = concat_last_layers
            trainer = load_trainer(data_cfg)

            # Run the model prediction
            preds = trainer.predict(model=predictor, datamodule=datamodule)
            if isinstance(preds[0], torch.Tensor):
                preds = [p.detach().cpu().numpy() for p in preds]
            preds = np.concatenate(preds, axis=0)

            # Generate output dataframe
            df = {"SMILES": datamodule.dataset.smiles}

            target = datamodule.dataset.labels
            for ii in range(target.shape[1]):
                df[f"Target-{ii}"] = target[:, ii]

            for ii in range(preds.shape[1]):
                df[f"Preds-{ii}"] = preds[:, ii]
            df = pd.DataFrame(df)
            mkdir(export_dir)
            df.to_csv(export_df_path)
            print(df)
            print(f"file saved to:`{export_df_path}`")


if __name__ == "__main__":
    main()


# expts/main_run_multitask.py

# General imports
import os
from os.path import dirname, abspath
from omegaconf import DictConfig, OmegaConf
import timeit
from loguru import logger
from datetime import datetime
from lightning.pytorch.utilities.model_summary import ModelSummary

# Current project imports
import graphium
from graphium.config._loader import (
    load_datamodule,
    load_metrics,
    load_architecture,
    load_predictor,
    load_trainer,
    save_params_to_wandb,
    load_accelerator,
)
from graphium.utils.safe_run import SafeRun

import hydra

# WandB
import wandb

# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(graphium.__file__)))
os.chdir(MAIN_DIR)


@hydra.main(version_base=None, config_path="hydra-configs", config_name="main")
def main(cfg: DictConfig) -> None:
    cfg = OmegaConf.to_container(cfg, resolve=True)

    run_name: str = "main"
    add_date_time: bool = True

    st = timeit.default_timer()

    date_time_suffix = ""
    if add_date_time:
        date_time_suffix = datetime.now().strftime("%d.%m.%Y_%H.%M.%S")

    wandb.init(entity=cfg["constants"]["entity"], project=cfg["constants"]["name"], config=cfg)

    # Initialize the accelerator
    cfg, accelerator_type = load_accelerator(cfg)

    # Load and initialize the dataset
    datamodule = load_datamodule(cfg, accelerator_type)

    # Initialize the network
    model_class, model_kwargs = load_architecture(
        cfg,
        in_dims=datamodule.in_dims,
    )

    datamodule.prepare_data()

    metrics = load_metrics(cfg)
    logger.info(metrics)

    predictor = load_predictor(
        cfg, model_class, model_kwargs, metrics, accelerator_type, datamodule.task_norms
    )

    logger.info(predictor.model)
    logger.info(ModelSummary(predictor, max_depth=4))

    trainer = load_trainer(cfg, run_name, accelerator_type, date_time_suffix)
    save_params_to_wandb(trainer.logger, cfg, predictor, datamodule)

    # Determine the max num nodes and edges in training and validation
    logger.info("About to set the max nodes etc.")
    predictor.set_max_nodes_edges_per_graph(datamodule, stages=["train", "val"])

    # Run the model training
    with SafeRun(name="TRAINING", raise_error=cfg["constants"]["raise_train_error"], verbose=True):
        trainer.fit(model=predictor, datamodule=datamodule)

    # Determine the max num nodes and edges in testing
    predictor.set_max_nodes_edges_per_graph(datamodule, stages=["test"])

    # Run the model testing
    with SafeRun(name="TESTING", raise_error=cfg["constants"]["raise_train_error"], verbose=True):
        trainer.test(model=predictor, datamodule=datamodule)  # , ckpt_path=ckpt_path)

    logger.info("--------------------------------------------")
    logger.info("total computation used", timeit.default_timer() - st)
    logger.info("--------------------------------------------")
    wandb.finish()

    return trainer.callback_metrics


if __name__ == "__main__":
    main()


# expts/main_run_predict.py

# General imports
import os
from os.path import dirname, abspath
import yaml
from copy import deepcopy
from omegaconf import DictConfig
import numpy as np
import pandas as pd
import torch
from lightning.pytorch.utilities.model_summary import ModelSummary

# Current project imports
import graphium
from graphium.config._loader import load_datamodule, load_trainer
from graphium.utils.fs import mkdir
from graphium.trainer.predictor import PredictorModule


# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(graphium.__file__)))
os.chdir(MAIN_DIR)

DATA_NAME = "molhiv"
MODEL_FILE = "models_checkpoints/ogb-molpcba/model-v2.ckpt"
CONFIG_FILE = f"expts/config_{DATA_NAME}_pretrained.yaml"

# MODEL_FILE = "models_checkpoints/micro_ZINC/model.ckpt"
# CONFIG_FILE = "expts/config_micro_ZINC.yaml"


def main(cfg: DictConfig) -> None:
    cfg = deepcopy(cfg)

    # Load and initialize the dataset
    datamodule = load_datamodule(cfg)
    print("\ndatamodule:\n", datamodule, "\n")

    export_df_path = f"predictions/preds-{DATA_NAME}.csv.gz"

    predictor = PredictorModule.load_from_checkpoint(MODEL_FILE)

    print(predictor.model)
    print(ModelSummary(predictor, max_depth=4))

    trainer = load_trainer(cfg)

    # Run the model prediction
    preds = trainer.predict(model=predictor, datamodule=datamodule)
    if isinstance(preds[0], torch.Tensor):
        preds = [p.detach().cpu().numpy() for p in preds]
    preds = np.concatenate(preds, axis=0)

    # Generate output dataframe
    df = {"SMILES": datamodule.dataset.smiles}

    target = datamodule.dataset.labels
    for ii in range(target.shape[1]):
        df[f"Target-{ii}"] = target[:, ii]

    for ii in range(preds.shape[1]):
        df[f"Preds-{ii}"] = preds[:, ii]
    df = pd.DataFrame(df)
    mkdir("predictions")
    df.to_csv(export_df_path)
    print(df)
    print(f"file saved to:`{export_df_path}`")


if __name__ == "__main__":
    with open(os.path.join(MAIN_DIR, CONFIG_FILE), "r") as f:
        cfg = yaml.safe_load(f)
    main(cfg)


# expts/main_run_test.py

# General imports
import os
from os.path import dirname, abspath
import yaml
from copy import deepcopy
from omegaconf import DictConfig
from lightning.pytorch.utilities.model_summary import ModelSummary

# Current project imports
import graphium
from graphium.config._loader import load_datamodule, load_metrics, load_trainer

from graphium.trainer.predictor import PredictorModule


# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(graphium.__file__)))
os.chdir(MAIN_DIR)

MODEL_FILE = "models_checkpoints/ogb-molpcba/model-v2.ckpt"

CONFIG_FILE = "expts/config_molPCBA.yaml"


def main(cfg: DictConfig) -> None:
    cfg = deepcopy(cfg)

    # Load and initialize the dataset
    datamodule = load_datamodule(cfg)
    print("\ndatamodule:\n", datamodule, "\n")

    metrics = load_metrics(cfg)
    print(metrics)

    predictor = PredictorModule.load_from_checkpoint(MODEL_FILE)
    predictor.metrics = metrics

    print(predictor.model)
    print(ModelSummary(predictor, max_depth=4))

    trainer = load_trainer(cfg)

    # Run the model testing
    trainer.test(model=predictor, datamodule=datamodule, ckpt_path=MODEL_FILE)


if __name__ == "__main__":
    with open(os.path.join(MAIN_DIR, CONFIG_FILE), "r") as f:
        cfg = yaml.safe_load(f)
    main(cfg)


# expts/run_validation_test.py

# General imports
import argparse
import os
from os.path import dirname, abspath
import yaml
from copy import deepcopy
from omegaconf import DictConfig
import timeit
from loguru import logger
from datetime import datetime
from pytorch_lightning.utilities.model_summary import ModelSummary

# Current project imports
import graphium
from graphium.config._loader import (
    load_datamodule,
    load_metrics,
    load_architecture,
    load_predictor,
    load_trainer,
    save_params_to_wandb,
    load_accelerator,
    load_yaml_config,
)
from graphium.utils.safe_run import SafeRun


# WandB
import wandb

# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(graphium.__file__)))

# CONFIG_FILE = "expts/configs/config_mpnn_10M_b3lyp.yaml"
# CONFIG_FILE = "expts/configs/config_mpnn_10M_pcqm4m.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_debug.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_large_mpnn.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_large_gcn.yaml"
CONFIG_FILE = "expts/neurips2023_configs/debug/config_large_gcn_debug.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_large_gin.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_large_gine.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_small_gcn.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_large_gcn.yaml"
# CONFIG_FILE = "exptas/neurips2023_configs/config_small_gin.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_small_gine.yaml"
os.chdir(MAIN_DIR)


def main(cfg: DictConfig, run_name: str = "main", add_date_time: bool = True) -> None:
    st = timeit.default_timer()

    date_time_suffix = ""
    if add_date_time:
        date_time_suffix = datetime.now().strftime("%d.%m.%Y_%H.%M.%S")

    cfg = deepcopy(cfg)
    wandb.init(project=cfg["constants"]["name"], config=cfg)

    # Initialize the accelerator
    cfg, accelerator_type = load_accelerator(cfg)

    # Load and initialize the dataset
    datamodule = load_datamodule(cfg, accelerator_type)

    # Initialize the network
    model_class, model_kwargs = load_architecture(
        cfg,
        in_dims=datamodule.in_dims,
    )

    datamodule.prepare_data()

    metrics = load_metrics(cfg)
    logger.info(metrics)

    predictor = load_predictor(
        cfg, model_class, model_kwargs, metrics, accelerator_type, datamodule.task_norms
    )
    logger.info(predictor.model)
    logger.info(ModelSummary(predictor, max_depth=4))

    trainer = load_trainer(cfg, run_name, accelerator_type, date_time_suffix)
    save_params_to_wandb(trainer.logger, cfg, predictor, datamodule)

    # Determine the max num nodes and edges in training and validation
    predictor.set_max_nodes_edges_per_graph(datamodule, stages=["val"])

    # Run the model validation
    with SafeRun(name="VALIDATING", raise_error=cfg["constants"]["raise_train_error"], verbose=True):
        trainer.validate(
            model=predictor,
            ckpt_path=f'{cfg["trainer"]["model_checkpoint"]["dirpath"]}{cfg["trainer"]["seed"]}/{cfg["trainer"]["model_checkpoint"]["filename"]}.ckpt',
            datamodule=datamodule,
        )

    # Run the model testing
    with SafeRun(name="TESTING", raise_error=cfg["constants"]["raise_train_error"], verbose=True):
        trainer.test(
            model=predictor,
            ckpt_path=f'{cfg["trainer"]["model_checkpoint"]["dirpath"]}{cfg["trainer"]["seed"]}/{cfg["trainer"]["model_checkpoint"]["filename"]}.ckpt',
            datamodule=datamodule,
        )

    logger.info("--------------------------------------------")
    logger.info("total computation used", timeit.default_timer() - st)
    logger.info("--------------------------------------------")
    wandb.finish()

    return trainer.callback_metrics


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", help="Path to the config file", default=None)

    args, unknown_args = parser.parse_known_args()
    if args.config is not None:
        CONFIG_FILE = args.config
    cfg = load_yaml_config(CONFIG_FILE, MAIN_DIR, unknown_args)

    main(cfg)


# profiling/profile_mol_to_graph.py

from tqdm import tqdm
import datamol as dm
import pickle

from graphium.data.utils import load_micro_zinc
from graphium.features.featurizer import mol_to_pyggraph, mol_to_adj_and_features, mol_to_graph_dict

# Check out this profiling tool: https://kirillstrelkov.medium.com/python-profiling-with-vscode-3a17c0407833


def main():
    df = load_micro_zinc()
    smiles = df["SMILES"].values.tolist()
    smiles = smiles * 10
    print("Num smiles: ", len(smiles))

    pos_enc = {
        "pos_type": "laplacian_eigvec",
        "num_pos": 3,
        "normalization": "none",
        "disconnected_comp": True,
    }
    atom_property_list_float1 = [
        "mass",
        "in-ring",
        "hybridization",
        "chirality",
        "aromatic",
        "degree",
        "formal-charge",
        "single-bond",
        "double-bond",
        "radical-electron",
    ]
    atom_property_list_float2 = ["electronegativity", "vdw-radius", "covalent-radius", "metal"]

    featurizer = {
        "atom_property_list_onehot": ["atomic-number", "valence"],
        "atom_property_list_float": atom_property_list_float1 + atom_property_list_float2,
        "edge_property_list": [
            "bond-type-onehot",
            "bond-type-float",
            "stereo",
            "in-ring",
            "conjugated",
            "estimated-bond-length",
        ],
        "add_self_loop": False,
        "explicit_H": False,
        "use_bonds_weights": False,
        "pos_encoding_as_features": pos_enc,
        # "on_error": "raise",
    }

    graphs = []
    for s in tqdm(smiles):
        mol = dm.to_mol(s)
        graphs.append(mol_to_graph_dict(mol, **featurizer))

    print(graphs[0])


if __name__ == "__main__":
    main()


# profiling/profile_one_of_k_encoding.py

from tqdm import tqdm
from graphium.utils.tensor import one_of_k_encoding


def main():
    CLASSES = ["AA", "BB", "CC", "DD", "EE"]
    CHOICES = CLASSES + ["FF"]

    for ii in tqdm(range(500000)):
        for choice in CHOICES:
            one_of_k_encoding(choice, CLASSES)

    print("DONE :)")


if __name__ == "__main__":
    main()


# profiling/profile_predictor.py

from tqdm import tqdm
import os
from time import time
import yaml
import fsspec

from graphium.config._loader import (
    load_datamodule,
    load_metrics,
    load_trainer,
    load_predictor,
    load_architecture,
)
from lightning import Trainer


def main():
    CONFIG_PATH = "expts/config_micro-PCBA.yaml"
    # DATA_PATH = "https://storage.googleapis.com/graphium-public/datasets/graphium-zinc-bench-gnn/smiles_score.csv.gz"

    with fsspec.open(CONFIG_PATH, "r") as f:
        cfg = yaml.safe_load(f)
    cfg["datamodule"]["args"]["cache_data_path"] = "graphium/data/cache/profiling/predictor_data.cache"
    # cfg["datamodule"]["args"]["df_path"] = DATA_PATH
    cfg["trainer"]["trainer"]["max_epochs"] = 5
    cfg["trainer"]["trainer"]["min_epochs"] = 5

    datamodule = load_datamodule(cfg)

    # Initialize the network
    model_class, model_kwargs = load_architecture(
        cfg,
        in_dim_nodes=datamodule.num_node_feats_with_positional_encoding,
        in_dim_edges=datamodule.num_edge_feats,
    )

    metrics = load_metrics(cfg)
    print(metrics)
    predictor = load_predictor(cfg, model_class, model_kwargs, metrics)
    trainer = load_trainer(cfg)
    trainer.fit(model=predictor, datamodule=datamodule)

    print("Done :)")


if __name__ == "__main__":
    main()


# scripts/convert_yml.py

"""
Convert the dependencies from conda's `env.yml` to pip `requirements.txt`
"""

import ruamel.yaml

yaml = ruamel.yaml.YAML()
data = yaml.load(open("env.yml"))

requirements = []
for dep in data["dependencies"]:
    if isinstance(dep, str):
        outputs = dep.split("=")
        if len(outputs) == 1:
            package = outputs[0]
            requirements.append(package)
        elif len(outputs) == 2:
            package, package_version = outputs[0], outputs[1]
            requirements.append(package + "==" + package_version)
        elif len(outputs) == 3:
            package, package_version, python_version = outputs[0], outputs[1], outputs[2]
            requirements.append(package + "==" + package_version)
    elif isinstance(dep, dict):
        for preq in dep.get("pip", []):
            requirements.append(preq)

with open("requirements.txt", "w") as fp:
    for requirement in requirements:
        print(requirement, file=fp)


# tests/conftest.py

import pathlib

import pytest

TEST_DIR_PATH = pathlib.Path(__file__).parent / "data"
DATA_DIR_PATH = TEST_DIR_PATH.absolute()


@pytest.fixture
def datadir(request):
    return DATA_DIR_PATH


# tests/test_architectures.py

"""
Unit tests for the different architectures of graphium/nn/architectures...

The layers are not thoroughly tested due to the difficulty of testing them
"""

import torch
import unittest as ut
from copy import deepcopy
import sys
import traceback

from graphium.nn.architectures import FeedForwardNN, FeedForwardPyg, FullGraphMultiTaskNetwork
from graphium.nn.base_layers import FCLayer
from graphium.nn.residual_connections import (
    ResidualConnectionConcat,
    ResidualConnectionDenseNet,
    ResidualConnectionNone,
    ResidualConnectionSimple,
    ResidualConnectionWeighted,
)
from torch_geometric.data import Data, Batch

from graphium.utils.spaces import LAYERS_DICT


class test_FeedForwardNN(ut.TestCase):
    kwargs = {
        "activation": "relu",
        "last_activation": "none",
        "normalization": "none",
        "dropout": 0.2,
        "name": "LNN",
        "layer_type": FCLayer,
    }

    norms = ["none", "batch_norm", "layer_norm"]

    def test_forward_no_residual(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [5, 6, 7]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="none",
            residual_skip_steps=1,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2])

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])

    def test_forward_simple_residual_1(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="simple",
            residual_skip_steps=1,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2])

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])

    def test_forward_norms(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        for normalization in self.norms:
            this_kwargs = deepcopy(self.kwargs)
            this_kwargs["normalization"] = normalization
            err_msg = f"normalization = {normalization}"
            lnn = FeedForwardNN(
                in_dim=in_dim,
                out_dim=out_dim,
                hidden_dims=hidden_dims,
                residual_type="simple",
                residual_skip_steps=1,
                **this_kwargs,
            )

            self.assertEqual(len(lnn.layers), len(hidden_dims) + 1, msg=err_msg)
            self.assertEqual(lnn.layers[0].in_dim, in_dim, msg=err_msg)
            self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0], msg=err_msg)
            self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1], msg=err_msg)
            self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2], msg=err_msg)

            feat = torch.FloatTensor(batch, in_dim)
            feat_out = lnn.forward(feat)

            self.assertListEqual(list(feat_out.shape), [batch, out_dim], msg=err_msg)

    def test_forward_simple_residual_2(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="simple",
            residual_skip_steps=2,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2])
        self.assertEqual(lnn.layers[4].in_dim, hidden_dims[3])
        self.assertEqual(lnn.layers[5].in_dim, hidden_dims[4])

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])

    def test_forward_concat_residual_1(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="concat",
            residual_skip_steps=1,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, 2 * hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, 2 * hidden_dims[2])
        self.assertEqual(lnn.layers[4].in_dim, 2 * hidden_dims[3])
        self.assertEqual(lnn.layers[5].in_dim, 2 * hidden_dims[4])

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])

    def test_forward_concat_residual_2(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="concat",
            residual_skip_steps=2,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, 1 * hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, 1 * hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, 2 * hidden_dims[2])
        self.assertEqual(lnn.layers[4].in_dim, 1 * hidden_dims[3])
        self.assertEqual(lnn.layers[5].in_dim, 2 * hidden_dims[4])

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])

    def test_forward_densenet_residual_1(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="densenet",
            residual_skip_steps=1,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, 2 * hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, 3 * hidden_dims[2])
        self.assertEqual(lnn.layers[4].in_dim, 4 * hidden_dims[3])
        self.assertEqual(lnn.layers[5].in_dim, 5 * hidden_dims[4])

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])

    def test_forward_densenet_residual_2(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="densenet",
            residual_skip_steps=2,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, 1 * hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, 2 * hidden_dims[2])
        self.assertEqual(lnn.layers[4].in_dim, 1 * hidden_dims[3])
        self.assertEqual(lnn.layers[5].in_dim, 3 * hidden_dims[4])

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])

    def test_forward_weighted_residual_1(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="weighted",
            residual_skip_steps=1,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2])
        self.assertEqual(lnn.layers[4].in_dim, hidden_dims[3])
        self.assertEqual(lnn.layers[5].in_dim, hidden_dims[4])

        self.assertEqual(len(lnn.residual_layer.residual_list), len(hidden_dims))

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])

    def test_forward_weighted_residual_2(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="weighted",
            residual_skip_steps=2,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2])
        self.assertEqual(lnn.layers[4].in_dim, hidden_dims[3])
        self.assertEqual(lnn.layers[5].in_dim, hidden_dims[4])

        self.assertEqual(len(lnn.residual_layer.residual_list), (len(hidden_dims) // 2 + 1))

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])


class test_FeedForwardGraph(ut.TestCase):
    kwargs = {
        "activation": "relu",
        "last_activation": "none",
        "dropout": 0.2,
        "name": "LNN",
    }

    in_dim = 7
    out_dim = 11
    in_dim_edges = 13
    hidden_dims = [6, 6, 6, 6, 6]

    edge_idx1 = (torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3]))
    edge_idx2 = (torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0]))
    num_edges1 = len(edge_idx1[0])
    num_nodes1 = max(edge_idx1[0].max(), edge_idx1[1].max()) + 1
    num_edges2 = len(edge_idx2[0])
    num_nodes2 = max(edge_idx2[0].max(), edge_idx2[1].max()) + 1
    h1 = torch.zeros(num_nodes1, in_dim, dtype=torch.float32)
    e1 = torch.ones(num_edges1, in_dim_edges, dtype=torch.float32)
    h2 = torch.ones(num_nodes2, in_dim, dtype=torch.float32)
    e2 = torch.zeros(num_edges2, in_dim_edges, dtype=torch.float32)

    g1 = Data(feat=h1, edge_index=torch.stack(edge_idx1), edge_feat=e1)
    g2 = Data(feat=h2, edge_index=torch.stack(edge_idx2), edge_feat=e2)
    data_list = [g1, g2, deepcopy(g1), deepcopy(g2)]
    batch_pyg = Batch.from_data_list(data_list)
    num_nodes = batch_pyg.num_nodes
    num_edges = batch_pyg.num_edges
    batch_size = len(data_list)

    virtual_nodes = ["none", "mean", "sum"]
    norms = ["none", "batch_norm", "layer_norm"]
    pna_kwargs = {"aggregators": ["mean", "max", "sum"], "scalers": ["identity", "amplification"]}

    layers_kwargs = {
        "pyg:gin": {},
        "pyg:gine": {"in_dim_edges": in_dim_edges},
        "pyg:gated-gcn": {"in_dim_edges": in_dim_edges, "hidden_dims_edges": hidden_dims},
        "pyg:pna-msgpass#1": {"layer_kwargs": pna_kwargs, "in_dim_edges": 0},
        "pyg:pna-msgpass#2": {"layer_kwargs": pna_kwargs, "in_dim_edges": in_dim_edges},
    }

    def test_forward_no_residual(self):
        for residual_skip_steps in [1, 2]:
            for virtual_node in self.virtual_nodes:
                for normalization in self.norms:
                    for layer_name, this_kwargs in self.layers_kwargs.items():
                        err_msg = f"virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                        layer_type = layer_name.split("#")[0]

                        # PYG
                        if layer_type.startswith("pyg:"):
                            layer_class = FeedForwardPyg
                            bg = deepcopy(self.batch_pyg)

                        gnn = layer_class(
                            in_dim=self.in_dim,
                            out_dim=self.out_dim,
                            hidden_dims=self.hidden_dims,
                            residual_type="none",
                            residual_skip_steps=residual_skip_steps,
                            layer_type=layer_type,
                            normalization=normalization,
                            **this_kwargs,
                            **self.kwargs,
                        )
                        # gnn.to(torch.float32)

                        self.assertIsInstance(gnn.residual_layer, ResidualConnectionNone)
                        self.assertEqual(len(gnn.layers), len(self.hidden_dims) + 1, msg=err_msg)
                        self.assertEqual(gnn.layers[0].out_dim, self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[1].out_dim, self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[2].out_dim, self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[3].out_dim, self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[4].out_dim, self.hidden_dims[4], msg=err_msg)
                        self.assertEqual(gnn.layers[5].out_dim, self.out_dim, msg=err_msg)

                        f = gnn.layers[0].out_dim_factor
                        self.assertEqual(gnn.layers[0].in_dim, self.in_dim, msg=err_msg)
                        self.assertEqual(gnn.layers[1].in_dim, f * self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[2].in_dim, f * self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[3].in_dim, f * self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[4].in_dim, f * self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[5].in_dim, f * self.hidden_dims[4], msg=err_msg)

                        out_g = gnn.forward(bg)
                        feat_out = out_g["feat"]
                        edge_feat_out = out_g["edge_feat"]

                        out_dim_edges = (
                            gnn.full_dims_edges[-1] if gnn.full_dims_edges is not None else self.in_dim_edges
                        )

                        self.assertListEqual(
                            list(feat_out.shape), [self.num_nodes, self.out_dim], msg=err_msg
                        )
                        self.assertListEqual(
                            list(edge_feat_out.shape), [self.num_edges, out_dim_edges], msg=err_msg
                        )

    def test_forward_simple_residual(self):
        for residual_skip_steps in [1, 2]:
            for virtual_node in self.virtual_nodes:
                for normalization in self.norms:
                    for layer_name, this_kwargs in self.layers_kwargs.items():
                        err_msg = f"virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                        layer_type = layer_name.split("#")[0]

                        # PYG
                        if layer_type.startswith("pyg:"):
                            layer_class = FeedForwardPyg
                            bg = deepcopy(self.batch_pyg)

                        gnn = layer_class(
                            in_dim=self.in_dim,
                            out_dim=self.out_dim,
                            hidden_dims=self.hidden_dims,
                            residual_type="simple",
                            residual_skip_steps=1,
                            layer_type=layer_type,
                            **this_kwargs,
                            **self.kwargs,
                        )
                        # gnn.to(torch.float32)

                        self.assertIsInstance(gnn.residual_layer, ResidualConnectionSimple)
                        self.assertEqual(len(gnn.layers), len(self.hidden_dims) + 1, msg=err_msg)
                        self.assertEqual(gnn.layers[0].out_dim, self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[1].out_dim, self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[2].out_dim, self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[3].out_dim, self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[4].out_dim, self.hidden_dims[4], msg=err_msg)
                        self.assertEqual(gnn.layers[5].out_dim, self.out_dim, msg=err_msg)

                        f = gnn.layers[0].out_dim_factor
                        self.assertEqual(gnn.layers[0].in_dim, self.in_dim, msg=err_msg)
                        self.assertEqual(gnn.layers[1].in_dim, f * self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[2].in_dim, f * self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[3].in_dim, f * self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[4].in_dim, f * self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[5].in_dim, f * self.hidden_dims[4], msg=err_msg)

                        out_g = gnn.forward(bg)
                        feat_out = out_g["feat"]
                        edge_feat_out = out_g["edge_feat"]

                        out_dim_edges = (
                            gnn.full_dims_edges[-1] if gnn.full_dims_edges is not None else self.in_dim_edges
                        )

                        self.assertListEqual(
                            list(feat_out.shape), [self.num_nodes, self.out_dim], msg=err_msg
                        )
                        self.assertListEqual(
                            list(edge_feat_out.shape), [self.num_edges, out_dim_edges], msg=err_msg
                        )

    def test_forward_weighted_residual(self):
        for residual_skip_steps in [1, 2]:
            for virtual_node in self.virtual_nodes:
                for normalization in self.norms:
                    for layer_name, this_kwargs in self.layers_kwargs.items():
                        err_msg = f"virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                        layer_type = layer_name.split("#")[0]

                        # PYG
                        if layer_type.startswith("pyg:"):
                            layer_class = FeedForwardPyg
                            bg = deepcopy(self.batch_pyg)

                        gnn = layer_class(
                            in_dim=self.in_dim,
                            out_dim=self.out_dim,
                            hidden_dims=self.hidden_dims,
                            residual_type="weighted",
                            residual_skip_steps=residual_skip_steps,
                            layer_type=layer_type,
                            **this_kwargs,
                            **self.kwargs,
                        )
                        # gnn.to(torch.float32)

                        self.assertIsInstance(gnn.residual_layer, ResidualConnectionWeighted)
                        self.assertEqual(len(gnn.layers), len(self.hidden_dims) + 1, msg=err_msg)
                        self.assertEqual(gnn.layers[0].out_dim, self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[1].out_dim, self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[2].out_dim, self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[3].out_dim, self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[4].out_dim, self.hidden_dims[4], msg=err_msg)
                        self.assertEqual(gnn.layers[5].out_dim, self.out_dim, msg=err_msg)

                        f = gnn.layers[0].out_dim_factor
                        self.assertEqual(gnn.layers[0].in_dim, self.in_dim, msg=err_msg)
                        self.assertEqual(gnn.layers[1].in_dim, f * self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[2].in_dim, f * self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[3].in_dim, f * self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[4].in_dim, f * self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[5].in_dim, f * self.hidden_dims[4], msg=err_msg)

                        out_g = gnn.forward(bg)
                        feat_out = out_g["feat"]
                        edge_feat_out = out_g["edge_feat"]

                        out_dim_edges = (
                            gnn.full_dims_edges[-1] if gnn.full_dims_edges is not None else self.in_dim_edges
                        )

                        self.assertListEqual(
                            list(feat_out.shape), [self.num_nodes, self.out_dim], msg=err_msg
                        )
                        self.assertListEqual(
                            list(edge_feat_out.shape), [self.num_edges, out_dim_edges], msg=err_msg
                        )

    def test_forward_concat_residual(self):
        for residual_skip_steps in [1, 2]:
            for virtual_node in self.virtual_nodes:
                for normalization in self.norms:
                    for layer_name, this_kwargs in self.layers_kwargs.items():
                        err_msg = f"virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                        layer_type = layer_name.split("#")[0]

                        # PYG
                        if layer_type.startswith("pyg:"):
                            layer_class = FeedForwardPyg
                            bg = deepcopy(self.batch_pyg)

                        gnn = layer_class(
                            in_dim=self.in_dim,
                            out_dim=self.out_dim,
                            hidden_dims=self.hidden_dims,
                            residual_type="concat",
                            residual_skip_steps=residual_skip_steps,
                            layer_type=layer_type,
                            **this_kwargs,
                            **self.kwargs,
                        )
                        # gnn.to(torch.float32)

                        self.assertIsInstance(gnn.residual_layer, ResidualConnectionConcat)
                        self.assertEqual(len(gnn.layers), len(self.hidden_dims) + 1, msg=err_msg)
                        self.assertEqual(gnn.layers[0].out_dim, self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[1].out_dim, self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[2].out_dim, self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[3].out_dim, self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[4].out_dim, self.hidden_dims[4], msg=err_msg)
                        self.assertEqual(gnn.layers[5].out_dim, self.out_dim, msg=err_msg)

                        f = gnn.layers[0].out_dim_factor
                        f2 = [2 * f if ((ii % residual_skip_steps) == 0 and ii > 0) else f for ii in range(6)]
                        self.assertEqual(gnn.layers[0].in_dim, self.in_dim, msg=err_msg)
                        self.assertEqual(gnn.layers[1].in_dim, f2[0] * self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[2].in_dim, f2[1] * self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[3].in_dim, f2[2] * self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[4].in_dim, f2[3] * self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[5].in_dim, f2[4] * self.hidden_dims[4], msg=err_msg)

                        out_g = gnn.forward(bg)
                        feat_out = out_g["feat"]
                        edge_feat_out = out_g["edge_feat"]

                        out_dim_edges = (
                            gnn.full_dims_edges[-1] if gnn.full_dims_edges is not None else self.in_dim_edges
                        )

                        self.assertListEqual(
                            list(feat_out.shape), [self.num_nodes, self.out_dim], msg=err_msg
                        )
                        self.assertListEqual(
                            list(edge_feat_out.shape), [self.num_edges, out_dim_edges], msg=err_msg
                        )

    def test_forward_densenet_residual(self):
        for residual_skip_steps in [1, 2]:
            for virtual_node in self.virtual_nodes:
                for normalization in self.norms:
                    for layer_name, this_kwargs in self.layers_kwargs.items():
                        err_msg = f"virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                        layer_type = layer_name.split("#")[0]

                        # PYG
                        if layer_type.startswith("pyg:"):
                            layer_class = FeedForwardPyg
                            bg = deepcopy(self.batch_pyg)

                        gnn = layer_class(
                            in_dim=self.in_dim,
                            out_dim=self.out_dim,
                            hidden_dims=self.hidden_dims,
                            residual_type="densenet",
                            residual_skip_steps=residual_skip_steps,
                            layer_type=layer_type,
                            **this_kwargs,
                            **self.kwargs,
                        )
                        # gnn.to(torch.float32)

                        self.assertIsInstance(gnn.residual_layer, ResidualConnectionDenseNet)
                        self.assertEqual(len(gnn.layers), len(self.hidden_dims) + 1, msg=err_msg)
                        self.assertEqual(gnn.layers[0].out_dim, self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[1].out_dim, self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[2].out_dim, self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[3].out_dim, self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[4].out_dim, self.hidden_dims[4], msg=err_msg)
                        self.assertEqual(gnn.layers[5].out_dim, self.out_dim, msg=err_msg)

                        f = gnn.layers[0].out_dim_factor
                        f2 = [
                            ((ii // residual_skip_steps) + 1) * f
                            if ((ii % residual_skip_steps) == 0 and ii > 0)
                            else f
                            for ii in range(6)
                        ]
                        self.assertEqual(gnn.layers[0].in_dim, self.in_dim, msg=err_msg)
                        self.assertEqual(gnn.layers[1].in_dim, f2[0] * self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[2].in_dim, f2[1] * self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[3].in_dim, f2[2] * self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[4].in_dim, f2[3] * self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[5].in_dim, f2[4] * self.hidden_dims[4], msg=err_msg)

                        out_g = gnn.forward(bg)
                        feat_out = out_g["feat"]
                        edge_feat_out = out_g["edge_feat"]

                        out_dim_edges = (
                            gnn.full_dims_edges[-1] if gnn.full_dims_edges is not None else self.in_dim_edges
                        )

                        self.assertListEqual(
                            list(feat_out.shape), [self.num_nodes, self.out_dim], msg=err_msg
                        )
                        self.assertListEqual(
                            list(edge_feat_out.shape), [self.num_edges, out_dim_edges], msg=err_msg
                        )


if __name__ == "__main__":
    ut.main()


# tests/test_attention.py

"""
Unit tests for the attention layer
"""

from ast import Assert
import numpy as np
import torch
import unittest as ut
from torch_geometric.data import Data, Batch
from copy import deepcopy
from graphium.ipu.to_dense_batch import to_dense_batch

from graphium.nn.base_layers import MultiheadAttentionMup


def seed_everything(seed: int):
    import random, os
    import numpy as np
    import torch

    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


class test_MultiHeadAttention(ut.TestCase):
    seed_everything(42)
    in_dim = 12
    out_dim = 12
    in_dim_edges = 10
    out_dim_edges = 10
    edge_idx1 = torch.stack([torch.tensor([0, 1, 2, 3, 2]), torch.tensor([1, 2, 3, 0, 0])])
    edge_idx2 = torch.stack([torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])])
    x1 = torch.randn(edge_idx1.max() + 1, in_dim, dtype=torch.float32)
    e1 = torch.randn(edge_idx1.shape[-1], in_dim_edges, dtype=torch.float32)
    x2 = torch.randn(edge_idx2.max() + 1, in_dim, dtype=torch.float32)
    e2 = torch.randn(edge_idx2.shape[-1], in_dim_edges, dtype=torch.float32)
    g1 = Data(feat=x1, edge_index=edge_idx1, edge_feat=e1)
    g2 = Data(feat=x2, edge_index=edge_idx2, edge_feat=e2)
    bg = Batch.from_data_list([g1, g2])

    attn_kwargs = {"embed_dim": in_dim, "num_heads": 2, "batch_first": True}

    def test_attention_class(self):
        bg = deepcopy(self.bg)
        seed_everything(42)
        attention_layer = MultiheadAttentionMup(biased_attention=False, **self.attn_kwargs)
        attention_layer.eval()
        seed_everything(42)
        attention_layer_bias = MultiheadAttentionMup(biased_attention=True, **self.attn_kwargs)
        attention_layer_bias.eval()

        h_dense, mask, _ = to_dense_batch(
            bg.feat,
            batch=bg.batch,
            batch_size=None,
            max_num_nodes_per_graph=None,
            drop_nodes_last_graph=False,
        )
        # attn_bias [batch, num_heads, nodes, nodes]
        nodes = h_dense.size()[1]
        attn_bias_3d = torch.zeros(2, 2, nodes, nodes)
        h_dense.nodepair_gaussian_bias_3d = attn_bias_3d
        # Apply attention layer and attention layer with bias.
        h_attn_output = attention_layer(
            h_dense,
            h_dense,
            h_dense,
            attn_bias=None,
            attn_mask=None,
            key_padding_mask=~mask,
            need_weights=False,
        )[0]
        h_attn_output_biased = attention_layer_bias(
            h_dense,
            h_dense,
            h_dense,
            attn_bias=attn_bias_3d,
            attn_mask=None,
            key_padding_mask=~mask,
            need_weights=False,
        )[0]
        self.assertEqual(h_attn_output.size(), h_attn_output_biased.size())
        np.testing.assert_array_almost_equal(
            h_attn_output.detach().numpy(), h_attn_output_biased.detach().numpy(), decimal=0
        )


if __name__ == "__main__":
    ut.main()


# tests/test_base_layers.py

"""
Unit tests for the different layers of graphium/nn/base_layers
"""

import torch
import unittest as ut
from torch_geometric.data import Data, Batch
from copy import deepcopy

from graphium.nn.base_layers import DropPath, TransformerEncoderLayerMup
from graphium.ipu.to_dense_batch import to_dense_batch, to_sparse_batch


class test_Base_Layers(ut.TestCase):
    in_dim = 21
    out_dim = 11
    in_dim_edges = 13
    out_dim_edges = 17

    edge_idx1 = torch.stack([torch.tensor([0, 1, 2, 3, 2]), torch.tensor([1, 2, 3, 0, 0])])
    edge_idx2 = torch.stack([torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])])
    x1 = torch.randn(edge_idx1.max() + 1, in_dim, dtype=torch.float32)
    e1 = torch.randn(edge_idx1.shape[-1], in_dim_edges, dtype=torch.float32)
    x2 = torch.randn(edge_idx2.max() + 1, in_dim, dtype=torch.float32)
    e2 = torch.randn(edge_idx2.shape[-1], in_dim_edges, dtype=torch.float32)

    g1 = Data(feat=x1, edge_index=edge_idx1, edge_feat=e1)
    g2 = Data(feat=x2, edge_index=edge_idx2, edge_feat=e2)
    bg = Batch.from_data_list([g1, g2])

    batch_size = 2
    max_num_nodes_per_graph = max(x1.shape[0], x2.shape[0])

    # for drop_rate=0.5, test if the output shape is correct
    def test_droppath_layer_0p5(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        layer = DropPath(drop_rate=0.5)
        h_out = layer.forward(feat_in, bg.batch)
        self.assertEqual(h_out.shape, feat_in.shape)

    # for drop_rate=1.0, test if the output are all zeros
    def test_droppath_layer_1p0(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        zero_tesor = torch.zeros(feat_in.shape)
        layer = DropPath(drop_rate=1.0)
        h_out = layer.forward(feat_in, bg.batch)
        self.assertTrue(torch.allclose(zero_tesor, h_out.detach()))

    # for drop_rate=0.0, test if the output matches the original output
    def test_droppath_layer_0p0(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        layer = DropPath(drop_rate=0.0)
        h_out = layer.forward(feat_in, bg.batch)
        self.assertTrue(torch.allclose(feat_in.detach(), h_out.detach()))

    # test output shape is correct for TransformerEncoderLayerMup
    def test_transformer_encoder_layer_mup(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        layer = TransformerEncoderLayerMup(
            biased_attention=False, d_model=self.in_dim, nhead=1, dim_feedforward=4 * self.in_dim
        )

        feat_dense, key_padding_mask, idx = to_dense_batch(
            feat_in,
            batch=bg.batch,
            batch_size=self.batch_size,
            max_num_nodes_per_graph=self.max_num_nodes_per_graph,
            drop_nodes_last_graph=False,
        )
        attn_mask = None
        key_padding_mask = ~key_padding_mask

        h_out_dense = layer.forward(feat_dense)

        h_out = to_sparse_batch(h_out_dense, mask_idx=idx)

        self.assertEqual(h_out.shape, feat_in.shape)


# tests/test_collate.py

"""
Unit tests for the collate
"""

import unittest as ut
from copy import deepcopy
import torch
import numpy as np
from torch_geometric.data import Data

from graphium.data.collate import collate_labels, graphium_collate_fn


class test_Collate(ut.TestCase):
    def test_collate_labels(self):
        # Create fake labels
        labels_size_dict = {
            "graph_label1": [1],
            "graph_label2": [3],
            "node_label2": [5],
            "edge_label3": [5, 2],
            "node_label4": [5, 1],
        }
        labels_dtype_dict = {
            "graph_label1": torch.float32,
            "graph_label2": torch.float16,
            "node_label2": torch.float32,
            "edge_label3": torch.float32,
            "node_label4": torch.float32,
        }
        fake_label = {
            "graph_label1": torch.FloatTensor([1]),
            "graph_label2": torch.HalfTensor([1, 2, 3]),
            "node_label2": torch.FloatTensor([1, 2, 3, 4, 5]),
            "edge_label3": torch.FloatTensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
            "node_label4": torch.FloatTensor([[1], [2], [3], [4], [5]]),
        }
        fake_labels = []
        num_labels = 10
        for i in range(num_labels):
            pyg_labels = Data(x=torch.empty(5, 0), edge_index=torch.empty(2, 5))
            for key, val in fake_label.items():
                pyg_labels[key] = val + 17 * 2
            fake_labels.append(pyg_labels)

        # Collate labels and check for the right shapes and dtypes
        collated_labels = collate_labels(
            deepcopy(fake_labels), deepcopy(labels_size_dict), deepcopy(labels_dtype_dict)
        )
        self.assertEqual(collated_labels["graph_label1"].shape, torch.Size([num_labels, 1]))  # , 1
        self.assertEqual(collated_labels["graph_label2"].shape, torch.Size([num_labels, 3]))  # , 1
        self.assertEqual(collated_labels["node_label2"].shape, torch.Size([num_labels * 5, 1]))  # , 5
        self.assertEqual(collated_labels["edge_label3"].shape, torch.Size([num_labels * 5, 2]))  # , 5, 2
        self.assertEqual(collated_labels["node_label4"].shape, torch.Size([num_labels * 5, 1]))  # , 5, 1

        self.assertEqual(collated_labels["graph_label1"].dtype, torch.float32)
        self.assertEqual(collated_labels["graph_label2"].dtype, torch.float16)

        # Check that the values are correct
        graph_label1_true = deepcopy(torch.stack([this_label["graph_label1"] for this_label in fake_labels]))
        graph_label2_true = deepcopy(torch.stack([this_label["graph_label2"] for this_label in fake_labels]))
        label2_true = deepcopy(torch.stack([this_label["node_label2"] for this_label in fake_labels]))
        label3_true = deepcopy(torch.stack([this_label["edge_label3"] for this_label in fake_labels]))
        label4_true = deepcopy(torch.stack([this_label["node_label4"] for this_label in fake_labels]))

        # NOTE: Flatten due to the way Data objects are collated (concat along first dim, instead of stacked)
        np.testing.assert_array_equal(collated_labels["graph_label1"].numpy(), graph_label1_true.numpy())
        np.testing.assert_array_equal(collated_labels["graph_label2"].numpy(), graph_label2_true.numpy())
        np.testing.assert_array_equal(
            collated_labels["node_label2"].numpy(), label2_true.flatten(0, 1).unsqueeze(1).numpy()
        )
        np.testing.assert_array_equal(
            collated_labels["edge_label3"].numpy(), label3_true.flatten(0, 1).numpy()
        )
        np.testing.assert_array_equal(
            collated_labels["node_label4"].numpy(), label4_true.flatten(0, 1).numpy()
        )

        # Remove some labels and check that the collation still works and puts `nan` in the right places
        missing_labels = {
            "graph_label1": [1, 3, 5, 7, 9],
            "graph_label2": [0, 4, 3, 1, 7],
            "node_label2": [0, 2, 4, 6, 8],
            "edge_label3": [0, 1, 2, 3, 4],
            "node_label4": [5, 1, 4, 9, 6],
        }
        for key, missing_idx in missing_labels.items():
            for idx in missing_idx:
                fake_labels[idx].pop(key)
        graph_label1_true[missing_labels["graph_label1"]] = float("nan")
        graph_label2_true[missing_labels["graph_label2"]] = float("nan")
        label2_true[missing_labels["node_label2"]] = float("nan")
        label3_true[missing_labels["edge_label3"]] = float("nan")
        label4_true[missing_labels["node_label4"]] = float("nan")

        # Collate labels and check for the right shapes
        labels_size_dict = {
            "graph_label1": [1],
            "graph_label2": [3],
            "node_label2": [5],
            "edge_label3": [5, 2],
            "node_label4": [5, 1],
        }
        collated_labels = collate_labels(
            deepcopy(fake_labels), deepcopy(labels_size_dict), deepcopy(labels_dtype_dict)
        )
        self.assertEqual(collated_labels["graph_label1"].shape, torch.Size([num_labels, 1]))  # , 1
        self.assertEqual(collated_labels["graph_label2"].shape, torch.Size([num_labels, 3]))  # , 1
        self.assertEqual(collated_labels["node_label2"].shape, torch.Size([num_labels * 5, 1]))  # , 5
        self.assertEqual(collated_labels["edge_label3"].shape, torch.Size([num_labels * 5, 2]))  # , 5, 2
        self.assertEqual(collated_labels["node_label4"].shape, torch.Size([num_labels * 5, 1]))  # , 5, 1

        # Check that the values are correct when some labels are missing
        # NOTE: Flatten due to the way Data objects are collated (concat along first dim, instead of stacked)
        np.testing.assert_array_equal(collated_labels["graph_label1"].numpy(), graph_label1_true.numpy())
        np.testing.assert_array_equal(collated_labels["graph_label2"].numpy(), graph_label2_true.numpy())
        np.testing.assert_array_equal(
            collated_labels["node_label2"].numpy(), label2_true.flatten(0, 1).unsqueeze(1).numpy()
        )
        np.testing.assert_array_equal(
            collated_labels["edge_label3"].numpy(), label3_true.flatten(0, 1).numpy()
        )
        np.testing.assert_array_equal(
            collated_labels["node_label4"].numpy(), label4_true.flatten(0, 1).numpy()
        )
        # Now test the `graphium_collate_fn` function when only labels are given
        fake_labels2 = [{"labels": this_label} for this_label in fake_labels]
        collated_labels = graphium_collate_fn(
            deepcopy(fake_labels2), labels_size_dict=labels_size_dict, labels_dtype_dict=labels_dtype_dict
        )["labels"]
        self.assertEqual(collated_labels["graph_label1"].shape, torch.Size([num_labels, 1]))
        self.assertEqual(collated_labels["graph_label2"].shape, torch.Size([num_labels, 3]))
        self.assertEqual(collated_labels["node_label2"].shape, torch.Size([num_labels * 5, 1]))  # , 5
        self.assertEqual(collated_labels["edge_label3"].shape, torch.Size([num_labels * 5, 2]))  # , 5, 2
        self.assertEqual(collated_labels["node_label4"].shape, torch.Size([num_labels * 5, 1]))  # , 5, 1

        # Check that the values are correct when some labels are missing
        # NOTE: Flatten due to the way Data objects are collated (concat along first dim, instead of stacked)
        np.testing.assert_array_equal(collated_labels["graph_label1"].numpy(), graph_label1_true.numpy())
        np.testing.assert_array_equal(collated_labels["graph_label2"].numpy(), graph_label2_true.numpy())
        np.testing.assert_array_equal(
            collated_labels["node_label2"].numpy(), label2_true.flatten(0, 1).unsqueeze(1).numpy()
        )
        np.testing.assert_array_equal(
            collated_labels["edge_label3"].numpy(), label3_true.flatten(0, 1).numpy()
        )
        np.testing.assert_array_equal(
            collated_labels["node_label4"].numpy(), label4_true.flatten(0, 1).numpy()
        )


if __name__ == "__main__":
    ut.main()


# tests/test_data_utils.py

import pandas as pd
import unittest as ut
import graphium
import tempfile


class TestDataUtils(ut.TestCase):
    def test_list_datasets(
        self,
    ):
        datasets = graphium.data.utils.list_graphium_datasets()
        assert isinstance(datasets, set)
        assert len(datasets) > 0

    def test_download_datasets(self):
        dataset_dir = tempfile.TemporaryDirectory()
        data_path = graphium.data.utils.download_graphium_dataset(
            "graphium-zinc-micro", output_path=dataset_dir.name
        )

        fpath = graphium.utils.fs.join(data_path, "ZINC-micro.csv")
        df = pd.read_csv(fpath)
        assert df.shape == (1000, 4)  # type: ignore
        assert df.columns.tolist() == ["SMILES", "SA", "logp", "score"]  # type: ignore


if __name__ == "__main__":
    ut.main()


# tests/test_datamodule.py

import unittest as ut
import numpy as np
import torch
import pandas as pd
import datamol as dm

import graphium
from graphium.utils.fs import rm, exists, get_size
from graphium.data import GraphOGBDataModule, MultitaskFromSmilesDataModule

TEMP_CACHE_DATA_PATH = "tests/temp_cache_0000"


class Test_DataModule(ut.TestCase):
    def test_ogb_datamodule(self):
        # other datasets are too large to be tested
        dataset_names = ["ogbg-molhiv", "ogbg-molpcba", "ogbg-moltox21", "ogbg-molfreesolv"]
        dataset_name = dataset_names[3]

        # Setup the featurization
        featurization_args = {}
        featurization_args["atom_property_list_float"] = []  # ["weight", "valence"]
        featurization_args["atom_property_list_onehot"] = ["atomic-number", "degree"]
        # featurization_args["conformer_property_list"] = ["positions_3d"]
        featurization_args["edge_property_list"] = ["bond-type-onehot"]
        featurization_args["add_self_loop"] = False
        featurization_args["use_bonds_weights"] = False
        featurization_args["explicit_H"] = False

        # Config for datamodule
        task_specific_args = {}
        task_specific_args["task_1"] = {"task_level": "graph", "dataset_name": dataset_name}
        dm_args = {}
        dm_args["cache_data_path"] = None
        dm_args["featurization"] = featurization_args
        dm_args["batch_size_training"] = 16
        dm_args["batch_size_inference"] = 16
        dm_args["num_workers"] = 0
        dm_args["pin_memory"] = True
        dm_args["featurization_n_jobs"] = 2
        dm_args["featurization_progress"] = True
        dm_args["featurization_backend"] = "loky"
        dm_args["featurization_batch_size"] = 50

        ds = GraphOGBDataModule(task_specific_args, **dm_args)

        ds.prepare_data()

        # Check the keys in the dataset
        ds.setup(save_smiles_and_ids=False)
        assert set(ds.train_ds[0].keys()) == {"features", "labels"}

        ds.setup(save_smiles_and_ids=True)
        assert set(ds.train_ds[0].keys()) == {"smiles", "mol_ids", "features", "labels"}

        # test module
        assert ds.num_edge_feats == 5
        assert ds.num_node_feats == 50
        assert len(ds) == 642

        # test batch loader
        batch = next(iter(ds.train_dataloader()))
        assert len(batch["smiles"]) == 16
        assert len(batch["labels"]["graph_task_1"]) == 16
        assert len(batch["mol_ids"]) == 16

    def test_none_filtering(self):
        # Create the objects to filter
        list_of_num = [ii for ii in range(100)]
        list_of_str = [str(ii) for ii in list_of_num]
        tuple_of_num = tuple(list_of_num)
        array_of_num = np.asarray(list_of_num)
        array_of_str = np.asarray(list_of_str)
        tensor_of_num = torch.as_tensor(array_of_num)
        arrays_of_num = np.stack([list_of_num, list_of_num, list_of_num], axis=1)
        arrays_of_str = np.stack([list_of_str, list_of_str, list_of_str], axis=1)
        tensors_of_num = torch.as_tensor(arrays_of_num)
        dic = {"str": list_of_str, "num": list_of_num}
        df = pd.DataFrame(dic)
        df_shuffled = df.sample(frac=1)
        series_num = df["num"]
        series_num_shuffled = df_shuffled["num"]

        # Create different indexes to use for filtering
        all_idx_none = [[3, 17, 88], [22, 33, 44, 55, 66, 77, 88], [], np.arange(len(list_of_num))]

        # Loop all the indexes and filter the objects.
        for ii, idx_none in enumerate(all_idx_none):
            msg = f"Failed for ii={ii}"

            # Create the true filtered sequences
            filtered_num = [ii for ii in range(100) if ii not in idx_none]
            filtered_str = [str(ii) for ii in filtered_num]
            assert len(filtered_num) == len(list_of_num) - len(idx_none)
            assert len(filtered_str) == len(list_of_str) - len(idx_none)

            # Filter the sequences from the Datamodule function
            (
                list_of_num_2,
                list_of_str_2,
                tuple_of_num_2,
                array_of_num_2,
                array_of_str_2,
                tensor_of_num_2,
                df_2,
                df_shuffled_2,
                dic_2,
                arrays_of_num_2,
                arrays_of_str_2,
                tensors_of_num_2,
                series_num_2,
                series_num_shuffled_2,
            ) = graphium.data.MultitaskFromSmilesDataModule._filter_none_molecules(
                idx_none,
                list_of_num,
                list_of_str,
                tuple_of_num,
                array_of_num,
                array_of_str,
                tensor_of_num,
                df,
                df_shuffled,
                dic,
                arrays_of_num,
                arrays_of_str,
                tensors_of_num,
                series_num,
                series_num_shuffled,
            )

            df_shuffled_2 = df_shuffled_2.sort_values(by="num", axis=0)
            series_num_shuffled_2 = series_num_shuffled_2.sort_values(axis=0)

            # Assert the filtering is done correctly
            self.assertListEqual(list_of_num_2, filtered_num, msg=msg)
            self.assertListEqual(list_of_str_2, filtered_str, msg=msg)
            self.assertListEqual(list(tuple_of_num_2), filtered_num, msg=msg)
            self.assertListEqual(array_of_num_2.tolist(), filtered_num, msg=msg)
            self.assertListEqual(array_of_str_2.tolist(), filtered_str, msg=msg)
            self.assertListEqual(tensor_of_num_2.tolist(), filtered_num, msg=msg)
            for jj in range(arrays_of_num.shape[1]):
                self.assertListEqual(arrays_of_num_2[:, jj].tolist(), filtered_num, msg=msg)
                self.assertListEqual(arrays_of_str_2[:, jj].tolist(), filtered_str, msg=msg)
                self.assertListEqual(tensors_of_num_2[:, jj].tolist(), filtered_num, msg=msg)
            self.assertListEqual(dic_2["num"], filtered_num, msg=msg)
            self.assertListEqual(dic_2["str"], filtered_str, msg=msg)
            self.assertListEqual(df_2["num"].tolist(), filtered_num, msg=msg)
            self.assertListEqual(df_2["str"].tolist(), filtered_str, msg=msg)
            self.assertListEqual(series_num_2.tolist(), filtered_num, msg=msg)

            # When the dataframe is shuffled, the lists are different because the filtering
            # is done on the row indexes, not the dataframe indexes.
            bool_to_check = (len(idx_none) == 0) or (len(idx_none) == len(df_shuffled))
            self.assertIs(df_shuffled_2["num"].tolist() == filtered_num, bool_to_check, msg=msg)
            self.assertIs(df_shuffled_2["str"].tolist() == filtered_str, bool_to_check, msg=msg)
            self.assertIs(series_num_shuffled_2.tolist() == filtered_num, bool_to_check, msg=msg)

    def test_caching(self):
        # other datasets are too large to be tested
        dataset_name = "ogbg-molfreesolv"

        # Setup the featurization
        featurization_args = {}
        featurization_args["atom_property_list_float"] = []  # ["weight", "valence"]
        featurization_args["atom_property_list_onehot"] = ["atomic-number", "degree"]
        # featurization_args["conformer_property_list"] = ["positions_3d"]
        featurization_args["edge_property_list"] = ["bond-type-onehot"]
        featurization_args["add_self_loop"] = False
        featurization_args["use_bonds_weights"] = False
        featurization_args["explicit_H"] = False

        # Config for datamodule
        task_specific_args = {}
        task_specific_args["task_1"] = {"task_level": "graph", "dataset_name": dataset_name}
        dm_args = {}
        dm_args["featurization"] = featurization_args
        dm_args["batch_size_training"] = 16
        dm_args["batch_size_inference"] = 16
        dm_args["num_workers"] = 0
        dm_args["pin_memory"] = True
        dm_args["featurization_n_jobs"] = 2
        dm_args["featurization_progress"] = True
        dm_args["featurization_backend"] = "loky"
        dm_args["featurization_batch_size"] = 50

        # Delete the cache if already exist
        if exists(TEMP_CACHE_DATA_PATH):
            rm(TEMP_CACHE_DATA_PATH, recursive=True)

        # Prepare the data. It should create the cache there
        assert not exists(TEMP_CACHE_DATA_PATH)
        ds = GraphOGBDataModule(task_specific_args, cache_data_path=TEMP_CACHE_DATA_PATH, **dm_args)
        assert not ds.load_data_from_cache(verbose=False)
        ds.prepare_data()

        # Check the keys in the dataset
        ds.setup(save_smiles_and_ids=False)
        assert set(ds.train_ds[0].keys()) == {"features", "labels"}

        ds.setup(save_smiles_and_ids=True)
        assert set(ds.train_ds[0].keys()) == {"smiles", "mol_ids", "features", "labels"}

        # Make sure that the cache is created
        full_cache_path = ds.get_data_cache_fullname(compress=False)
        assert exists(full_cache_path)
        assert get_size(full_cache_path) > 10000

        # Check that the data is loaded correctly from cache
        assert ds.load_data_from_cache(verbose=False)

        # test module
        assert ds.num_edge_feats == 5
        assert ds.num_node_feats == 50
        assert len(ds) == 642

        # test batch loader
        batch = next(iter(ds.train_dataloader()))
        assert len(batch["smiles"]) == 16
        assert len(batch["labels"]["graph_task_1"]) == 16
        assert len(batch["mol_ids"]) == 16

    def test_datamodule_with_none_molecules(self):
        # Setup the featurization
        featurization_args = {}
        featurization_args["atom_property_list_float"] = []  # ["weight", "valence"]
        featurization_args["atom_property_list_onehot"] = ["atomic-number", "degree"]
        featurization_args["edge_property_list"] = ["bond-type-onehot"]

        # Config for datamodule
        bad_csv = "tests/data/micro_ZINC_corrupt.csv"
        task_specific_args = {}
        task_kwargs = {"df_path": bad_csv, "split_val": 0.0, "split_test": 0.0}
        task_specific_args["task_1"] = {
            "task_level": "graph",
            "label_cols": "SA",
            "smiles_col": "SMILES1",
            **task_kwargs,
        }
        task_specific_args["task_2"] = {
            "task_level": "graph",
            "label_cols": "logp",
            "smiles_col": "SMILES2",
            **task_kwargs,
        }
        task_specific_args["task_3"] = {
            "task_level": "graph",
            "label_cols": "score",
            "smiles_col": "SMILES3",
            **task_kwargs,
        }

        # Read the corrupted dataset and get stats
        df = pd.read_csv(bad_csv)
        bad_smiles = (df["SMILES1"] == "XXX") & (df["SMILES2"] == "XXX") & (df["SMILES3"] == "XXX")
        num_bad_smiles = sum(bad_smiles)

        # Test the datamodule
        datamodule = MultitaskFromSmilesDataModule(
            task_specific_args=task_specific_args,
            featurization_args=featurization_args,
            featurization_n_jobs=0,
            featurization_batch_size=1,
        )
        datamodule.prepare_data()
        datamodule.setup(save_smiles_and_ids=True)

        # Check that the number of molecules is correct
        smiles = df["SMILES1"].tolist() + df["SMILES2"].tolist() + df["SMILES3"].tolist()
        num_unique_smiles = len(set(smiles)) - 1  # -1 because of the XXX
        # self.assertEqual(len(datamodule.train_ds), num_unique_smiles - num_bad_smiles)

        # Change the index of the dataframe
        index_smiles = []
        for ii in range(len(df)):
            if df["SMILES1"][ii] != "XXX":
                smiles = df["SMILES1"][ii]
            elif df["SMILES2"][ii] != "XXX":
                smiles = df["SMILES2"][ii]
            elif df["SMILES3"][ii] != "XXX":
                smiles = df["SMILES3"][ii]
            else:
                smiles = "XXX"
            index_smiles.append(smiles)
        df["idx_smiles"] = index_smiles
        df = df.set_index("idx_smiles")

        # Convert the smilies from the train_ds to a list, and check the content
        train_smiles = [d["smiles"] for d in datamodule.train_ds]

        # Check that the set of smiles are the same
        train_smiles_flat = list(set([item for sublist in train_smiles for item in sublist]))
        train_smiles_flat.sort()
        index_smiles_filt = list(set([smiles for smiles in index_smiles if smiles != "XXX"]))
        index_smiles_filt.sort()
        self.assertListEqual(train_smiles_flat, index_smiles_filt)

        # Check that the smiles are correct for each datapoint in the dataset
        for smiles in train_smiles:
            self.assertEqual(len(set(smiles)), 1)  # Check that all smiles are the same
            this_smiles = smiles[0]
            true_smiles = df.loc[this_smiles][["SMILES1", "SMILES2", "SMILES3"]]
            num_true_smiles = sum(true_smiles != "XXX")
            self.assertEqual(len(smiles), num_true_smiles)  # Check that the number of smiles is correct
            self.assertEqual(
                this_smiles, true_smiles[true_smiles != "XXX"].values[0]
            )  # Check that the smiles are correct

        # Convert the labels from the train_ds to a dataframe
        train_labels = [{task: val[0] for task, val in d["labels"].items()} for d in datamodule.train_ds]
        train_labels_df = pd.DataFrame(train_labels)
        train_labels_df = train_labels_df.rename(
            columns={"graph_task_1": "graph_SA", "graph_task_2": "graph_logp", "graph_task_3": "graph_score"}
        )
        train_labels_df["smiles"] = [s[0] for s in datamodule.train_ds.smiles]
        train_labels_df = train_labels_df.set_index("smiles")
        train_labels_df = train_labels_df.sort_index()

        # Check that the labels are correct
        df2 = df.reset_index()[~bad_smiles].set_index("idx_smiles").sort_index()
        labels = train_labels_df[["graph_SA", "graph_logp", "graph_score"]].values
        nans = np.isnan(labels)
        true_nans = df2[["SMILES1", "SMILES2", "SMILES3"]].values == "XXX"
        true_labels = df2[["SA", "logp", "score"]].values
        true_labels[true_nans] = np.nan
        np.testing.assert_array_equal(nans, true_nans)  # Check that the nans are correct
        np.testing.assert_array_almost_equal(
            labels, true_labels, decimal=5
        )  # Check that the label values are correct

    def test_datamodule_multiple_data_files(self):
        # Test single CSV files
        csv_file = "tests/data/micro_ZINC_shard_1.csv"
        task_kwargs = {"df_path": csv_file, "split_val": 0.0, "split_test": 0.0}
        task_specific_args = {
            "task": {"task_level": "graph", "label_cols": ["score"], "smiles_col": "SMILES", **task_kwargs}
        }

        ds = MultitaskFromSmilesDataModule(task_specific_args)
        ds.prepare_data()
        ds.setup()

        self.assertEqual(len(ds.train_ds), 10)

        # Test multi CSV files
        csv_file = "tests/data/micro_ZINC_shard_*.csv"
        task_kwargs = {"df_path": csv_file, "split_val": 0.0, "split_test": 0.0}
        task_specific_args = {
            "task": {"task_level": "graph", "label_cols": ["score"], "smiles_col": "SMILES", **task_kwargs}
        }

        ds = MultitaskFromSmilesDataModule(task_specific_args)
        ds.prepare_data()
        ds.setup()

        self.assertEqual(len(ds.train_ds), 20)

        # Test single Parquet files
        parquet_file = "tests/data/micro_ZINC_shard_1.parquet"
        task_kwargs = {"df_path": parquet_file, "split_val": 0.0, "split_test": 0.0}
        task_specific_args = {
            "task": {"task_level": "graph", "label_cols": ["score"], "smiles_col": "SMILES", **task_kwargs}
        }

        ds = MultitaskFromSmilesDataModule(task_specific_args)
        ds.prepare_data()
        ds.setup()

        self.assertEqual(len(ds.train_ds), 10)

        # Test multi Parquet files
        parquet_file = "tests/data/micro_ZINC_shard_*.parquet"
        task_kwargs = {"df_path": parquet_file, "split_val": 0.0, "split_test": 0.0}
        task_specific_args = {
            "task": {"task_level": "graph", "label_cols": ["score"], "smiles_col": "SMILES", **task_kwargs}
        }

        ds = MultitaskFromSmilesDataModule(task_specific_args)
        ds.prepare_data()
        ds.setup()

        self.assertEqual(len(ds.train_ds), 20)


if __name__ == "__main__":
    ut.main()

    # Delete the cache
    if exists(TEMP_CACHE_DATA_PATH):
        rm(TEMP_CACHE_DATA_PATH, recursive=True)


# tests/test_dataset.py

import unittest as ut

from graphium.data import load_micro_zinc
from graphium.data.dataset import SingleTaskDataset, MultitaskDataset
from graphium.data.smiles_transform import smiles_to_unique_mol_ids
from graphium.data.utils import get_keys


class Test_Multitask_Dataset(ut.TestCase):
    # Then we can choose different rows and columns for the tests as we see fit.
    # Remember tests are supposed to be FAST, and reading from the file system multiple times slows things down.

    # Make sure that the inputs to single task datasets are always lists!
    # Do not pass a data frame itself, but turn it into a list to satisfy the type required.

    def test_multitask_dataset_case_1(self):
        """Case: different tasks, all with the same smiles set.
        - Check that for each task, all smiles are received from the initial DF.
        - Check that for each task, you have the same label values as the initial DF.
        """

        df_micro_zinc = load_micro_zinc()  # Has about 1000 molecules
        df = df_micro_zinc.iloc[0:4]
        num_unique_mols = 4

        # Here we take the microzinc dataset and split the labels up into 'SA', 'logp' and 'score' in order to simulate having multiple single-task datasets
        df_micro_zinc_SA = df[["SMILES", "SA"]]
        df_micro_zinc_logp = df[["SMILES", "logp"]]
        df_micro_zinc_score = df[["SMILES", "score"]]

        # We need to turn these dataframes into single-task datasets.
        # We don't need to do featurization yet.
        ds_micro_zinc_SA = SingleTaskDataset(
            smiles=df_micro_zinc_SA.loc[:, "SMILES"].tolist(), labels=df_micro_zinc_SA.loc[:, "SA"].tolist()
        )

        ds_micro_zinc_logp = SingleTaskDataset(
            smiles=df_micro_zinc_logp.loc[:, "SMILES"].tolist(),
            labels=df_micro_zinc_logp.loc[:, "logp"].tolist(),
        )
        ds_micro_zinc_score = SingleTaskDataset(
            smiles=df_micro_zinc_score.loc[:, "SMILES"].tolist(),
            labels=df_micro_zinc_score.loc[:, "score"].tolist(),
        )

        # Create the multitask dataset
        datasets_dict = {"SA": ds_micro_zinc_SA, "logp": ds_micro_zinc_logp, "score": ds_micro_zinc_score}
        multitask_microzinc = MultitaskDataset(
            datasets_dict, save_smiles_and_ids=True
        )  # Can optionally have features

        # Check: The number of unique molecules equals the number of datapoints in the multitask dataset.
        self.assertEqual(num_unique_mols, multitask_microzinc.__len__())

        # Check that for each task, you have the same label values as the initial DF.
        for idx in range(multitask_microzinc.__len__()):
            smiles = df[["SMILES"]].iloc[idx].values[0]
            # label = df[['SA']].iloc[idx]
            label_SA = ds_micro_zinc_SA.labels[idx]
            label_logp = ds_micro_zinc_logp.labels[idx]
            label_score = ds_micro_zinc_score.labels[idx]

            # Search for the mol id in the multitask dataset
            mol_ids = smiles_to_unique_mol_ids([smiles])
            mol_id = mol_ids[0]
            found_idx = -1
            for i, id in enumerate(multitask_microzinc.mol_ids):
                if mol_id == id:
                    found_idx = i

            # Compare labels
            self.assertEqual(label_SA, multitask_microzinc.labels[found_idx]["SA"])
            self.assertEqual(label_logp, multitask_microzinc.labels[found_idx]["logp"])
            self.assertEqual(label_score, multitask_microzinc.labels[found_idx]["score"])

    def test_multitask_dataset_case_2(self):
        """Case: Different tasks, but with no intersection in the smiles (each task has a unique set of smiles)
        - Check that the total dataset has as much smiles as all tasks together
        - Check that, for each task, only the smiles related to that task have values, and ensure the value is what's expected from the initial DF
        """
        df = load_micro_zinc()  # Has about 1000 molecules

        # Choose non-overlapping smiles by choosing specific rows from the original dataframe.
        df_rows_SA = df.iloc[0:200]  # 200 data points
        df_rows_logp = df.iloc[200:400]  # 200 data points
        df_rows_score = df.iloc[400:750]  # 350 data points
        total_data_points = 750

        # Here we split the data according to the task we care about.
        df_micro_zinc_SA = df_rows_SA[["SMILES", "SA"]]
        df_micro_zinc_logp = df_rows_logp[["SMILES", "logp"]]
        df_micro_zinc_score = df_rows_score[["SMILES", "score"]]

        # We need to turn these dataframes into single-task datasets.
        # We don't need to do featurization yet.
        ds_micro_zinc_SA = SingleTaskDataset(
            smiles=df_micro_zinc_SA.loc[:, "SMILES"].tolist(), labels=df_micro_zinc_SA.loc[:, "SA"].tolist()
        )
        ds_micro_zinc_logp = SingleTaskDataset(
            smiles=df_micro_zinc_logp.loc[:, "SMILES"].tolist(),
            labels=df_micro_zinc_logp.loc[:, "logp"].tolist(),
        )
        ds_micro_zinc_score = SingleTaskDataset(
            smiles=df_micro_zinc_score.loc[:, "SMILES"].tolist(),
            labels=df_micro_zinc_score.loc[:, "score"].tolist(),
        )

        # Create the multitask dataset
        datasets_dict = {"SA": ds_micro_zinc_SA, "logp": ds_micro_zinc_logp, "score": ds_micro_zinc_score}
        multitask_microzinc = MultitaskDataset(
            datasets_dict, save_smiles_and_ids=True
        )  # Can optionally have features

        # The total dataset has as many molecules as there are smiles in all tasks put together
        self.assertEqual(total_data_points, multitask_microzinc.__len__())

        # For each task, only the smiles related to that task have values, and the value is what's expected from the initial DF.
        for idx in range(len(ds_micro_zinc_SA)):
            smiles = df[["SMILES"]].iloc[idx].values[0]

            task = "task"
            if idx in range(0, 200):
                task = "SA"
            elif idx in range(200, 400):
                task = "logp"
            elif idx in range(400, 750):
                task = "score"

            # Labels of that molecule
            label_SA = df[["SA"]].iloc[idx].values[0]
            label_logp = df[["logp"]].iloc[idx].values[0]
            label_score = df[["score"]].iloc[idx].values[0]

            # Search for that molecule in the multitask dataset
            mol_ids = smiles_to_unique_mol_ids([smiles])
            mol_id = mol_ids[0]
            found_idx = -1
            for i, id in enumerate(multitask_microzinc.mol_ids):
                if mol_id == id:
                    found_idx = i
            multitask_microzinc_labels = get_keys(multitask_microzinc.labels[found_idx])
            if task == "SA":
                self.assertEqual(label_SA, multitask_microzinc.labels[found_idx]["SA"])
                self.assertFalse("score" in multitask_microzinc_labels)
                self.assertFalse("logp" in multitask_microzinc_labels)
            elif task == "logp":
                self.assertEqual(label_logp, multitask_microzinc.labels[found_idx]["logp"])
                self.assertFalse("score" in multitask_microzinc_labels)
                self.assertFalse("SA" in multitask_microzinc_labels)
            elif task == "score":
                self.assertEqual(label_score, multitask_microzinc.labels[found_idx]["score"])
                self.assertFalse("SA" in multitask_microzinc_labels)
                self.assertFalse("logp" in multitask_microzinc_labels)

    def test_multitask_dataset_case_3(self):
        """Case: Different tasks, but with semi-intersection (some smiles unique per task, some intersect)
        - Check that the total dataset has as much smiles as the unique number of smiles.
        - Check that for each task, you retrieve the same smiles as expected from the initial DF
        """
        df_micro_zinc = load_micro_zinc()  # Has about 1000 molecules
        df = df_micro_zinc.iloc[0:5]

        # Choose OVERLAPPING smiles by choosing specific rows from the original dataframe. The tasks will not necessarily have unique smiles.
        df_rows_SA = df.iloc[0:3]
        df_rows_logp = df.iloc[1:4]
        df_rows_score = df.iloc[3:5]
        total_data_points = 5

        # Here we split the data according to the task we care about.
        df_micro_zinc_SA = df_rows_SA[["SMILES", "SA"]]
        df_micro_zinc_logp = df_rows_logp[["SMILES", "logp"]]
        df_micro_zinc_score = df_rows_score[["SMILES", "score"]]

        # We need to turn these dataframes into single-task datasets.
        # We don't need to do featurization yet.
        ds_micro_zinc_SA = SingleTaskDataset(
            smiles=df_micro_zinc_SA.loc[:, "SMILES"].tolist(), labels=df_micro_zinc_SA.loc[:, "SA"].tolist()
        )
        ds_micro_zinc_logp = SingleTaskDataset(
            smiles=df_micro_zinc_logp.loc[:, "SMILES"].tolist(),
            labels=df_micro_zinc_logp.loc[:, "logp"].tolist(),
        )
        ds_micro_zinc_score = SingleTaskDataset(
            smiles=df_micro_zinc_score.loc[:, "SMILES"].tolist(),
            labels=df_micro_zinc_score.loc[:, "score"].tolist(),
        )

        # Create the multitask dataset
        datasets_dict = {"SA": ds_micro_zinc_SA, "logp": ds_micro_zinc_logp, "score": ds_micro_zinc_score}
        multitask_microzinc = MultitaskDataset(
            datasets_dict, save_smiles_and_ids=True
        )  # Can optionally have features

        # The multitask dataset has as many molecules as there are unique smiles across the single task datasets.
        self.assertEqual(total_data_points, multitask_microzinc.__len__())


if __name__ == "__main__":
    ut.main()


# tests/test_dict_tensor.py

"""
Unit tests for the metrics and wrappers of graphium/utils/...
"""

from typing import Dict
from graphium.utils.dict_tensor import DictTensor
import torch
from torch import Tensor

import unittest as ut
from copy import deepcopy


class test_dict_tensor(ut.TestCase):
    def _assert_dict_tensor_equal(self, dict_ten1, dict_ten2, msg=""):
        self.assertSetEqual(set(dict_ten1.keys()), set(dict_ten2.keys()), msg=msg)
        for key, tensor1 in dict_ten1.items():
            tensor2 = dict_ten2[key]
            if isinstance(tensor1, Tensor):
                msg2 = msg + f"\ntensor1: \n{tensor1}\n\ntensor2: \n{tensor2}\n\n"
                self.assertTrue(torch.all(tensor1 == tensor2), msg=msg2)
            else:
                for ii in range(len(tensor1)):
                    msg2 = msg + f"\ntensor1: \n{tensor1[ii]}\n\ntensor2: \n{tensor2[ii]}\n\n"
                    self.assertTrue(torch.all(tensor1[ii] == tensor2[ii]), msg=msg2)

    def test_tensor_funcs(self):
        dict_ten = DictTensor(
            {
                "a": torch.randn(5, 6),
                "b": torch.randn(2, 3, 4),
                "c": torch.randn(1, 2, 3, 5),
            }
        )

        # Check `DictTensor.func` returns right value
        self.assertDictEqual(dict_ten.shape, {k: v.shape for k, v in dict_ten.items()})
        self._assert_dict_tensor_equal(
            dict_ten.pinverse(), {k: v.pinverse() for k, v in dict_ten.items()}, msg="dict_ten.pinverse()"
        )
        self._assert_dict_tensor_equal(
            dict_ten.sum(), {k: v.sum() for k, v in dict_ten.items()}, msg="dict_ten.sum()"
        )
        self._assert_dict_tensor_equal(
            dict_ten.sum(axis=0), {k: v.sum(axis=0) for k, v in dict_ten.items()}, msg="dict_ten.sum(axis=0)"
        )
        self._assert_dict_tensor_equal(
            dict_ten.sum(0), {k: v.sum(0) for k, v in dict_ten.items()}, msg="dict_ten.sum(0)"
        )
        self._assert_dict_tensor_equal(
            dict_ten.sum(axis=1), {k: v.sum(axis=1) for k, v in dict_ten.items()}, msg="dict_ten.sum(axis=1)"
        )
        self._assert_dict_tensor_equal(
            dict_ten.to(dtype=int),
            {k: v.to(dtype=int) for k, v in dict_ten.items()},
            msg="dict_ten.to(dtype=int)",
        )
        self._assert_dict_tensor_equal(
            dict_ten.max(axis=0),
            {k: v.max(axis=0) for k, v in dict_ten.items()},
            msg="dict_ten.max(axis=0)[0]",
        )
        self._assert_dict_tensor_equal(
            dict_ten.transpose(0, 1),
            {k: v.transpose(0, 1) for k, v in dict_ten.items()},
            msg="dict_ten.max(axis=0)[0]",
        )

        # Check `DictTensor.func` returns right type
        self.assertIsInstance(dict_ten.shape, dict)
        self.assertIsInstance(dict_ten.shape["a"], torch.Size)
        self.assertIsInstance(dict_ten.pinverse(), DictTensor)
        self.assertIsInstance(dict_ten.sum(), DictTensor)
        self.assertIsInstance(dict_ten.sum(axis=0), DictTensor)
        self.assertIsInstance(dict_ten.sum(0), DictTensor)
        self.assertIsInstance(dict_ten.sum(axis=1), DictTensor)
        self.assertIsInstance(dict_ten.to(dtype=int), DictTensor)
        self.assertIsInstance(dict_ten.max(axis=0), dict)
        self.assertIsInstance(dict_ten.max(axis=0)["a"][0], Tensor)

    def test_torch_funcs(self):
        dict_ten = DictTensor(
            {
                "a": torch.randn(5, 6),
                "b": torch.randn(2, 3, 4),
                "c": torch.randn(1, 2, 3, 5),
            }
        )

        # Check `torch.func(DictTensor, *)` returns right value
        self._assert_dict_tensor_equal(
            torch.pinverse(
                dict_ten,
            ),
            {
                k: torch.pinverse(
                    v,
                )
                for k, v in dict_ten.items()
            },
            msg="torch.pinverse(dict_ten, )",
        )
        self._assert_dict_tensor_equal(
            torch.sum(
                dict_ten,
            ),
            {
                k: torch.sum(
                    v,
                )
                for k, v in dict_ten.items()
            },
            msg="torch.sum(dict_ten, )",
        )
        self._assert_dict_tensor_equal(
            torch.sum(dict_ten, 0),
            {k: torch.sum(v, 0) for k, v in dict_ten.items()},
            msg="torch.sum(dict_ten, 0)",
        )
        self._assert_dict_tensor_equal(
            torch.sum(dict_ten, axis=0),
            {k: torch.sum(v, axis=0) for k, v in dict_ten.items()},
            msg="torch.sum(dict_ten, axis=0)",
        )
        self._assert_dict_tensor_equal(
            torch.sum(dict_ten, axis=1),
            {k: torch.sum(v, axis=1) for k, v in dict_ten.items()},
            msg="torch.sum(dict_ten, axis=1)",
        )
        self._assert_dict_tensor_equal(
            torch.max(dict_ten, axis=0),
            {k: torch.max(v, axis=0) for k, v in dict_ten.items()},
            msg="torch.max(dict_ten, axis=0)[0]",
        )

        # Check `torch.func(DictTensor, *)` returns right type
        self.assertIsInstance(
            torch.pinverse(
                dict_ten,
            ),
            DictTensor,
        )
        self.assertIsInstance(
            torch.sum(
                dict_ten,
            ),
            DictTensor,
        )
        self.assertIsInstance(torch.sum(dict_ten, axis=0), DictTensor)
        self.assertIsInstance(torch.sum(dict_ten, 0), DictTensor)
        self.assertIsInstance(torch.sum(dict_ten, axis=1), DictTensor)
        self.assertIsInstance(torch.max(dict_ten, axis=0), dict)
        self.assertIsInstance(torch.max(dict_ten, axis=0)["a"][0], Tensor)

    def test_apply(self):
        dict_ten = DictTensor(
            {
                "a": torch.randn(5, 6),
                "b": torch.randn(2, 3, 4),
                "c": torch.randn(1, 2, 3, 5),
            }
        )

        # Check `DictTensor.apply` function returns right value
        self._assert_dict_tensor_equal(
            dict_ten.apply(
                torch.pinverse,
            ),
            {
                k: torch.pinverse(
                    v,
                )
                for k, v in dict_ten.items()
            },
            msg="torch.pinverse(dict_ten, )",
        )
        self._assert_dict_tensor_equal(
            dict_ten.apply(
                torch.sum,
            ),
            {
                k: torch.sum(
                    v,
                )
                for k, v in dict_ten.items()
            },
            msg="torch.sum(dict_ten, )",
        )
        self._assert_dict_tensor_equal(
            dict_ten.apply(torch.sum, 0),
            {k: torch.sum(v, 0) for k, v in dict_ten.items()},
            msg="torch.sum(dict_ten, 0)",
        )
        self._assert_dict_tensor_equal(
            dict_ten.apply(torch.sum, axis=0),
            {k: torch.sum(v, axis=0) for k, v in dict_ten.items()},
            msg="torch.sum(dict_ten, axis=0)",
        )
        self._assert_dict_tensor_equal(
            dict_ten.apply(torch.sum, axis=1),
            {k: torch.sum(v, axis=1) for k, v in dict_ten.items()},
            msg="torch.sum(dict_ten, axis=1)",
        )
        self._assert_dict_tensor_equal(
            dict_ten.apply(torch.max, axis=0),
            {k: torch.max(v, axis=0) for k, v in dict_ten.items()},
            msg="torch.max(dict_ten, axis=0)[0]",
        )

        # Check `DictTensor.apply` function returns right type
        self.assertIsInstance(
            dict_ten.apply(
                torch.pinverse,
            ),
            DictTensor,
        )
        self.assertIsInstance(
            dict_ten.apply(
                torch.sum,
            ),
            DictTensor,
        )
        self.assertIsInstance(dict_ten.apply(torch.sum, axis=0), DictTensor)
        self.assertIsInstance(dict_ten.apply(torch.sum, 0), DictTensor)
        self.assertIsInstance(dict_ten.apply(torch.sum, axis=1), DictTensor)
        self.assertIsInstance(dict_ten.apply(torch.max, axis=0), dict)
        self.assertIsInstance(dict_ten.apply(torch.max, axis=0)["a"][0], Tensor)

    def test_tensor_operators(self):
        dict_ten = DictTensor(
            {
                "a": torch.randn(5, 6),
                "b": torch.randn(4, 5, 6),
                "c": torch.randn(1, 2, 5, 6),
            }
        )

        # Check product, division, sum, subtraction by a constant
        FACTOR = 0.5
        self._assert_dict_tensor_equal(
            dict_ten * FACTOR, {k: FACTOR * v for k, v in dict_ten.items()}, msg="DT * factor"
        )
        self._assert_dict_tensor_equal(
            FACTOR * dict_ten, {k: FACTOR * v for k, v in dict_ten.items()}, msg="factor * DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten + FACTOR, {k: FACTOR + v for k, v in dict_ten.items()}, msg="DT + factor"
        )
        self._assert_dict_tensor_equal(
            FACTOR + dict_ten, {k: FACTOR + v for k, v in dict_ten.items()}, msg="factor + DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten - FACTOR, {k: v - FACTOR for k, v in dict_ten.items()}, msg="DT - factor"
        )
        self._assert_dict_tensor_equal(
            FACTOR - dict_ten, {k: FACTOR - v for k, v in dict_ten.items()}, msg="factor - DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten / FACTOR, {k: v / FACTOR for k, v in dict_ten.items()}, msg="DT / factor"
        )
        self._assert_dict_tensor_equal(
            dict_ten // FACTOR, {k: v // FACTOR for k, v in dict_ten.items()}, msg="DT // factor"
        )

        # Check product, division, sum, subtraction by a Tensor of shape [6]
        FACTOR = torch.rand(6)
        self._assert_dict_tensor_equal(
            dict_ten * FACTOR, {k: FACTOR * v for k, v in dict_ten.items()}, msg="DT * tensor"
        )
        self._assert_dict_tensor_equal(
            FACTOR * dict_ten, {k: FACTOR * v for k, v in dict_ten.items()}, msg="tensor * DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten + FACTOR, {k: FACTOR + v for k, v in dict_ten.items()}, msg="DT + tensor"
        )
        self._assert_dict_tensor_equal(
            FACTOR + dict_ten, {k: FACTOR + v for k, v in dict_ten.items()}, msg="tensor + DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten - FACTOR, {k: v - FACTOR for k, v in dict_ten.items()}, msg="DT - tensor"
        )
        self._assert_dict_tensor_equal(
            FACTOR - dict_ten, {k: FACTOR - v for k, v in dict_ten.items()}, msg="tensor - DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten / FACTOR, {k: v / FACTOR for k, v in dict_ten.items()}, msg="DT / tensor"
        )
        self._assert_dict_tensor_equal(
            dict_ten // FACTOR, {k: v // FACTOR for k, v in dict_ten.items()}, msg="DT // tensor"
        )

        # Check product, division, sum, subtraction by a Tensor of shape [5, 6]
        FACTOR = torch.rand(5, 6)
        self._assert_dict_tensor_equal(
            dict_ten * FACTOR, {k: FACTOR * v for k, v in dict_ten.items()}, msg="DT * tensor"
        )
        self._assert_dict_tensor_equal(
            FACTOR * dict_ten, {k: FACTOR * v for k, v in dict_ten.items()}, msg="tensor * DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten + FACTOR, {k: FACTOR + v for k, v in dict_ten.items()}, msg="DT + tensor"
        )
        self._assert_dict_tensor_equal(
            FACTOR + dict_ten, {k: FACTOR + v for k, v in dict_ten.items()}, msg="tensor + DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten - FACTOR, {k: v - FACTOR for k, v in dict_ten.items()}, msg="DT - tensor"
        )
        self._assert_dict_tensor_equal(
            FACTOR - dict_ten, {k: FACTOR - v for k, v in dict_ten.items()}, msg="tensor - DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten / FACTOR, {k: v / FACTOR for k, v in dict_ten.items()}, msg="DT / tensor"
        )
        self._assert_dict_tensor_equal(
            dict_ten // FACTOR, {k: v // FACTOR for k, v in dict_ten.items()}, msg="DT // tensor"
        )

    def test_comparison_operators(self):
        dict_ten = DictTensor(
            {
                "a": torch.randn(5, 6),
                "b": torch.randn(4, 5, 6),
                "c": torch.randn(1, 2, 5, 6),
            }
        )
        dict_ten2 = deepcopy(dict_ten).abs()

        # Comparison operators with float
        self._assert_dict_tensor_equal(
            dict_ten > 0.2, {k: v > 0.2 for k, v in dict_ten.items()}, msg="DT > 0.2"
        )
        self._assert_dict_tensor_equal(
            dict_ten < 0.2, {k: v < 0.2 for k, v in dict_ten.items()}, msg="DT < 0.2"
        )
        self._assert_dict_tensor_equal(
            dict_ten >= 0.2, {k: v >= 0.2 for k, v in dict_ten.items()}, msg="DT >= 0.2"
        )
        self._assert_dict_tensor_equal(
            dict_ten <= 0.2, {k: v <= 0.2 for k, v in dict_ten.items()}, msg="DT <= 0.2"
        )
        self._assert_dict_tensor_equal(
            dict_ten.to(int) == 0, {k: v.to(int) == 0 for k, v in dict_ten.items()}, msg="DT.int == 0"
        )
        self._assert_dict_tensor_equal(
            dict_ten.to(int) != 0, {k: v.to(int) != 0 for k, v in dict_ten.items()}, msg="DT.int != 0"
        )

        # Comparison operators with Tensor
        tensor = torch.rand(5, 6)
        self._assert_dict_tensor_equal(
            dict_ten > tensor, {k: v > tensor for k, v in dict_ten.items()}, msg="DT > DT2"
        )
        self._assert_dict_tensor_equal(
            dict_ten < tensor, {k: v < tensor for k, v in dict_ten.items()}, msg="DT < DT2"
        )
        self._assert_dict_tensor_equal(
            dict_ten >= tensor, {k: v >= tensor for k, v in dict_ten.items()}, msg="DT >= DT2"
        )
        self._assert_dict_tensor_equal(
            dict_ten <= tensor, {k: v <= tensor for k, v in dict_ten.items()}, msg="DT <= DT2"
        )
        self._assert_dict_tensor_equal(
            dict_ten.to(int) == tensor,
            {k: v.to(int) == tensor for k, v in dict_ten.items()},
            msg="DT.int == DT2",
        )
        self._assert_dict_tensor_equal(
            dict_ten.to(int) != tensor,
            {k: v.to(int) != tensor for k, v in dict_ten.items()},
            msg="DT.int != DT2",
        )

        # Comparison operators with DictTensor
        self._assert_dict_tensor_equal(
            dict_ten > dict_ten2, {k: v > dict_ten2[k] for k, v in dict_ten.items()}, msg="DT > DT2"
        )
        self._assert_dict_tensor_equal(
            dict_ten < dict_ten2, {k: v < dict_ten2[k] for k, v in dict_ten.items()}, msg="DT < DT2"
        )
        self._assert_dict_tensor_equal(
            dict_ten >= dict_ten2, {k: v >= dict_ten2[k] for k, v in dict_ten.items()}, msg="DT >= DT2"
        )
        self._assert_dict_tensor_equal(
            dict_ten <= dict_ten2, {k: v <= dict_ten2[k] for k, v in dict_ten.items()}, msg="DT <= DT2"
        )
        self._assert_dict_tensor_equal(
            dict_ten.to(int) == dict_ten2,
            {k: v.to(int) == dict_ten2[k] for k, v in dict_ten.items()},
            msg="DT.int == DT2",
        )
        self._assert_dict_tensor_equal(
            dict_ten.to(int) != dict_ten2,
            {k: v.to(int) != dict_ten2[k] for k, v in dict_ten.items()},
            msg="DT.int != DT2",
        )

    def test_dict_functions(self):
        dict1 = {
            "a": torch.randn(5, 6),
            "b": torch.randn(4, 5, 6),
            "c": torch.randn(1, 2, 5, 6),
        }
        dict_ten1 = DictTensor(deepcopy(dict1))
        dict2 = {
            "c": torch.randn(1, 2, 6),
            "d": torch.randn(1, 2, 3, 4),
        }
        dict_ten2 = DictTensor(deepcopy(dict2))

        # Check update
        dict1_temp = deepcopy(dict1)
        dict_ten1_temp = deepcopy(dict_ten1)
        dict1_temp.update(dict2)
        truth = DictTensor(dict1_temp)
        dict_ten1_temp.update(dict_ten2)
        self._assert_dict_tensor_equal(dict_ten1_temp, truth)

        # Check pop
        dict1_temp = deepcopy(dict1)
        dict_ten1_temp = deepcopy(dict_ten1)
        truth_a = dict1_temp.pop("a")
        truth = DictTensor(dict1_temp)
        val_a = dict_ten1_temp.pop("a")
        self._assert_dict_tensor_equal(dict_ten1_temp, truth)
        self.assertTrue(torch.all(truth_a == val_a))

        # Check indexing
        dict1_temp = deepcopy(dict1)
        dict_ten1_temp = deepcopy(dict_ten1)
        truth_b = dict1_temp["b"]
        val_b = dict_ten1_temp["b"]
        self.assertTrue(torch.all(truth_b == val_b))

        # Check adding element
        dict1_temp = deepcopy(dict1)
        dict_ten1_temp = deepcopy(dict_ten1)
        dict1_temp["new"] = dict2["d"]
        truth = DictTensor(dict1_temp)
        dict_ten1_temp["new"] = dict_ten2["d"]
        self._assert_dict_tensor_equal(dict_ten1_temp, truth)


if __name__ == "__main__":
    ut.main()


# tests/test_featurizer.py

"""
Unit tests for the different datasets of graphium/features/featurizer.py
"""

import numpy as np
import unittest as ut
from copy import deepcopy
from rdkit import Chem
import datamol as dm

from graphium.features.featurizer import (
    get_mol_atomic_features_onehot,
    get_mol_atomic_features_float,
    get_mol_edge_features,
    mol_to_adj_and_features,
    mol_to_pyggraph,
)


class test_featurizer(ut.TestCase):
    smiles = [
        "C",
        "CC",
        "C1(C[N]CCC1)=O",
        "CC(C)CC1=CC=C(C=C1)C(C)C(=O)O",
        "OCCc1c(C)[n+](cs1)Cc2cnc(C)nc2N",
        "O1C=C[C@H]([C@H]1O2)c3c2cc(OC)c4c3OC(=O)C5=C4CCC(=O)5",
        "CC1(C2C1C(N(C2)C(=O)C(C(C)=B)NC(=O)C(F)(F)F)C(=O)NC(C(C3CCNC3=O)[Cl])C#N)C",
    ]

    smiles_noble = ["[He].[He]", "[He][He]", "[Kr][Kr]"]

    atomic_onehot_props = [
        "atomic-number",
        "valence",
        "degree",
        "implicit-valence",
        "hybridization",
        "chirality",
        "phase",
        "type",
        "group",
        "period",
    ]

    atomic_float_props = [
        "atomic-number",
        "mass",
        "valence",
        "implicit-valence",
        "hybridization",
        "chirality",
        "aromatic",
        "in-ring",
        "min-ring",
        "max-ring",
        "num-ring",
        "degree",
        "radical-electron",
        "formal-charge",
        "vdw-radius",
        "covalent-radius",
        "electronegativity",
        "ionization",
        "melting-point",
        "metal",
        "single-bond",
        "aromatic-bond",
        "double-bond",
        "triple-bond",
        "is-carbon",
        "group",
        "period",
    ]

    edge_props = [
        "bond-type-onehot",
        "bond-type-float",
        "stereo",
        "in-ring",
        "conjugated",
        "estimated-bond-length",
        "conformer-bond-length",
    ]

    def test_get_mol_atomic_features_onehot(self):
        props = deepcopy(self.atomic_onehot_props)
        bad_props = ["bob"]

        all_smiles = self.smiles + self.smiles_noble

        for s in all_smiles:
            err_msg = f"\n\tError for params:\n\t\tSMILES: {s}"
            mol = dm.to_mol(s)

            for ii in range(len(props)):
                this_props = props[:ii]
                err_msg2 = err_msg + f"\n\t\tprops: {this_props}"
                prop_dict = get_mol_atomic_features_onehot(mol, property_list=this_props)
                self.assertListEqual(list(prop_dict.keys()), this_props, msg=err_msg)
                for key, val in prop_dict.items():
                    err_msg3 = err_msg2 + f"\n\t\tkey: {key}"
                    self.assertEqual(val.shape[0], mol.GetNumAtoms(), msg=err_msg3)
                    self.assertGreater(val.shape[1], 1, msg=err_msg3)
                    self.assertTrue(np.all((val == 0) | (val == 1)), msg=err_msg3)

            with self.assertRaises(ValueError, msg=err_msg):
                get_mol_atomic_features_onehot(mol, property_list=bad_props)

    def test_get_mol_atomic_features_float(self):
        props = deepcopy(self.atomic_float_props)

        bad_props = ["bob"]

        all_smiles = self.smiles + self.smiles_noble
        for s in all_smiles:
            err_msg = f"\n\tError for params:\n\t\tSMILES: {s}"
            mol = dm.to_mol(s)

            for ii in range(len(props)):
                this_props = props[:ii]
                err_msg2 = err_msg + f"\n\t\tprops: {this_props}"
                prop_dict = get_mol_atomic_features_float(mol, property_list=this_props, mask_nan=None)
                self.assertListEqual(list(prop_dict.keys()), this_props, msg=err_msg)
                for key, val in prop_dict.items():
                    err_msg3 = err_msg2 + f"\n\t\tkey: {key}"
                    self.assertListEqual(list(val.shape), [mol.GetNumAtoms()], msg=err_msg3)

            with self.assertRaises(ValueError, msg=err_msg):
                get_mol_atomic_features_float(mol, property_list=bad_props)

    def test_get_mol_atomic_features_float_nan_mask(self):
        for s in self.smiles_noble:
            mol = dm.to_mol(s)

            # Nothing happens when `mask_nan = None`, nans are still in the property array
            prop_dict = get_mol_atomic_features_float(
                mol, property_list=self.atomic_float_props, mask_nan=None
            )
            prop_array = np.concatenate(list(prop_dict.values()), axis=0)
            nans = np.isnan(prop_array)

            # Capture a raised error when `mask_nan = "raise"`
            with self.assertRaises(ValueError):
                prop_dict = get_mol_atomic_features_float(
                    mol, property_list=self.atomic_float_props, mask_nan="raise"
                )

            # Not sure how to Capture a logged warning when `mask_nan = "warn"`
            # Here, I'm testing a behaviour similar to `mask_nan = None`
            prop_dict = get_mol_atomic_features_float(
                mol, property_list=self.atomic_float_props, mask_nan="warn"
            )
            prop_array = np.concatenate(list(prop_dict.values()), axis=0)
            self.assertEqual(len(self.atomic_float_props), len(prop_dict))
            self.assertTrue(any(np.isnan(prop_array)))

            # NaNs are replaced by `42` when `mask_nan=42`
            prop_dict = get_mol_atomic_features_float(mol, property_list=self.atomic_float_props, mask_nan=42)
            prop_array = np.concatenate(list(prop_dict.values()), axis=0)
            self.assertEqual(len(self.atomic_float_props), len(prop_dict))
            self.assertFalse(any(np.isnan(prop_array)))
            self.assertTrue(all(prop_array[nans] == 42))

    def test_get_mol_edge_features(self):
        props = deepcopy(self.edge_props)
        bad_props = ["bob"]

        all_smiles = self.smiles + self.smiles_noble
        for s in all_smiles:
            err_msg = f"\n\tError for params:\n\t\tSMILES: {s}"
            mol = dm.to_mol(s)
            for ii in range(len(props)):
                this_props = props[: ii + 1]
                err_msg2 = err_msg + f"\n\t\tprops: {this_props}"
                prop_dict = get_mol_edge_features(mol, property_list=this_props)
                self.assertListEqual(list(prop_dict.keys()), this_props, msg=err_msg)
                for key, val in prop_dict.items():
                    err_msg3 = err_msg2 + f"\n\t\tkey: {key}"
                    self.assertEqual(val.shape[0], mol.GetNumBonds(), msg=err_msg3)

            if mol.GetNumBonds() > 0:
                with self.assertRaises(ValueError, msg=err_msg):
                    get_mol_edge_features(mol, property_list=bad_props)

    def test_mol_to_adj_and_features(self):
        np.random.seed(42)

        for s in self.smiles:
            err_msg = f"\n\tError for params:\n\t\tSMILES: {s}"
            mol = dm.to_mol(s)
            mol_Hs = Chem.AddHs(mol)  # type: ignore
            mol_No_Hs = Chem.RemoveHs(mol)  # type: ignore

            for explicit_H in [True, False]:
                this_mol = mol_Hs if explicit_H else mol_No_Hs
                for ii in np.arange(0, 5, 0.2):
                    num_props = int(round(ii))
                    err_msg2 = err_msg + f"\n\t\texplicit_H: {explicit_H}\n\t\tii: {ii}"

                    adj, ndata, edata, _, _ = mol_to_adj_and_features(
                        mol=mol,
                        atom_property_list_onehot=np.random.choice(
                            self.atomic_onehot_props, size=num_props, replace=False
                        ),
                        atom_property_list_float=np.random.choice(
                            self.atomic_float_props, size=num_props, replace=False
                        ),
                        edge_property_list=np.random.choice(self.edge_props, size=num_props, replace=False),
                        add_self_loop=False,
                        explicit_H=explicit_H,
                        use_bonds_weights=False,
                    )

                    self.assertEqual(adj.shape[0], this_mol.GetNumAtoms(), msg=err_msg2)
                    if num_props > 0:
                        self.assertEqual(ndata.shape[0], this_mol.GetNumAtoms(), msg=err_msg2)
                        if this_mol.GetNumBonds() > 0:
                            self.assertEqual(edata.shape[0], this_mol.GetNumBonds(), msg=err_msg2)
                            self.assertGreaterEqual(edata.shape[1], num_props, msg=err_msg2)
                        self.assertGreaterEqual(ndata.shape[1], num_props, msg=err_msg2)

    def test_mol_to_pyggraph(self):
        np.random.seed(42)

        for s in self.smiles:
            err_msg = f"\n\tError for params:\n\t\tSMILES: {s}"
            mol = dm.to_mol(s)
            mol_Hs = Chem.AddHs(mol)  # type: ignore
            mol_No_Hs = Chem.RemoveHs(mol)  # type: ignore

            graph = mol_to_pyggraph(
                mol=mol,
                atom_property_list_onehot=[],
                atom_property_list_float=["atomic-number"],
                edge_property_list=["bond-type-float"],
                add_self_loop=False,
                explicit_H=False,
                use_bonds_weights=False,
                on_error="raise",
            )

            # Check the number of nodes and edges
            self.assertListEqual(list(graph["feat"].shape), [mol.GetNumAtoms(), 1], msg=err_msg)
            self.assertListEqual(list(graph["edge_feat"].shape), [2 * mol.GetNumBonds(), 1], msg=err_msg)

            # Check the node features
            feat = graph["feat"].to_dense().numpy() * 5 + 6  # Undo the scaling
            atom_nums = np.asarray([atom.GetAtomicNum() for atom in mol.GetAtoms()])
            np.testing.assert_array_almost_equal(feat[:, 0], atom_nums, decimal=5, err_msg=err_msg)

            # Check the edge features
            edge_feat = graph["edge_feat"].to_dense().numpy()
            bond_types = np.asarray([bond.GetBondTypeAsDouble() for bond in mol.GetBonds()]).repeat(2)
            np.testing.assert_array_almost_equal(edge_feat[:, 0], bond_types, decimal=5, err_msg=err_msg)

            # Check the edge indices
            if mol.GetNumBonds() > 0:
                edge_index = graph["edge_index"].to_dense().numpy()
                true_edge_index = []
                for bond in mol.GetBonds():
                    true_edge_index.append([bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()])
                    true_edge_index.append([bond.GetEndAtomIdx(), bond.GetBeginAtomIdx()])
                true_edge_index = np.asarray(true_edge_index).T
                np.testing.assert_array_equal(edge_index, true_edge_index, err_msg=err_msg)

            # Loop over many possible combinations of properties
            for explicit_H in [True, False]:
                this_mol = mol_Hs if explicit_H else mol_No_Hs
                for ii in np.arange(0, 5, 0.2):
                    num_props = int(round(ii))
                    err_msg2 = err_msg + f"\n\t\texplicit_H: {explicit_H}\n\t\tii: {ii}"

                    graph = mol_to_pyggraph(
                        mol=mol,
                        atom_property_list_onehot=np.random.choice(
                            self.atomic_onehot_props, size=num_props, replace=False
                        ),
                        atom_property_list_float=np.random.choice(
                            self.atomic_float_props, size=num_props, replace=False
                        ),
                        edge_property_list=np.random.choice(self.edge_props, size=num_props, replace=False),
                        add_self_loop=False,
                        explicit_H=explicit_H,
                        use_bonds_weights=False,
                        on_error="raise",
                    )

                    self.assertEqual(graph.num_nodes, this_mol.GetNumAtoms(), msg=err_msg2)
                    self.assertEqual(graph.num_edges, 2 * this_mol.GetNumBonds(), msg=err_msg2)
                    if num_props > 0:
                        ndata = graph["feat"]
                        edata = graph["edge_feat"]
                        self.assertEqual(ndata.shape[0], this_mol.GetNumAtoms(), msg=err_msg2)
                        self.assertEqual(edata.shape[0], 2 * this_mol.GetNumBonds(), msg=err_msg2)
                        self.assertGreaterEqual(ndata.shape[1], num_props, msg=err_msg2)
                        self.assertGreaterEqual(edata.shape[1], num_props, msg=err_msg2)


if __name__ == "__main__":
    ut.main()


# tests/test_ipu_dataloader.py

# General imports
import yaml
import unittest as ut
import numpy as np
from copy import deepcopy
from warnings import warn
from lightning import Trainer, LightningModule
from lightning_graphcore import IPUStrategy
from functools import partial
import pytest
import torch
from torch.utils.data.dataloader import default_collate

# Current library imports
from graphium.config._loader import (
    load_datamodule,
    load_metrics,
    load_architecture,
    load_accelerator,
    load_predictor,
    load_trainer,
)
from graphium.utils.safe_run import SafeRun


def random_packing(num_nodes, batch_size):
    ipu_batch_size = int(len(num_nodes) / batch_size)
    indices = np.arange(len(num_nodes))
    np.random.shuffle(indices)
    indices = np.reshape(indices, (ipu_batch_size, batch_size)).tolist()
    return indices


def global_batch_collator(batch_size, batches):
    packs = []
    for pack_idx in range(0, len(batches), batch_size):
        packs.append(default_collate(batches[pack_idx : pack_idx + batch_size]))
    global_batch = default_collate(packs)
    global_batch = (global_batch[0], tuple(global_batch[1]))
    return global_batch


@pytest.mark.ipu
class test_DataLoading(ut.TestCase):
    class TestSimpleLightning(LightningModule):
        # Create a basic Ligthning for testing the batch sizes
        def __init__(self, batch_size, node_feat_size, edge_feat_size, num_batch) -> None:
            super().__init__()
            self.batch_size = batch_size
            self.node_feat_size = node_feat_size
            self.edge_feat_size = edge_feat_size
            self.layer = torch.nn.Linear(node_feat_size, 1)
            self.loss_fn = torch.nn.L1Loss()
            self.num_batch = num_batch

        def validation_step(self, batch, batch_idx):
            self.assert_shapes(batch, batch_idx, "val")
            loss = self.forward(batch)
            return loss

        def training_step(self, batch, batch_idx):
            self.assert_shapes(batch, batch_idx, "train")
            loss = self.forward(batch)
            return loss

        def forward(self, batch):
            out = self.layer(batch[1][0]).squeeze(-1)
            loss = self.loss_fn(out, batch[0])
            return loss

        def assert_shapes(self, batch, batch_idx, step):
            # Test the shape of the labels
            this_shape = list(batch[0].shape)
            true_shape = [1, self.batch_size]
            assert (
                this_shape == true_shape
            ), f"Shape of the labels is `{this_shape}` but should be {true_shape}"

            # Test the shape of the first feature
            this_shape = list(batch[1][0].shape)
            true_shape = [1, self.batch_size, self.node_feat_size]
            assert (
                this_shape == true_shape
            ), f"Shape of the feature 0 is `{this_shape}` but should be {true_shape}"

            # Test the shape of the second feature
            this_shape = list(batch[1][1].shape)
            true_shape = [1, self.batch_size, self.edge_feat_size]
            assert (
                this_shape == true_shape
            ), f"Shape of the feature 0 is `{this_shape}` but should be {true_shape}"

        def configure_optimizers(self):
            return torch.optim.Adam(self.parameters(), lr=1e-3)

    class TestDataset(torch.utils.data.Dataset):
        # Create a simple dataset for testing the Lightning integration
        def __init__(self, labels, node_features, edge_features):
            self.labels = labels
            self.node_features = node_features
            self.edge_features = edge_features

        def __len__(self):
            return len(self.labels)

        def __getitem__(self, idx):
            # [label, [feat1, feat2]]
            return [self.labels[idx], [self.node_features[idx], self.edge_features[idx]]]

    @pytest.mark.skip
    def test_poptorch_simple_deviceiterations_gradient_accumulation(self):
        """
        Test a simple version of the device-iterations and gradient accumulation
        to make sure that the dataloader and models handle them correcly.
        """

        # Run this test only if poptorch is available
        try:
            import poptorch
        except Exception as e:
            warn(f"Skipping this test because poptorch is not available.\n{e}")
            return

        # Initialize constants
        gradient_accumulation = 2
        device_iterations = 3
        batch_size = 5
        num_replicate = 7
        node_feat_size = 11
        edge_feat_size = 13

        # Initialize the batch info and poptorch options
        opts = poptorch.Options()
        opts.useIpuModel(True)
        opts.deviceIterations(device_iterations)
        training_opts = deepcopy(opts)
        training_opts.Training.gradientAccumulation(gradient_accumulation)
        inference_opts = deepcopy(opts)

        # Initialize the dataset
        num_batch = device_iterations * gradient_accumulation * num_replicate
        data_size = num_batch * batch_size
        dataset = self.TestDataset(
            labels=np.random.rand(data_size).astype(np.float32),
            node_features=[np.random.rand(node_feat_size).astype(np.float32) for ii in range(data_size)],
            edge_features=[np.random.rand(edge_feat_size).astype(np.float32) for ii in range(data_size)],
        )

        # Initialize the dataloader
        train_dataloader = poptorch.DataLoader(
            options=training_opts,
            dataset=deepcopy(dataset),
            batch_size=batch_size,
            collate_fn=partial(global_batch_collator, batch_size),
        )

        val_dataloader = poptorch.DataLoader(
            options=inference_opts,
            dataset=deepcopy(dataset),
            batch_size=batch_size,
            collate_fn=partial(global_batch_collator, batch_size),
        )

        # Build the model, and run it on IPU
        model = self.TestSimpleLightning(batch_size, node_feat_size, edge_feat_size, num_batch)
        strategy = IPUStrategy(training_opts=training_opts, inference_opts=inference_opts)
        trainer = Trainer(
            logger=False,
            enable_checkpointing=False,
            max_epochs=2,
            strategy=strategy,
            num_sanity_val_steps=0,
            accelerator="ipu",
            devices=1,
        )
        trainer.fit(model=model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader)

    @pytest.mark.skip
    def test_poptorch_graphium_deviceiterations_gradient_accumulation(self):
        """
        Test the device-iterations and gradient accumulation in a way
        that is very similar to the Graphium code
        to make sure that the dataloader and models handle them correcly.
        """

        try:
            import poptorch
        except Exception as e:
            warn(f"Skipping this test because poptorch is not available.\n{e}")
            return

        # Simplified testing config - reflecting the toymix requirements
        # CONFIG_FILE = "tests/config_test_ipu_dataloader.yaml"
        CONFIG_FILE = "tests/config_test_ipu_dataloader_multitask.yaml"
        with open(CONFIG_FILE, "r") as f:
            cfg = yaml.safe_load(f)
        cfg, accelerator = load_accelerator(cfg)

        # Load the datamodule, and prepare the data
        datamodule = load_datamodule(cfg, accelerator_type=accelerator)
        datamodule.prepare_data()
        metrics = load_metrics(cfg)
        model_class, model_kwargs = load_architecture(cfg, in_dims=datamodule.in_dims)
        # datamodule.setup()
        predictor = load_predictor(
            cfg, model_class, model_kwargs, metrics, accelerator, datamodule.task_norms
        )
        trainer = load_trainer(cfg, "test", accelerator, "date_time_suffix")
        # Run the model training
        with SafeRun(name="TRAINING", raise_error=cfg["constants"]["raise_train_error"], verbose=True):
            trainer.fit(model=predictor, datamodule=datamodule)


if __name__ == "__main__":
    ut.main()


# tests/test_ipu_losses.py

import unittest as ut
import torch
from torch.nn import BCELoss, MSELoss, L1Loss, BCEWithLogitsLoss
from copy import deepcopy
import pytest

from graphium.ipu.ipu_losses import BCELossIPU, MSELossIPU, L1LossIPU, BCEWithLogitsLossIPU, HybridCELossIPU
from graphium.trainer.losses import HybridCELoss


@pytest.mark.ipu
class test_Losses(ut.TestCase):
    torch.manual_seed(42)
    preds = torch.rand((100, 10), dtype=torch.float32)
    target = torch.rand((100, 10), dtype=torch.float32)

    th = 0.7
    nan_th = 0.2
    preds_greater = preds > th
    target_greater = (target > th).to(torch.float32)
    target_greater_nan = deepcopy(target_greater)
    is_nan = target < nan_th
    target_greater_nan[target < nan_th] = torch.nan
    target_nan = deepcopy(target)
    target_nan[target < nan_th] = torch.nan

    def test_bce(self):
        preds = deepcopy(self.preds)
        target = deepcopy(self.target_greater)
        target_nan = deepcopy(self.target_greater_nan)

        # Regular loss
        loss_true = BCELoss()(preds, target)
        loss_ipu = BCELossIPU()(preds, target)
        self.assertFalse(loss_true.isnan(), "Regular BCELoss is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular BCELoss is different"
        )

        # Weighted loss
        weight = torch.rand(preds.shape[1], dtype=torch.float32)
        loss_true = BCELoss(weight=weight)(preds, target)
        loss_ipu = BCELossIPU(weight=weight)(preds, target)
        self.assertFalse(loss_true.isnan(), "Regular BCELoss is NaN")
        self.assertAlmostEqual(loss_true.item(), loss_ipu.item(), msg="Weighted BCELoss is different")

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        loss_true = BCELoss()(preds[not_nan], target[not_nan])
        loss_ipu = BCELossIPU()(preds, target_nan)
        self.assertFalse(loss_true.isnan(), "Regular BCELoss with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Regular BCELossIPU with target_nan is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular BCELoss with NaN is different"
        )

        # Weighted loss with NaNs in target
        not_nan = ~target_nan.isnan()
        weight = torch.rand(preds.shape, dtype=torch.float32)
        loss_true = BCELoss(weight=weight[not_nan])(preds[not_nan], target_nan[not_nan])
        loss_ipu = BCELossIPU(weight=weight)(preds, target_nan)
        self.assertFalse(loss_true.isnan(), "Weighted BCELoss with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Weighted BCELossIPU with target_nan is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Weighted BCELoss with NaN is different"
        )

    def test_mse(self):
        preds = deepcopy(self.preds)
        target = deepcopy(self.target)
        target_nan = deepcopy(self.target_nan)

        # Regular loss
        loss_true = MSELoss()(preds, target)
        loss_ipu = MSELossIPU()(preds, target)
        self.assertFalse(loss_true.isnan(), "Regular MSELoss is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular MSELoss is different"
        )

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        loss_true = MSELoss()(preds[not_nan], target[not_nan])
        loss_ipu = MSELossIPU()(preds, target_nan)
        self.assertFalse(loss_true.isnan(), "Regular MSELoss with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Regular MSELossIPU with target_nan is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular MSELoss with NaN is different"
        )

    def test_l1(self):
        preds = deepcopy(self.preds)
        target = deepcopy(self.target)
        target_nan = deepcopy(self.target_nan)

        # Regular loss
        loss_true = L1Loss()(preds, target)
        loss_ipu = L1LossIPU()(preds, target)
        self.assertFalse(loss_true.isnan(), "Regular MAELoss is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular MAELoss is different"
        )

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        loss_true = L1Loss()(preds[not_nan], target[not_nan])
        loss_ipu = L1LossIPU()(preds, target_nan)
        self.assertFalse(loss_true.isnan(), "Regular MAELoss with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Regular MAELossIPU with target_nan is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular MAELoss with NaN is different"
        )

    def test_bce_logits(self):
        preds = deepcopy(self.preds)
        target = deepcopy(self.target_greater)
        target_nan = deepcopy(self.target_greater_nan)

        # Regular loss
        loss_true = BCEWithLogitsLoss()(preds, target)
        loss_ipu = BCEWithLogitsLossIPU()(preds, target)
        self.assertFalse(loss_true.isnan(), "Regular BCEWithLogitsLoss is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular BCEWithLogitsLoss is different"
        )

        # Weighted loss
        weight = torch.rand(preds.shape[1], dtype=torch.float32)
        loss_true = BCEWithLogitsLoss(weight=weight)(preds, target)
        loss_ipu = BCEWithLogitsLossIPU(weight=weight)(preds, target)
        self.assertFalse(loss_true.isnan(), "Regular BCEWithLogitsLoss is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), msg="Weighted BCEWithLogitsLoss is different"
        )

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        loss_true = BCEWithLogitsLoss()(preds[not_nan], target[not_nan])
        loss_ipu = BCEWithLogitsLossIPU()(preds, target_nan)
        self.assertFalse(loss_true.isnan(), "Regular test_bce_logits with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Regular test_bce_logits with target_nan is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular BCELoss with NaN is different"
        )

        # Weighted loss with NaNs in target
        not_nan = ~target_nan.isnan()
        weight = torch.rand(preds.shape, dtype=torch.float32)
        loss_true = BCEWithLogitsLoss(weight=weight[not_nan])(preds[not_nan], target_nan[not_nan])
        loss_ipu = BCEWithLogitsLossIPU(weight=weight)(preds, target_nan)
        self.assertFalse(loss_true.isnan(), "Weighted test_bce_logits with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Weighted test_bce_logits with target_nan is NaN")
        self.assertAlmostEqual(
            loss_true.item(),
            loss_ipu.item(),
            places=6,
            msg="Weighted BCEWithLogitsLoss with NaN is different",
        )


# tests/test_ipu_metrics.py

import unittest as ut
import torch
from torchmetrics.functional import (
    auroc,
    average_precision,
    precision,
    accuracy,
    recall,
    pearson_corrcoef,
    spearman_corrcoef,
    r2_score,
    f1_score,
    fbeta_score,
    mean_squared_error,
    mean_absolute_error,
)
from copy import deepcopy
import pytest

from graphium.ipu.ipu_metrics import (
    auroc_ipu,
    average_precision_ipu,
    precision_ipu,
    accuracy_ipu,
    recall_ipu,
    pearson_ipu,
    spearman_ipu,
    r2_score_ipu,
    f1_score_ipu,
    fbeta_score_ipu,
    mean_squared_error_ipu,
    mean_absolute_error_ipu,
)


@pytest.mark.ipu
class test_Metrics(ut.TestCase):
    torch.manual_seed(42)
    preds = torch.rand((100, 10), dtype=torch.float32)
    target = torch.rand((100, 10), dtype=torch.float32)

    th = 0.7
    nan_th = 0.2
    preds_greater = preds > th
    target_greater = (target > th).to(torch.float32)
    target_greater_nan = deepcopy(target_greater)
    is_nan = target < nan_th
    target_greater_nan[target < nan_th] = torch.nan
    target_nan = deepcopy(target)
    target_nan[target < nan_th] = torch.nan

    def test_auroc(self):
        preds = deepcopy(self.preds)[:, 0]
        target = deepcopy(self.target)[:, 0]
        target_nan = deepcopy(self.target_nan)[:, 0]

        target[target < 0.5] = 0
        target[target >= 0.5] = 1

        target_nan[target_nan < 0.5] = 0
        target_nan[target_nan >= 0.5] = 1

        # Regular loss
        score_true = auroc(preds, target.to(int))
        score_ipu = auroc_ipu(preds, target)
        self.assertFalse(score_true.isnan(), "Regular AUROC score is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Regular AUROC score is different"
        )

        # Weighted loss (As in BCE)
        sample_weights = torch.rand(preds.shape[0], dtype=torch.float32)
        score_true = auroc(preds, target.to(int), sample_weights=sample_weights)
        score_ipu = auroc_ipu(preds, target, sample_weights=sample_weights)
        self.assertFalse(score_true.isnan(), "Regular AUROC score is NaN")
        self.assertAlmostEqual(score_true.item(), score_ipu.item(), msg="Weighted AUROC score is different")

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = auroc(preds[not_nan], target[not_nan].to(int))
        score_ipu = auroc_ipu(preds, target_nan)
        self.assertFalse(score_true.isnan(), "Regular AUROC score with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Regular AUROCIPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Regular AUROC score with NaN is different"
        )

        # Weighted loss with NaNs in target (As in BCE)
        not_nan = ~target_nan.isnan()
        sample_weights = torch.rand(preds.shape, dtype=torch.float32)
        loss_true = auroc(preds[not_nan], target_nan[not_nan].to(int), sample_weights=sample_weights[not_nan])
        loss_ipu = auroc_ipu(preds, target_nan, sample_weights=sample_weights)
        self.assertFalse(loss_true.isnan(), "Weighted AUROC score with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Weighted AUROC IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            # AssertionError: 0.6603766679763794 != 0.6234951615333557 within 2 places
            loss_true.item(),
            loss_ipu.item(),
            places=6,
            msg="Weighted AUROC with NaN is different",
        )

    def test_average_precision(self):  # TODO: Make work with multi-class
        preds = deepcopy(self.preds)[:, 0]
        target = deepcopy(self.target)[:, 0]
        target_nan = deepcopy(self.target_nan)[:, 0]

        target[target < 0.5] = 0
        target[target >= 0.5] = 1

        target_nan[target_nan < 0.5] = 0
        target_nan[target_nan >= 0.5] = 1

        # Regular loss
        score_true = average_precision(preds, target.to(int), task="binary")
        score_ipu = average_precision_ipu(preds, target.to(int), task="binary")
        self.assertFalse(score_true.isnan(), "Regular Average Precision is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Regular Average Precision is different"
        )

        # Regular average precision with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = average_precision(preds[not_nan], target[not_nan].to(int), task="binary")
        score_ipu = average_precision_ipu(preds, target_nan, task="binary")
        self.assertFalse(score_true.isnan(), "Regular Average Precision with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Regular Average Precision IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Regular Average Precision with NaN is different",
        )

    def test_precision(self):
        preds = deepcopy(self.preds)[:, :4]
        target = deepcopy(self.target)[:, 0]
        t = deepcopy(target)

        target[t < 0.4] = 0
        target[(t >= 0.4) & (t < 0.6)] = 1
        target[(t >= 0.6) & (t < 0.8)] = 2
        target[(t >= 0.8)] = 3

        target_nan = deepcopy(target)
        target_nan[self.is_nan[:, 0]] = float("nan")
        target_nan_bin = deepcopy(target_nan)
        target_nan_bin[target_nan > 0] = 1

        # Micro precision binary
        score_true = precision(preds[:, 0], target.to(int) > 0, average="micro")
        score_ipu = precision_ipu(preds[:, 0], target > 0, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Precision binary is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Precision binary is different"
        )

        # Micro precision binary with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = precision(preds[:, 0][not_nan], target_nan_bin[not_nan].to(int), average="micro")
        score_ipu = precision_ipu(preds[:, 0], target_nan_bin, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Precision binary with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro Precision binary IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Precision with NaN is different"
        )

        # Micro precision
        score_true = precision(preds, target.to(int), average="micro")
        score_ipu = precision_ipu(preds, target, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Precision is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Precision is different"
        )

        # Micro precision with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = precision(preds[not_nan], target[not_nan].to(int), average="micro")
        score_ipu = precision_ipu(preds, target_nan, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Precision with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro Precision IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Precision with NaN is different"
        )

        # Macro precision
        score_true = precision(preds, target.to(int), average="macro", num_classes=4)
        score_ipu = precision_ipu(preds, target, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro Precision is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Macro Precision is different"
        )

        # Macro precision multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = precision(preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4)
        score_ipu = precision_ipu(preds, target_nan, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro Precision multiclass with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Macro Precision multiclass IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Macro Precision multiclass with NaN is different",
        )

        # Macro precision multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = precision(preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4)
        score_ipu = precision_ipu(preds, target_nan, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro Precision multiclass with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Macro Precision multiclass IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Macro Precision multiclass with NaN is different",
        )

        # Weighted precision multiclass
        score_true = precision(preds, target.to(int), average="weighted", num_classes=4)
        score_ipu = precision_ipu(preds, target, average="weighted", num_classes=4)
        self.assertFalse(score_true.isnan(), "Weighted Precision multiclass is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Weighted Precision multiclass is different"
        )

        # Weighted precision multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = precision(preds[not_nan], target[not_nan].to(int), average="weighted", num_classes=4)
        score_ipu = precision_ipu(preds, target_nan, average="weighted", num_classes=4)
        self.assertFalse(score_true.isnan(), "Weighted Precision multiclass with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Weighted Precision multiclass IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Regular Average Precision multiclass with NaN is different",
        )

    def test_accuracy(self):
        preds = deepcopy(self.preds)[:, :4]
        target = deepcopy(self.target)[:, 0]
        t = deepcopy(target)

        target[t < 0.4] = 0
        target[(t >= 0.4) & (t < 0.6)] = 1
        target[(t >= 0.6) & (t < 0.8)] = 2
        target[(t >= 0.8)] = 3

        target_nan = deepcopy(target)
        target_nan[self.is_nan[:, 0]] = float("nan")
        target_nan_bin = deepcopy(target_nan)
        target_nan_bin[target_nan > 0] = 1

        # Micro accuracy binary
        score_true = accuracy(preds[:, 0], target.to(int) > 0, average="micro")
        score_ipu = accuracy_ipu(preds[:, 0], target > 0, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Accuracy binary is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Accuracy binary is different"
        )

        # Micro accuracy binary with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = accuracy(preds[:, 0][not_nan], target_nan_bin[not_nan].to(int), average="micro")
        score_ipu = accuracy_ipu(preds[:, 0], target_nan_bin, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Accuracy binary with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro Accuracy binary IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Accuracy with NaN is different"
        )

        # Micro accuracy
        score_true = accuracy(preds, target.to(int), average="micro")
        score_ipu = accuracy_ipu(preds, target, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Accuracy is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Accuracy is different"
        )

        # Micro accuracy with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = accuracy(preds[not_nan], target[not_nan].to(int), average="micro")
        score_ipu = accuracy_ipu(preds, target_nan, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Accuracy with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro Accuracy IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Accuracy with NaN is different"
        )

        # Macro accuracy
        score_true = accuracy(preds, target.to(int), average="macro", num_classes=4)
        score_ipu = accuracy_ipu(preds, target, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro Accuracy is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Macro Accuracy is different"
        )

        # Macro accuracy with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = accuracy(preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4)
        score_ipu = accuracy_ipu(preds, target_nan, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro Accuracy with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Macro Accuracy IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Macro Accuracy with NaN is different"
        )

        # Weighted accuracy
        score_true = accuracy(preds, target.to(int), average="weighted", num_classes=4)
        score_ipu = accuracy_ipu(preds, target, average="weighted", num_classes=4)
        self.assertFalse(score_true.isnan(), "Weighted Accuracy is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Weighted Accuracy is different"
        )

        # Weighted accuracy with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = accuracy(preds[not_nan], target[not_nan].to(int), average="weighted", num_classes=4)
        score_ipu = accuracy_ipu(preds, target_nan, average="weighted", num_classes=4)
        self.assertFalse(score_true.isnan(), "Weighted Accuracy with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Weighted Accuracy IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Regular Accuracy with NaN is different"
        )

    def test_recall(self):
        preds = deepcopy(self.preds)[:, :4]
        target = deepcopy(self.target)[:, 0]
        t = deepcopy(target)

        target[t < 0.4] = 0
        target[(t >= 0.4) & (t < 0.6)] = 1
        target[(t >= 0.6) & (t < 0.8)] = 2
        target[(t >= 0.8)] = 3

        target_nan = deepcopy(target)
        target_nan[self.is_nan[:, 0]] = float("nan")
        target_nan_bin = deepcopy(target_nan)
        target_nan_bin[target_nan > 0] = 1

        # Micro recall binary
        score_true = recall(preds[:, 0], target.to(int) > 0, average="micro")
        score_ipu = recall_ipu(preds[:, 0], target > 0, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Recall binary is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Recall binary is different"
        )

        # Micro recall binary with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = recall(preds[:, 0][not_nan], target_nan_bin[not_nan].to(int), average="micro")
        score_ipu = recall_ipu(preds[:, 0], target_nan_bin, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Recall binary with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro Recall binary IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Recall binary with NaN is different"
        )

        # Micro recall
        score_true = recall(preds, target.to(int), average="micro")
        score_ipu = recall_ipu(preds, target, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Recall is NaN")
        self.assertAlmostEqual(score_true.item(), score_ipu.item(), places=6, msg="Micro Recall is different")

        # Micro recall with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = recall(preds[not_nan], target[not_nan].to(int), average="micro")
        score_ipu = recall_ipu(preds, target_nan, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Recall with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro Recall IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Recall with NaN is different"
        )

        # Macro recall multiclass
        score_true = recall(preds, target.to(int), average="macro", num_classes=4)
        score_ipu = recall_ipu(preds, target, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro Recall is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Macro Recall multiclass is different"
        )

        # Macro recall multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = recall(preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4)
        score_ipu = recall_ipu(preds, target_nan, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro Recall multiclass with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Macro Recall multiclass IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Macro Recall multiclass with NaN is different"
        )

        # Weighted recallmulticlass
        score_true = recall(preds, target.to(int), average="weighted", num_classes=4)
        score_ipu = recall_ipu(preds, target, average="weighted", num_classes=4)
        self.assertFalse(score_true.isnan(), "Weighted Recall is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Weighted Recall is different"
        )

        # Weighted recall multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = recall(preds[not_nan], target[not_nan].to(int), average="weighted", num_classes=4)
        score_ipu = recall_ipu(preds, target_nan, average="weighted", num_classes=4)
        self.assertFalse(score_true.isnan(), "Weighted Recall multiclass with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Weighted Recall multiclass IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Regular Recall multiclass with NaN is different",
        )

    def test_pearsonr(self):
        preds = deepcopy(self.preds)[:, 0]
        target = deepcopy(self.target)[:, 0] + preds
        target_nan = deepcopy(target)
        target_nan[self.is_nan[:, 0]] = float("nan")

        # Regular loss
        score_true = pearson_corrcoef(preds, target)
        score_ipu = pearson_ipu(preds, target)
        self.assertFalse(score_true.isnan(), "Pearson is NaN")
        self.assertAlmostEqual(score_true.item(), score_ipu.item(), places=4, msg="Pearson is different")

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = pearson_corrcoef(preds[not_nan], target[not_nan])
        score_ipu = pearson_ipu(preds, target_nan)
        self.assertFalse(score_true.isnan(), "Regular PearsonR with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "IPU PearsonR score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=4, msg="Pearson with NaN is different"
        )

    def test_spearmanr(self):
        preds = deepcopy(self.preds)[:, 0]
        target = deepcopy(self.target)[:, 0] + preds
        target_nan = deepcopy(target)
        target_nan[self.is_nan[:, 0]] = float("nan")

        # Regular loss
        score_true = spearman_corrcoef(preds, target)
        score_ipu = spearman_ipu(preds, target)
        self.assertFalse(score_true.isnan(), "Spearman is NaN")
        self.assertAlmostEqual(score_true.item(), score_ipu.item(), places=4, msg="Spearman is different")

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = spearman_corrcoef(preds[not_nan], target[not_nan])
        score_ipu = spearman_ipu(preds, target_nan)
        self.assertFalse(score_true.isnan(), "Regular Spearman with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "IPU Spearman score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=4, msg="Spearman with NaN is different"
        )

    def test_r2_score(self):
        preds = deepcopy(self.preds)
        target = deepcopy(self.target) + preds
        target_nan = deepcopy(target)
        target_nan[self.is_nan] = float("nan")

        # Regular loss
        score_true = r2_score(preds, target)
        score_ipu = r2_score_ipu(preds, target)
        self.assertFalse(score_true.isnan(), "r2_score is NaN")
        self.assertAlmostEqual(score_true.item(), score_ipu.item(), places=4, msg="r2_score is different")

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        score_ipu = r2_score_ipu(preds, target_nan, multioutput="raw_values")
        for ii in range(preds.shape[1]):
            score_true = r2_score(
                preds[:, ii][not_nan[:, ii]], target_nan[:, ii][not_nan[:, ii]], multioutput="raw_values"
            )
            self.assertFalse(score_true.isnan().any(), f"{ii}: r2_score with target_nan is NaN")
            self.assertFalse(score_ipu[ii].isnan().any(), f"{ii}: IPU r2_score with target_nan is NaN")
            self.assertAlmostEqual(
                score_true.item(), score_ipu[ii].item(), places=4, msg=f"{ii}: r2_score with NaN is different"
            )

    def test_fbeta_score(self):
        preds = deepcopy(self.preds)[:, :4]
        target = deepcopy(self.target)[:, 0]
        t = deepcopy(target)

        target[t < 0.4] = 0
        target[(t >= 0.4) & (t < 0.6)] = 1
        target[(t >= 0.6) & (t < 0.8)] = 2
        target[(t >= 0.8)] = 3

        target_nan = deepcopy(target)
        target_nan[self.is_nan[:, 0]] = float("nan")
        target_nan_bin = deepcopy(target_nan)
        target_nan_bin[target_nan > 0] = 1

        # Micro fbeta_score binary
        score_true = fbeta_score(preds[:, 0], target.to(int) > 0, average="micro", beta=0.5)
        score_ipu = fbeta_score_ipu(preds[:, 0], target > 0, average="micro", beta=0.5)
        self.assertFalse(score_true.isnan(), "Micro FBETA_score binary is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro FBETA_score binary is different"
        )

        # Micro fbeta_score binary with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = fbeta_score(
            preds[:, 0][not_nan], target_nan_bin[not_nan].to(int), average="micro", beta=0.5
        )
        score_ipu = fbeta_score_ipu(preds[:, 0], target_nan_bin, average="micro", beta=0.5)
        self.assertFalse(score_true.isnan(), "Micro FBETA_score binary with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro FBETA_score binary IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Micro FBETA_score binary with NaN is different",
        )

        # Micro fbeta_score
        score_true = fbeta_score(preds, target.to(int), average="micro", beta=0.5)
        score_ipu = fbeta_score_ipu(preds, target, average="micro", beta=0.5)
        self.assertFalse(score_true.isnan(), "Micro FBETA_score is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro FBETA_score is different"
        )

        # Micro fbeta_score with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = fbeta_score(preds[not_nan], target[not_nan].to(int), average="micro", beta=0.5)
        score_ipu = fbeta_score_ipu(preds, target_nan, average="micro", beta=0.5)
        self.assertFalse(score_true.isnan(), "Micro FBETA_score with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro FBETA_score IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro FBETA_score with NaN is different"
        )

        # Macro fbeta_score multiclass
        score_true = fbeta_score(preds, target.to(int), average="macro", num_classes=4, beta=0.5)
        score_ipu = fbeta_score_ipu(preds, target, average="macro", num_classes=4, beta=0.5)
        self.assertFalse(score_true.isnan(), "Macro FBETA_score is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Macro FBETA_score multiclass is different"
        )

        # Macro fbeta_score multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = fbeta_score(
            preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4, beta=0.5
        )
        score_ipu = fbeta_score_ipu(preds, target_nan, average="macro", num_classes=4, beta=0.5)
        self.assertFalse(score_true.isnan(), "Macro FBETA_score multiclass with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Macro FBETA_score multiclass IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Macro FBETA_score multiclass with NaN is different",
        )

        # Weighted fbeta_scoremulticlass
        score_true = fbeta_score(preds, target.to(int), average="weighted", num_classes=4, beta=0.5)
        score_ipu = fbeta_score_ipu(preds, target, average="weighted", num_classes=4, beta=0.5)
        self.assertFalse(score_true.isnan(), "Weighted FBETA_score is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Weighted FBETA_score is different"
        )

        # Weighted fbeta_score multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = fbeta_score(
            preds[not_nan], target[not_nan].to(int), average="weighted", num_classes=4, beta=0.5
        )
        score_ipu = fbeta_score_ipu(preds, target_nan, average="weighted", num_classes=4, beta=0.5)
        self.assertFalse(score_true.isnan(), "Weighted FBETA_score multiclass with target_nan is NaN")
        self.assertFalse(
            score_ipu.isnan(), "Weighted FBETA_score multiclass IPU score with target_nan is NaN"
        )
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Regular FBETA_score multiclass with NaN is different",
        )

    def test_f1_score(self):
        preds = deepcopy(self.preds)[:, :4]
        target = deepcopy(self.target)[:, 0]
        t = deepcopy(target)

        target[t < 0.4] = 0
        target[(t >= 0.4) & (t < 0.6)] = 1
        target[(t >= 0.6) & (t < 0.8)] = 2
        target[(t >= 0.8)] = 3

        target_nan = deepcopy(target)
        target_nan[self.is_nan[:, 0]] = float("nan")
        target_nan_bin = deepcopy(target_nan)
        target_nan_bin[target_nan > 0] = 1

        # Micro f1_score binary
        score_true = f1_score(preds[:, 0], target.to(int) > 0, average="micro")
        score_ipu = f1_score_ipu(preds[:, 0], target > 0, average="micro")
        self.assertFalse(score_true.isnan(), "Micro F1_score binary is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro F1_score binary is different"
        )

        # Micro f1_score binary with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = f1_score(preds[:, 0][not_nan], target_nan_bin[not_nan].to(int), average="micro")
        score_ipu = f1_score_ipu(preds[:, 0], target_nan_bin, average="micro")
        self.assertFalse(score_true.isnan(), "Micro F1_score binary with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro F1_score binary IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro F1_score binary with NaN is different"
        )

        # Micro f1_score
        score_true = f1_score(preds, target.to(int), average="micro")
        score_ipu = f1_score_ipu(preds, target, average="micro")
        self.assertFalse(score_true.isnan(), "Micro F1_score is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro F1_score is different"
        )

        # Micro f1_score with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = f1_score(preds[not_nan], target[not_nan].to(int), average="micro")
        score_ipu = f1_score_ipu(preds, target_nan, average="micro")
        self.assertFalse(score_true.isnan(), "Micro F1_score with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro F1_score IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro F1_score with NaN is different"
        )

        # Macro f1_score multiclass
        score_true = f1_score(preds, target.to(int), average="macro", num_classes=4)
        score_ipu = f1_score_ipu(preds, target, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro F1_score is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Macro F1_score multiclass is different"
        )

        # Macro f1_score multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = f1_score(preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4)
        score_ipu = f1_score_ipu(preds, target_nan, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro F1_score multiclass with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Macro F1_score multiclass IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Macro F1_score multiclass with NaN is different",
        )

        # Weighted f1_scoremulticlass
        score_true = f1_score(preds, target.to(int), average="weighted", num_classes=4)
        score_ipu = f1_score_ipu(preds, target, average="weighted", num_classes=4)
        self.assertFalse(score_true.isnan(), "Weighted F1_score is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Weighted F1_score is different"
        )

        # Weighted f1_score multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = f1_score(preds[not_nan], target[not_nan].to(int), average="weighted", num_classes=4)
        score_ipu = f1_score_ipu(preds, target_nan, average="weighted", num_classes=4)
        self.assertFalse(score_true.isnan(), "Weighted F1_score multiclass with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Weighted F1_score multiclass IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Regular F1_score multiclass with NaN is different",
        )

    def test_mse(self):
        preds = deepcopy(self.preds)
        target = deepcopy(self.target)
        target_nan = deepcopy(self.target_nan)
        squared = True

        # Regular loss
        loss_true = mean_squared_error(preds, target, squared)
        loss_ipu = mean_squared_error_ipu(preds=preds, target=target, squared=squared)
        self.assertFalse(loss_true.isnan(), "Regular Mean Squared Error is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular Mean Squared Error is different"
        )

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        loss_true = mean_squared_error(preds[not_nan], target[not_nan], squared)
        loss_ipu = mean_squared_error_ipu(preds=preds, target=target_nan, squared=squared)
        self.assertFalse(loss_true.isnan(), "Regular Mean Squared Error with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Regular Mean Squared Error IPU with target_nan is NaN")
        self.assertAlmostEqual(
            loss_true.item(),
            loss_ipu.item(),
            places=6,
            msg="Regular Mean Squared Error with NaN is different",
        )

        squared = False

        # Regular loss
        loss_true = mean_squared_error(preds, target, squared)
        loss_ipu = mean_squared_error_ipu(preds=preds, target=target, squared=squared)
        self.assertFalse(loss_true.isnan(), "Regular Mean Squared Error is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular Mean Squared Error is different"
        )

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        loss_true = mean_squared_error(preds[not_nan], target[not_nan], squared)
        loss_ipu = mean_squared_error_ipu(preds=preds, target=target_nan, squared=squared)
        self.assertFalse(loss_true.isnan(), "Regular Mean Squared Error with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Regular Mean Squared Error IPU with target_nan is NaN")
        self.assertAlmostEqual(
            loss_true.item(),
            loss_ipu.item(),
            places=6,
            msg="Regular Mean Squared Error with NaN is different",
        )

    def test_mae(self):
        preds = deepcopy(self.preds)
        target = deepcopy(self.target)
        target_nan = deepcopy(self.target_nan)

        # Regular loss
        loss_true = mean_absolute_error(preds, target)
        loss_ipu = mean_absolute_error_ipu(preds=preds, target=target)
        self.assertFalse(loss_true.isnan(), "Regular Mean Absolute Error is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular Mean Absolute Error is different"
        )

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        loss_true = mean_absolute_error(preds[not_nan], target[not_nan])
        loss_ipu = mean_absolute_error_ipu(preds=preds, target=target_nan)
        self.assertFalse(loss_true.isnan(), "Regular Mean Absolute Error with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Regular Mean Absolute Error IPU with target_nan is NaN")
        self.assertAlmostEqual(
            loss_true.item(),
            loss_ipu.item(),
            places=6,
            msg="Regular Mean Absolute Error with NaN is different",
        )


if __name__ == "__main__":
    ut.main()


# tests/test_ipu_options.py

import pytest
from graphium.config._loader import _get_ipu_opts, load_ipu_options
from graphium.ipu.ipu_utils import ipu_options_list_to_file

import tempfile
from typing import Optional, List
import os

CONFIG_EXTRACT = {
    "trainer": {"trainer": {"accumulate_grad_batches": 10}},
    "accelerator": {
        "type": "ipu",
        "config_override": {
            "datamodule": {
                "args": {
                    "ipu_dataloader_training_opts": {
                        "mode": "async",
                        "max_num_nodes_per_graph": 44,
                        "max_num_edges_per_graph": 80,
                    },
                    "ipu_dataloader_inference_opts": {
                        "mode": "async",
                        "max_num_nodes_per_graph": 44,
                        "max_num_edges_per_graph": 80,
                    },
                    "batch_size_training": 50,
                    "batch_size_inference": 50,
                }
            },
            "predictor": {"optim_kwargs": {"loss_scaling": 1024}},
            "trainer": {"trainer": {"precision": 16, "accumulate_grad_batches": 4}},
        },
        "ipu_config": [
            "deviceIterations(5)",
            "replicationFactor(16)",
            "TensorLocations.numIOTiles(128)",
            '_Popart.set("defaultBufferingDepth", 128)',
            "Precision.enableStochasticRounding(True)",
        ],
        "ipu_inference_config": [
            "deviceIterations(1)",
            "replicationFactor(4)",
            "TensorLocations.numIOTiles(32)",
            '_Popart.set("defaultBufferingDepth", 16)',
            "Precision.enableStochasticRounding(True)",
        ],
    },
}


@pytest.mark.ipu
def test_ipu_options():
    try:
        import poptorch

        ipu_opts, ipu_inference_opts = _get_ipu_opts(CONFIG_EXTRACT)

        # Define the expected IPU options for comparison
        expected_ipu_opts = [
            "deviceIterations(5)",
            "replicationFactor(16)",
            "TensorLocations.numIOTiles(128)",
            '_Popart.set("defaultBufferingDepth", 128)',
            "Precision.enableStochasticRounding(True)",
        ]
        expected_ipu_inference_opts = [
            "deviceIterations(1)",
            "replicationFactor(4)",
            "TensorLocations.numIOTiles(32)",
            '_Popart.set("defaultBufferingDepth", 16)',
            "Precision.enableStochasticRounding(True)",
        ]

        # Test the _get_ipu_opts method
        ipu_opts, ipu_inference_opts = _get_ipu_opts(CONFIG_EXTRACT)
        assert ipu_opts == expected_ipu_opts, f"Expected {expected_ipu_opts}, but got {ipu_opts}"
        assert (
            ipu_inference_opts == expected_ipu_inference_opts
        ), f"Expected {expected_ipu_inference_opts}, but got {ipu_inference_opts}"

        # Test the load_ipu_options method
        ipu_training_opts, ipu_inference_opts = load_ipu_options(
            ipu_opts=ipu_opts,
            seed=42,
            model_name="test_model",
            gradient_accumulation=CONFIG_EXTRACT["trainer"]["trainer"].get("accumulate_grad_batches", None),
            ipu_inference_opts=ipu_inference_opts,
        )

        # Ensure that the options objects are not None
        assert ipu_training_opts is not None, "Expected ipu_training_opts not to be None"
        assert ipu_inference_opts is not None, "Expected ipu_inference_opts not to be None"

        # Test the properties of the options objects
        assert (
            ipu_training_opts.replication_factor == 16
        ), "Expected replication_factor of ipu_training_opts to be 16"
        assert (
            ipu_inference_opts.replication_factor == 4
        ), "Expected replication_factor of ipu_inference_opts to be 4"
        assert ipu_training_opts._popart, "Expected _popart of ipu_training_opts to be True"
        assert ipu_inference_opts._popart, "Expected _popart of ipu_inference_opts to be True"

    except ImportError:
        pytest.skip("Skipping this test because poptorch is not available")


@pytest.mark.ipu
def test_ipu_options_list_to_file():
    # Define a list of IPU options
    ipu_options = [
        "deviceIterations(5)",
        "replicationFactor(16)",
        "TensorLocations.numIOTiles(128)",
        '_Popart.set("defaultBufferingDepth", 128)',
        "Precision.enableStochasticRounding(True)",
    ]

    # Call the function with the list of IPU options
    tmp_file = ipu_options_list_to_file(ipu_options)

    # Check that the function returns a temporary file object
    assert isinstance(tmp_file, tempfile._TemporaryFileWrapper)

    # Check that the temporary file exists
    assert os.path.exists(tmp_file.name)

    # Check the contents of the temporary file
    with open(tmp_file.name, "r") as f:
        contents = f.read().splitlines()
    assert contents == ipu_options

    # Check the behavior when the input is None
    tmp_file = ipu_options_list_to_file(None)
    assert tmp_file is None


# tests/test_ipu_to_dense_batch.py

import pytest
import torch
from torch_geometric.data import Data, Batch
from graphium.ipu.to_dense_batch import to_dense_batch
from warnings import warn


# General imports
import yaml
import unittest as ut
import numpy as np
from copy import deepcopy
from warnings import warn
from lightning import Trainer, LightningModule
from lightning_graphcore import IPUStrategy
from functools import partial

import torch
from torch.utils.data.dataloader import default_collate

# Current library imports
from graphium.config._loader import load_datamodule, load_metrics, load_architecture, load_accelerator


@pytest.mark.ipu
class TestIPUBatch:
    @pytest.fixture(autouse=True)
    def setup_class(self):
        self.in_dim = 12
        self.out_dim = 12
        self.in_dim_edges = 10
        self.out_dim_edges = 10
        self.edge_idx1 = torch.stack(
            [torch.tensor([0, 1, 2, 3, 2], dtype=torch.int), torch.tensor([1, 2, 3, 0, 0], dtype=torch.int)]
        )
        self.edge_idx2 = torch.stack(
            [torch.tensor([0, 0, 0, 1], dtype=torch.int), torch.tensor([0, 1, 2, 0], dtype=torch.int)]
        )
        self.x1 = torch.randn(self.edge_idx1.max().item() + 1, self.in_dim, dtype=torch.float32)
        self.e1 = torch.randn(self.edge_idx1.shape[-1], self.in_dim_edges, dtype=torch.float32)
        self.x2 = torch.randn(self.edge_idx2.max().item() + 1, self.in_dim, dtype=torch.float32)
        self.e2 = torch.randn(self.edge_idx2.shape[-1], self.in_dim_edges, dtype=torch.float32)
        self.g1 = Data(feat=self.x1, edge_index=self.edge_idx1, edge_feat=self.e1)
        self.g2 = Data(feat=self.x2, edge_index=self.edge_idx2, edge_feat=self.e2)
        self.bg = Batch.from_data_list([self.g1, self.g2])
        self.attn_kwargs = {"embed_dim": self.in_dim, "num_heads": 2, "batch_first": True}

    @pytest.mark.parametrize("max_num_nodes_per_graph, batch_size", [(10, 5), (20, 10), (30, 15)])
    def test_ipu_to_dense_batch(self, max_num_nodes_per_graph, batch_size):
        # Run this test only if poptorch is available
        try:
            import poptorch

            opts = poptorch.Options()
            opts.useIpuModel(True)

            class MyModel(torch.nn.Module):
                def __init__(self):
                    super(MyModel, self).__init__()

                def forward(self, x, batch):
                    return to_dense_batch(
                        x,
                        batch=batch,
                        batch_size=batch_size,
                        max_num_nodes_per_graph=max_num_nodes_per_graph,
                        drop_nodes_last_graph=False,
                    )

            model = MyModel()
            model = model.eval()
            poptorch_model_inf = poptorch.inferenceModel(model, options=opts)
            # for data in train_dataloader:
            out, mask, idx = poptorch_model_inf(self.bg.feat, self.bg.batch)
            # Check the output sizes
            assert out.size() == torch.Size([batch_size, max_num_nodes_per_graph, 12])
            # Check the mask for true / false values
            assert mask.size() == torch.Size([batch_size, max_num_nodes_per_graph])
            assert torch.sum(mask) == 7
            assert (mask[0][:4] == True).all()
            assert (mask[0][4:] == False).all()
            assert (mask[1][:3] == True).all()
            assert (mask[1][3:] == False).all()
            assert (mask[2:] == False).all()

            # Check the idx are all the true values in the mask
            assert (mask.flatten()[idx] == True).all()
        except ImportError:
            pytest.skip("Skipping this test because poptorch is not available")

    def test_ipu_to_dense_batch_no_batch_no_max_nodes(self):
        h_dense, mask = to_dense_batch(
            self.bg.feat,
            batch=None,
            batch_size=None,
            max_num_nodes_per_graph=None,
            drop_nodes_last_graph=False,
        )
        # Add assertions to check the output as needed
        assert torch.allclose(h_dense, self.bg.feat.unsqueeze(0), atol=1e-5), "Tensors are not equal"
        assert mask.size(1) == h_dense.size(1)
        assert mask.all().item(), "Not all values in the tensor are True"

    def test_ipu_to_dense_batch_no_batch(self):
        max_nodes_per_graph = 10
        h_dense, mask, id = to_dense_batch(
            self.bg.feat,
            batch=None,
            batch_size=None,
            max_num_nodes_per_graph=max_nodes_per_graph,
            drop_nodes_last_graph=False,
        )
        assert mask.size() == (1, max_nodes_per_graph)
        assert torch.sum(mask) == 7
        assert torch.equal(id, torch.arange(7))
        assert h_dense.size() == (1, max_nodes_per_graph, self.bg.feat.size(-1))

    def test_ipu_to_dense_batch_drop_last(self):
        out, mask, idx = to_dense_batch(
            self.bg.feat,
            batch=None,
            batch_size=None,
            max_num_nodes_per_graph=3,
            drop_nodes_last_graph=True,
        )
        # Add assertions to check the output as needed
        assert mask.size(1) == out.size(1)
        # Check the mask and output have been clipped
        assert mask.size() == torch.Size([1, 3])
        assert mask.all().item(), "Not all values in the tensor are True"


# tests/test_loaders.py

from graphium.config._loader import merge_dicts
from copy import deepcopy
import unittest as ut


class TestLoader(ut.TestCase):
    def test_merge_dicts(self):
        dict_a = {"a": {"b": {"c": 1, "d": 2}, "e": 3}, "f": 4}

        dict_b = {"a": {"b": {"g": 6}}, "h": 7}

        dict_c = {
            "a": {
                "b": {
                    "c": 1,
                },
            }
        }

        # Check that the new keys are added correctly
        merge_dicts(dict_a, dict_b)
        self.assertEqual(dict_a["a"]["b"]["g"], dict_b["a"]["b"]["g"])
        self.assertEqual(dict_a["h"], dict_b["h"])

        # Check that no error is thrown if a key exists, but the value is identical
        merge_dicts(dict_a, dict_c)
        self.assertEqual(dict_a["a"]["b"]["c"], dict_c["a"]["b"]["c"])

        # Check that an error is thrown if a key exists, but the value is different
        dict_d = deepcopy(dict_c)
        dict_d["a"]["b"]["c"] = 2
        with self.assertRaises(ValueError):
            merge_dicts(dict_a, dict_d)


# Main
if __name__ == "__main__":
    ut.main()


# tests/test_losses.py

"""
Unit tests for the metrics and wrappers of graphium/trainer/metrics/...
"""

import torch
import unittest as ut
from torch.nn import functional as F

from graphium.trainer.losses import HybridCELoss
from graphium.trainer.predictor_options import EvalOptions


def _parse(loss_fun):
    eval_options = EvalOptions(loss_fun=loss_fun, metrics_on_progress_bar=None)
    return eval_options.parse_loss_fun(loss_fun)


class test_HybridCELoss(ut.TestCase):
    input = torch.Tensor([[0.1, 0.1, 0.3, 0.5, 0.0], [0.1, 0.0, 0.7, 0.2, 0.0]])
    target = torch.Tensor([3, 0]).long()
    brackets = torch.Tensor([0, 1, 2, 3, 4])
    regression_input = torch.Tensor([2.0537, 2.0017])  # inner product of input and brackets
    regression_target = torch.Tensor([3, 0]).float()

    def test_pure_ce_loss(self):
        loss = HybridCELoss(n_brackets=len(self.brackets), alpha=1.0, reduction="none")
        assert torch.allclose(
            loss(self.input, self.target),
            F.cross_entropy(self.input, self.target, reduction="none"),
        )
        assert loss(self.input, self.target).shape == (2,)

    def test_pure_mae_loss(self):
        loss = HybridCELoss(
            n_brackets=len(self.brackets),
            alpha=0.0,
            regression_loss="mae",
            reduction="none",
        )
        assert torch.allclose(
            loss(self.input, self.target),
            F.l1_loss(self.regression_input, self.regression_target, reduction="none"),
            rtol=1e-04,
            atol=1e-07,
        )
        assert loss(self.input, self.target).shape == (2,)

    def test_pure_mse_loss(self):
        loss = HybridCELoss(
            n_brackets=len(self.brackets),
            alpha=0.0,
            regression_loss="mse",
            reduction="none",
        )

        assert torch.allclose(
            loss(self.input, self.target),
            F.mse_loss(self.regression_input, self.regression_target, reduction="none"),
            rtol=1e-04,
            atol=1e-07,
        )
        assert loss(self.input, self.target).shape == (2,)

    def test_hybrid_loss(self):
        loss = HybridCELoss(n_brackets=len(self.brackets), alpha=0.5, regression_loss="mse")

        ce_loss = F.cross_entropy(self.input, self.target)
        mse_loss = F.mse_loss(self.regression_input, self.regression_target)

        assert torch.allclose(
            loss(self.input, self.target), 0.5 * ce_loss + 0.5 * mse_loss, rtol=1e-04, atol=1e-07
        )
        assert loss(self.input, self.target).shape == torch.Size([])

    def test_loss_parser(self):
        # HybridCE cannot be parsed from a string because it requires specifying n_brackets
        loss_fun = "hybrid_ce"
        self.assertRaises(TypeError, _parse, loss_fun=loss_fun)

        # HybridCE requires n_brackets to be specified
        loss_fun = {"name": "hybrid_ce"}
        self.assertRaises(TypeError, _parse, loss_fun=loss_fun)

        loss_fun = {"name": "hybrid_ce", "n_brackets": 3}
        parsed = _parse(loss_fun)

        assert isinstance(parsed, HybridCELoss)
        assert len(parsed.brackets) == 3
        assert parsed.regression_loss == F.mse_loss

        loss_fun = {"name": "hybrid_ce", "n_brackets": 5, "regression_loss": "mae"}
        parsed = _parse(loss_fun)

        assert isinstance(parsed, HybridCELoss)
        assert len(parsed.brackets) == 5
        assert parsed.regression_loss == F.l1_loss


class test_BCELoss(ut.TestCase):
    def test_loss_parser(self):
        loss_fun = "bce"
        parsed = _parse(loss_fun)

        assert isinstance(parsed, torch.nn.BCELoss)

        loss_fun = {"name": "bce"}
        parsed = _parse(loss_fun)

        assert isinstance(parsed, torch.nn.BCELoss)
        assert parsed.reduction != "sum"

        loss_fun = {"name": "bce", "reduction": "sum"}
        parsed = _parse(loss_fun)

        assert isinstance(parsed, torch.nn.BCELoss)
        assert parsed.reduction == "sum"


if __name__ == "__main__":
    ut.main()


# tests/test_metrics.py

"""
Unit tests for the metrics and wrappers of graphium/trainer/metrics/...
"""

import torch
import unittest as ut
import tempfile
import os
import operator as op

from graphium.trainer.metrics import (
    MetricWrapper,
    Thresholder,
)

from torchmetrics.functional import mean_squared_error


class test_Metrics(ut.TestCase):
    def test_thresholder(self):
        torch.manual_seed(42)
        preds = torch.rand(100, dtype=torch.float32)
        target = torch.rand(100, dtype=torch.float32)

        th = 0.7
        preds_greater = preds > th
        target_greater = target > th

        # Test thresholder greater
        for th_on_preds in [True, False]:
            for th_on_target in [True, False]:
                thresholder = Thresholder(
                    threshold=th, operator="greater", th_on_target=th_on_target, th_on_preds=th_on_preds
                )
                preds2, target2 = thresholder(preds, target)
                if th_on_preds:
                    self.assertListEqual(preds2.tolist(), preds_greater.tolist())
                else:
                    self.assertListEqual(preds2.tolist(), preds.tolist())
                if th_on_target:
                    self.assertListEqual(target2.tolist(), target_greater.tolist())
                else:
                    self.assertListEqual(target2.tolist(), target.tolist())

        # Test thresholder lower
        for th_on_preds in [True, False]:
            for th_on_target in [True, False]:
                thresholder = Thresholder(
                    threshold=th, operator="lower", th_on_target=th_on_target, th_on_preds=th_on_preds
                )
                preds2, target2 = thresholder(preds, target)
                if th_on_preds:
                    self.assertListEqual(preds2.tolist(), (~preds_greater).tolist())
                else:
                    self.assertListEqual(preds2.tolist(), preds.tolist())
                if th_on_target:
                    self.assertListEqual(target2.tolist(), (~target_greater).tolist())
                else:
                    self.assertListEqual(target2.tolist(), target.tolist())


class test_MetricWrapper(ut.TestCase):
    def test_target_nan_mask(self):
        torch.random.manual_seed(42)

        for sz in [(100,), (100, 1), (100, 10)]:
            err_msg = f"Error for `sz = {sz}`"

            # Generate prediction and target matrices
            target = torch.rand(sz, dtype=torch.float32)
            preds = (0.5 * target) + (0.5 * torch.rand(sz, dtype=torch.float32))
            is_nan = torch.rand(sz) > 0.3
            target = (target > 0.5).to(torch.float32)
            target[is_nan] = float("nan")

            # Compute score with different ways of ignoring NaNs
            metric = MetricWrapper(metric="mse", target_nan_mask=None)
            score1 = metric(preds, target)
            self.assertTrue(torch.isnan(score1), msg=err_msg)

            # Replace NaNs by 0
            metric = MetricWrapper(metric="mse", target_nan_mask=0)
            score2 = metric(preds, target)

            this_target = target.clone()
            this_target[is_nan] = 0
            this_preds = preds.clone()
            self.assertAlmostEqual(score2, mean_squared_error(this_preds, this_target), msg=err_msg)

            # Replace NaNs by 1.5
            metric = MetricWrapper(metric="mse", target_nan_mask=1.5)
            score3 = metric(preds, target)

            this_target = target.clone()
            this_target[is_nan] = 1.5
            this_preds = preds.clone()
            self.assertAlmostEqual(score3, mean_squared_error(this_preds, this_target), msg=err_msg)

            # Flatten matrix and ignore NaNs
            metric = MetricWrapper(metric="mse", target_nan_mask="ignore", multitask_handling="flatten")
            score4 = metric(preds, target)

            this_target = target.clone()[~is_nan]
            this_preds = preds.clone()[~is_nan]
            self.assertAlmostEqual(score4, mean_squared_error(this_preds, this_target), msg=err_msg)

            # Ignore NaNs in each column and average the score
            metric = MetricWrapper(
                metric="mse", target_nan_mask="ignore", multitask_handling="mean-per-label"
            )
            score5 = metric(preds, target)

            this_target = target.clone()
            this_preds = preds.clone()
            this_is_nan = is_nan.clone()
            if len(sz) == 1:
                this_target = target.unsqueeze(-1)
                this_preds = preds.unsqueeze(-1)
                this_is_nan = is_nan.unsqueeze(-1)

            this_target = [this_target[:, ii][~this_is_nan[:, ii]] for ii in range(this_target.shape[1])]
            this_preds = [this_preds[:, ii][~this_is_nan[:, ii]] for ii in range(this_preds.shape[1])]
            mse = []
            for ii in range(len(this_preds)):
                mse.append(mean_squared_error(this_preds[ii], this_target[ii]))
            mse = torch.mean(torch.stack(mse))
            self.assertAlmostEqual(score5.tolist(), mse.tolist(), msg=err_msg)

    def test_pickling(self):
        pickle_file = os.path.join(tempfile.gettempdir(), "test_metric_pickled.pkl")
        metrics = ["mae", "mse", mean_squared_error]
        target_nan_masks = [None, 2, "ignore"]
        multitask_handlings = [None, "flatten", "mean-per-label"]
        squeeze_targets = [True, False]
        target_to_ints = [True, False]
        other_kwargs = [{}, {"squared": False}]
        thresholds = [
            None,
            {"threshold": 0.2, "operator": "greater"},
            {"threshold": 0.3, "operator": op.lt},
            {"threshold": 0.4, "operator": "lower"},
            {"threshold": 0.5, "operator": "lower", "th_on_preds": False, "th_on_target": True},
            {"threshold": 0.6, "operator": "lower"},
        ]

        for metric in metrics:
            for target_nan_mask in target_nan_masks:
                for kwargs in other_kwargs:
                    for threshold_kwargs in thresholds:
                        for multitask_handling in multitask_handlings:
                            for squeeze_target in squeeze_targets:
                                for target_to_int in target_to_ints:
                                    err_msg = f"{metric} - {target_nan_mask} - {kwargs} - {threshold_kwargs}"

                                    if (multitask_handling is None) and (target_nan_mask == "ignore"):
                                        # Raise with incompatible options
                                        with self.assertRaises(ValueError):
                                            MetricWrapper(
                                                metric=metric,
                                                threshold_kwargs=threshold_kwargs,
                                                target_nan_mask=target_nan_mask,
                                                multitask_handling=multitask_handling,
                                                squeeze_target=squeeze_target,
                                                target_to_int=target_to_int,
                                                **kwargs,
                                            )

                                    else:
                                        metric_wrapper = MetricWrapper(
                                            metric=metric,
                                            threshold_kwargs=threshold_kwargs,
                                            target_nan_mask=target_nan_mask,
                                            multitask_handling=multitask_handling,
                                            squeeze_target=squeeze_target,
                                            target_to_int=target_to_int,
                                            **kwargs,
                                        )

                                        # Check that the metric can be saved and re-loaded without error
                                        torch.save(metric_wrapper, pickle_file)
                                        metric_wrapper2 = torch.load(pickle_file)
                                        self.assertTrue(metric_wrapper == metric_wrapper2, msg=err_msg)

                                        # Check that the metric only contains primitive types
                                        state = metric_wrapper.__getstate__()
                                        if state["threshold_kwargs"] is not None:
                                            self.assertIsInstance(
                                                state["threshold_kwargs"], dict, msg=err_msg
                                            )
                                        if isinstance(metric, str):
                                            self.assertIsInstance(state["metric"], str, msg=err_msg)

    def test_classifigression_target_squeezing(self):
        preds = torch.Tensor([[0.1, 0.1, 0.3, 0.5, 0.0, 0.1, 0.0, 0.7, 0.2, 0.0]])
        target = torch.Tensor([3, 0])
        expected_scores = [0.5, 0.75]
        n_brackets = 5
        metrics = ["accuracy", "averageprecision"]
        other_kwargs = [
            {"task": "multiclass", "num_classes": n_brackets, "top_k": 1},
            {"task": "multiclass", "num_classes": n_brackets},
        ]

        for metric, kwargs, expected_score in zip(metrics, other_kwargs, expected_scores):
            metric_wrapper = MetricWrapper(
                metric=metric,
                multitask_handling="mean-per-label",
                squeeze_targets=True,
                target_to_int=True,
                **kwargs,
            )
            score = metric_wrapper(preds, target)

            assert score == expected_score


if __name__ == "__main__":
    ut.main()


# tests/test_mtl_architecture.py

"""
Unit tests for the different architectures of graphium/nn/architectures...

The layers are not thoroughly tested due to the difficulty of testing them
"""

import torch
import unittest as ut
from torch_geometric.data import Data, Batch
from copy import deepcopy
from itertools import product
import sys
import traceback

import graphium
from graphium.config._loader import load_architecture
from graphium.nn.architectures import TaskHeads, FullGraphMultiTaskNetwork, GraphOutputNN
from graphium.nn.base_layers import FCLayer

from graphium.utils.spaces import LAYERS_DICT

kwargs = {
    "activation": "relu",
    "last_activation": "none",
    "normalization": "none",
    "dropout": 0.2,
    "name": "LNN",
    "layer_type": FCLayer,
    "residual_type": "none",
    "residual_skip_steps": 1,
}
# task kwargs
task_1_kwargs = {
    "out_dim": 5,
    "task_level": "node",
    "hidden_dims": [5, 6, 7],
}
task_2_kwargs = {
    "out_dim": 3,
    "task_level": "edge",
    "hidden_dims": [8, 9, 10],
}
task_3_kwargs = {
    "out_dim": 4,
    "task_level": "graph",
    "hidden_dims": [2, 2, 2],
}
task_4_kwargs = {
    "out_dim": 4,
    "task_level": "nodepair",
    "hidden_dims": [2, 2, 2],
}

# level-wise kwargs
node_level_kwargs = {
    "out_dim": 8,
    "hidden_dims": [8, 9, 10],
    "activation": "relu",
    "last_activation": "none",
    "normalization": "none",
    "dropout": 0.2,
    "name": "LNN",
    "layer_type": FCLayer,
    "residual_type": "none",
    "residual_skip_steps": 1,
}
graph_level_kwargs = {
    "pooling": ["sum", "max"],
    "out_dim": 8,
    "hidden_dims": [8, 9, 10],
    "activation": "relu",
    "last_activation": "none",
    "normalization": "none",
    "dropout": 0.2,
    "name": "LNN",
    "layer_type": FCLayer,
    "residual_type": "none",
    "residual_skip_steps": 1,
}
edge_level_kwargs = {
    "out_dim": 8,
    "hidden_dims": [8, 9, 10],
    "activation": "relu",
    "last_activation": "none",
    "normalization": "none",
    "dropout": 0.2,
    "name": "LNN",
    "layer_type": FCLayer,
    "residual_type": "none",
    "residual_skip_steps": 1,
}

nodepair_level_kwargs = {
    "out_dim": 8,
    "hidden_dims": [8, 9, 10],
    "activation": "relu",
    "last_activation": "none",
    "normalization": "none",
    "dropout": 0.2,
    "name": "LNN",
    "layer_type": FCLayer,
    "residual_type": "none",
    "residual_skip_steps": 1,
}

task_1_params = {}
task_1_params.update(task_1_kwargs)
task_1_params.update(kwargs)
task_2_params = {}
task_2_params.update(task_2_kwargs)
task_2_params.update(kwargs)
task_3_params = {}
task_3_params.update(task_3_kwargs)
task_3_params.update(kwargs)
task_4_params = {}
task_4_params.update(task_4_kwargs)
task_4_params.update(kwargs)


def toy_test_data(in_dim=7, in_dim_edges=3):
    edge_idx1 = torch.stack([torch.tensor([0, 1, 2, 3, 2]), torch.tensor([1, 2, 3, 0, 0])])
    edge_idx2 = torch.stack([torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])])
    x1 = torch.randn(edge_idx1.max() + 1, in_dim, dtype=torch.float32)
    e1 = torch.randn(edge_idx1.shape[-1], in_dim_edges, dtype=torch.float32)
    x2 = torch.randn(edge_idx2.max() + 1, in_dim, dtype=torch.float32)
    e2 = torch.randn(edge_idx2.shape[-1], in_dim_edges, dtype=torch.float32)
    # edge_idx1, e1 = add_self_loops(edge_idx1, e1)
    # edge_idx2, e2 = add_self_loops(edge_idx2, e2)
    g1 = Data(feat=x1, edge_index=edge_idx1, edge_feat=e1)
    g2 = Data(feat=x2, edge_index=edge_idx2, edge_feat=e2)
    bg = Batch.from_data_list([g1, g2])
    num_nodes_per_graph = bg.batch.bincount()
    batch_node = bg["feat"].size()[0]
    batch_edge = bg["edge_feat"].size()[0]
    batch_graph = 2
    batch_nodepair = ((num_nodes_per_graph.max().item() ** 2) - num_nodes_per_graph.max().item()) // 2
    return bg, batch_node, batch_edge, batch_graph, batch_nodepair


class test_GraphOutputNN(ut.TestCase):
    def generate_test_data(self):
        g_1 = torch.tensor([[2, 3, 4], [0, 1, 0], [0, 1, 1], [1, 0, 0]])
        g_2 = torch.tensor([[3, 4, 0], [0, 1, 1], [1, 0, 0]])
        g_3 = torch.tensor([[3, 4, 0], [0, 1, 1], [1, 0, 0], [2, 2, 2], [4, 4, 4]])

        batch = torch.tensor([0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 2, 2])
        x = torch.concat([g_1, g_2, g_3]).float()

        nan = -1
        expected_result = [
            # graph 1
            [
                [2.0, 4.0, 4.0, 2.0, 2.0, 4.0],
                [2.0, 4.0, 5.0, 2.0, 2.0, 3.0],
                [3.0, 3.0, 4.0, 1.0, 3.0, 4.0],
                [nan, nan, nan, nan, nan, nan],
                [0.0, 2.0, 1.0, 0.0, 0.0, 1.0],
                [1.0, 1.0, 0.0, 1.0, 1.0, 0.0],
                [nan, nan, nan, nan, nan, nan],
                [1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
                [nan, nan, nan, nan, nan, nan],
                [nan, nan, nan, nan, nan, nan],
            ],
            # graph 2
            [
                [3.0, 5.0, 1.0, 3.0, 3.0, 1.0],
                [4.0, 4.0, 0.0, 2.0, 4.0, 0.0],
                [nan, nan, nan, nan, nan, nan],
                [nan, nan, nan, nan, nan, nan],
                [1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
                [nan, nan, nan, nan, nan, nan],
                [nan, nan, nan, nan, nan, nan],
                [nan, nan, nan, nan, nan, nan],
                [nan, nan, nan, nan, nan, nan],
                [nan, nan, nan, nan, nan, nan],
            ],
            # graph 3
            [
                [3.0, 5.0, 1.0, 3.0, 3.0, 1.0],
                [4.0, 4.0, 0.0, 2.0, 4.0, 0.0],
                [5.0, 6.0, 2.0, 1.0, 2.0, 2.0],
                [7.0, 8.0, 4.0, 1.0, 0.0, 4.0],
                [1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
                [2.0, 3.0, 3.0, 2.0, 1.0, 1.0],
                [4.0, 5.0, 5.0, 4.0, 3.0, 3.0],
                [3.0, 2.0, 2.0, 1.0, 2.0, 2.0],
                [5.0, 4.0, 4.0, 3.0, 4.0, 4.0],
                [6.0, 6.0, 6.0, 2.0, 2.0, 2.0],
            ],
        ]

        return x, batch, expected_result

    def test_nodepair_max_num_nodes_not_set(self):
        in_dim = 3
        in_dim_edges = 8

        graph_output_nn_kwargs = {
            "node": node_level_kwargs,
            "edge": edge_level_kwargs,
            "graph": graph_level_kwargs,
            "nodepair": nodepair_level_kwargs,
        }

        graph_output_nn = GraphOutputNN(
            in_dim=in_dim,
            in_dim_edges=in_dim_edges,
            task_level="nodepair",
            graph_output_nn_kwargs=graph_output_nn_kwargs,
        )

        x, batch, expected_result = self.generate_test_data()

        out = graph_output_nn.compute_nodepairs(node_feats=x, batch=batch)
        out = torch.nan_to_num(out, nan=-1)
        self.assertListEqual(expected_result, out.tolist())

    def test_nodepair_with_max_num_nodes(self):
        in_dim = 3
        in_dim_edges = 8

        graph_output_nn_kwargs = {
            "node": node_level_kwargs,
            "edge": edge_level_kwargs,
            "graph": graph_level_kwargs,
            "nodepair": nodepair_level_kwargs,
        }

        graph_output_nn = GraphOutputNN(
            in_dim=in_dim,
            in_dim_edges=in_dim_edges,
            task_level="nodepair",
            graph_output_nn_kwargs=graph_output_nn_kwargs,
        )

        max_num_nodes = 5  # if we change this value, we also have to change the expected result.
        x, batch, expected_result = self.generate_test_data()

        out = graph_output_nn.compute_nodepairs(node_feats=x, batch=batch, max_num_nodes=max_num_nodes)
        out = torch.nan_to_num(out, nan=-1)
        self.assertListEqual(expected_result, out.tolist())


class test_TaskHeads(ut.TestCase):
    def test_task_heads_forward(self):
        in_dim = 8  # Dimension of the incoming data
        in_dim_edges = 8

        task_heads_params = {
            "task_1": task_1_params,
            "task_2": task_2_params,
            "task_3": task_3_params,
            "task_4": task_4_params,
        }
        graph_output_nn_kwargs = {
            "node": node_level_kwargs,
            "edge": edge_level_kwargs,
            "graph": graph_level_kwargs,
            "nodepair": nodepair_level_kwargs,
        }
        # Create the "multitask" network. Really it's just an input going to various FFNNs since there's nothing shared.
        multi_head_nn = TaskHeads(
            in_dim=in_dim,
            in_dim_edges=in_dim_edges,
            task_heads_kwargs=task_heads_params,
            graph_output_nn_kwargs=graph_output_nn_kwargs,
        )

        # Test the sizes of the MLPs for each head
        # Head for task_1
        task_1_head = multi_head_nn.task_heads["task_1"]

        # Check the dimensions
        self.assertEqual(len(task_1_head.layers), len(task_1_kwargs["hidden_dims"]) + 1)
        self.assertEqual(task_1_head.layers[0].in_dim, in_dim)
        self.assertEqual(task_1_head.layers[1].in_dim, task_1_kwargs["hidden_dims"][0])
        self.assertEqual(task_1_head.layers[2].in_dim, task_1_kwargs["hidden_dims"][1])
        self.assertEqual(task_1_head.layers[3].in_dim, task_1_kwargs["hidden_dims"][2])

        # Head for task_2
        task_2_head = multi_head_nn.task_heads["task_2"]

        # Check the dimensions
        self.assertEqual(len(task_2_head.layers), len(task_2_kwargs["hidden_dims"]) + 1)
        self.assertEqual(task_2_head.layers[0].in_dim, in_dim)
        self.assertEqual(task_2_head.layers[1].in_dim, task_2_kwargs["hidden_dims"][0])
        self.assertEqual(task_2_head.layers[2].in_dim, task_2_kwargs["hidden_dims"][1])
        self.assertEqual(task_2_head.layers[3].in_dim, task_2_kwargs["hidden_dims"][2])

        # Head for task_3
        task_3_head = multi_head_nn.task_heads["task_3"]

        # Check the dimensions
        self.assertEqual(len(task_3_head.layers), len(task_3_kwargs["hidden_dims"]) + 1)
        self.assertEqual(task_3_head.layers[0].in_dim, in_dim)
        self.assertEqual(task_3_head.layers[1].in_dim, task_3_kwargs["hidden_dims"][0])
        self.assertEqual(task_3_head.layers[2].in_dim, task_3_kwargs["hidden_dims"][1])
        self.assertEqual(task_3_head.layers[3].in_dim, task_3_kwargs["hidden_dims"][2])

        # Head for task_4
        task_4_head = multi_head_nn.task_heads["task_3"]

        # Check the dimensions
        self.assertEqual(len(task_4_head.layers), len(task_4_kwargs["hidden_dims"]) + 1)
        self.assertEqual(task_4_head.layers[0].in_dim, in_dim)
        self.assertEqual(task_4_head.layers[1].in_dim, task_4_kwargs["hidden_dims"][0])
        self.assertEqual(task_4_head.layers[2].in_dim, task_4_kwargs["hidden_dims"][1])
        self.assertEqual(task_4_head.layers[3].in_dim, task_4_kwargs["hidden_dims"][2])

        # Check the output: It's a per-task prediction!
        bg, batch_node, batch_edge, batch_graph, batch_nodepair = toy_test_data(
            in_dim=in_dim, in_dim_edges=in_dim_edges
        )
        feat_out = multi_head_nn.forward(bg)

        self.assertListEqual(
            list(feat_out["task_1"].shape), [batch_node, task_1_kwargs["out_dim"]]
        )  # node level task
        self.assertListEqual(
            list(feat_out["task_2"].shape), [batch_edge, task_2_kwargs["out_dim"]]
        )  # edge level task
        self.assertListEqual(
            list(feat_out["task_3"].shape), [batch_graph, task_3_kwargs["out_dim"]]
        )  # graph level task
        self.assertListEqual(
            list(feat_out["task_4"].shape), [batch_graph, batch_nodepair, task_4_kwargs["out_dim"]]
        )  # nodepair level task

    def test_task_heads_non_supported_level(self):
        in_dim = 8  # Dimension of the incoming data
        in_dim_edges = 8

        fake_task = deepcopy(task_1_params)
        fake_task["task_level"] = "certainly_not_supported_level"
        task_heads_params = {
            "task_1": fake_task,
            "task_2": task_2_params,
            "task_3": task_3_params,
            "task_4": task_4_params,
        }
        graph_output_nn_kwargs = {
            "certainly_not_supported_level": node_level_kwargs,
            "edge": edge_level_kwargs,
            "graph": graph_level_kwargs,
            "nodepair": nodepair_level_kwargs,
        }

        with self.assertRaises(ValueError):
            TaskHeads(
                in_dim=in_dim,
                in_dim_edges=in_dim_edges,
                task_heads_kwargs=task_heads_params,
                graph_output_nn_kwargs=graph_output_nn_kwargs,
            )


class test_Multitask_NN(ut.TestCase):
    pyg_kwargs = {
        "activation": "relu",
        "last_activation": "none",
        "normalization": "none",
        "dropout": 0.2,
        "name": "LNN",
    }

    in_dim = 7
    fullgraph_out_dim = 11
    in_dim_edges = 3
    hidden_dims = [6, 6, 6, 6, 6]

    virtual_nodes = ["none", "mean", "sum"]
    norms = ["none", None, "batch_norm", "layer_norm"]
    pna_kwargs = {"aggregators": ["mean", "max", "sum"], "scalers": ["identity", "amplification"]}

    gnn_layers_kwargs = {
        "pyg:gin": {},
        "pyg:gine": {"in_dim_edges": in_dim_edges},
        "pyg:gated-gcn": {"in_dim_edges": in_dim_edges, "hidden_dims_edges": hidden_dims},
        "pyg:pna-msgpass#1": {"layer_kwargs": pna_kwargs, "in_dim_edges": 0},
        "pyg:pna-msgpass#2": {"layer_kwargs": pna_kwargs, "in_dim_edges": in_dim_edges},
    }

    def test_full_graph_multitask_forward(self):
        # Params for the network
        temp_dim_1 = 5
        temp_dim_2 = 7
        temp_dim_edges = 21

        default_graph_output_nn_kwargs = {
            "node": dict(out_dim=10, hidden_dims=[3, 3, 3, 3]),
            "edge": dict(out_dim=11, hidden_dims=[4, 4, 4, 4]),
            "graph": dict(out_dim=12, hidden_dims=[5, 5, 5, 5]),
            "nodepair": dict(out_dim=13, hidden_dims=[6, 6, 6, 6]),
        }

        # TODO: Test all combinations?
        all_task_heads_kwargs = dict(
            task_node=task_1_params,
            task_edge=task_2_params,
            task_graph=task_3_params,
            task_nodepair=task_4_params,
        )
        only_node_and_graph_task_heads_kwargs = dict(task_node=task_1_params, task_graph=task_3_params)

        # TODO: Configure properly in pytest best-practice
        options = product(
            [["sum"], ["mean", "max"]],
            [1, 2],
            self.virtual_nodes,
            self.norms,
            self.gnn_layers_kwargs.items(),
            [None, all_task_heads_kwargs, only_node_and_graph_task_heads_kwargs],
            [None, dict(in_dim=self.in_dim, out_dim=temp_dim_1, hidden_dims=[4, 4, 4, 4, 4])],
            [None, dict(in_dim=self.in_dim_edges, out_dim=temp_dim_edges, hidden_dims=[4, 4])],
        )
        for (
            pooling,
            residual_skip_steps,
            virtual_node,
            normalization,
            (layer_name, this_kwargs),
            task_heads_kwargs,
            pre_nn_kwargs,
            pre_nn_edges_kwargs,
        ) in options:
            err_msg = f"pooling={pooling}, virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}, task_heads={task_heads_kwargs}, pre_nn_kwargs={pre_nn_kwargs}, graph_output_nn_kwargs={default_graph_output_nn_kwargs}, pre_nn_edges_kwargs={pre_nn_edges_kwargs}"
            layer_type = layer_name.split("#")[0]

            # TODO: graph_output_nn is currently non-optional, should it be optional?
            if task_heads_kwargs is not None:
                continue

            # TODO: Allow to pass single graph_output_nn_kwargs to apply to all levels?

            bg, num_nodes, num_edges, num_graphs, num_nodepairs = toy_test_data(
                in_dim=self.in_dim, in_dim_edges=self.in_dim_edges
            )

            expected_dim_edges = self.in_dim_edges
            this_kwargs2 = deepcopy(this_kwargs)
            if pre_nn_edges_kwargs is not None:
                this_kwargs2["in_dim_edges"] = expected_dim_edges = temp_dim_edges

            gnn_kwargs = dict(
                in_dim=self.in_dim if pre_nn_kwargs is None else temp_dim_1,
                out_dim=temp_dim_2,
                hidden_dims=self.hidden_dims,
                residual_type="densenet",
                residual_skip_steps=residual_skip_steps,
                layer_type=layer_type,
                **this_kwargs2,
                **self.pyg_kwargs,
            )

            graph_output_nn_kwargs = deepcopy(default_graph_output_nn_kwargs)
            graph_output_nn_kwargs["graph"]["pooling"] = pooling

            expectFailure = False

            if (not LAYERS_DICT[layer_type].layer_supports_edges) and (pre_nn_edges_kwargs is not None):
                expectFailure = True

            if (not LAYERS_DICT[layer_type].layer_supports_edges) and (
                task_heads_kwargs is not None and "task_edge" in task_heads_kwargs
            ):
                expectFailure = True

            if (
                "in_dim_edges" in this_kwargs2
                and this_kwargs2["in_dim_edges"] < 1
                and (task_heads_kwargs is not None and "task_edge" in task_heads_kwargs)
            ):
                expectFailure = True

            if task_heads_kwargs is not None and "task_graph" in task_heads_kwargs:
                expectFailure = True

            if expectFailure:
                with self.assertRaises(ValueError):
                    multitask_graph_nn = FullGraphMultiTaskNetwork(
                        gnn_kwargs=gnn_kwargs,
                        pre_nn_kwargs=pre_nn_kwargs,
                        pre_nn_edges_kwargs=pre_nn_edges_kwargs,
                        task_heads_kwargs=task_heads_kwargs,
                        graph_output_nn_kwargs=graph_output_nn_kwargs,
                    )
                continue

            try:
                multitask_graph_nn = FullGraphMultiTaskNetwork(
                    gnn_kwargs=gnn_kwargs,
                    pre_nn_kwargs=pre_nn_kwargs,
                    pre_nn_edges_kwargs=pre_nn_edges_kwargs,
                    task_heads_kwargs=task_heads_kwargs,
                    graph_output_nn_kwargs=graph_output_nn_kwargs,
                )

                batch_out = multitask_graph_nn.forward(bg)
            except Exception as e:
                exc_type, exc_value, exc_traceback = sys.exc_info()
                msg = err_msg + "\n" + str(traceback.format_exception(exc_type, exc_value, exc_traceback))
                self.fail(msg)

            if task_heads_kwargs is not None:
                if "task_node" in task_heads_kwargs:
                    self.assertListEqual(
                        list(batch_out["task_node"].shape),
                        [num_nodes, task_1_kwargs["out_dim"]],
                        msg=err_msg,
                    )
                if "task_edge" in task_heads_kwargs:
                    self.assertListEqual(
                        list(batch_out["task_edge"].shape),
                        [num_edges, task_2_kwargs["out_dim"]],
                        msg=err_msg,
                    )
                if "task_graph" in task_heads_kwargs:
                    self.assertListEqual(
                        list(batch_out["task_graph"].shape),
                        [num_graphs, task_3_kwargs["out_dim"]],
                        msg=err_msg,
                    )
                if "task_nodepair" in task_heads_kwargs:
                    self.assertListEqual(
                        list(batch_out["task_nodepair"].shape),
                        [num_nodepairs, task_4_kwargs["out_dim"]],
                        msg=err_msg,
                    )
            else:
                feat_out = batch_out["feat"]
                edge_feat_out = batch_out["edge_feat"]

                out_dim_edges = (
                    multitask_graph_nn.gnn.full_dims_edges[-1]
                    if multitask_graph_nn.gnn.full_dims_edges is not None
                    else expected_dim_edges
                )

                if not LAYERS_DICT[layer_type].layer_supports_edges or this_kwargs2["in_dim_edges"] == 0:
                    self.assertEqual(multitask_graph_nn.out_dim_edges, 0)
                else:
                    self.assertEqual(multitask_graph_nn.out_dim_edges, out_dim_edges)
                self.assertListEqual(list(feat_out.shape), [num_nodes, temp_dim_2], msg=err_msg)
                self.assertListEqual(list(edge_feat_out.shape), [num_edges, out_dim_edges], msg=err_msg)

    def test_full_graph_multi_task_from_config(self):
        cfg = graphium.load_config(name="zinc_default_multitask_pyg")

        # Initialize the network
        in_dims = {"feat": self.in_dim, "edge_feat": self.in_dim_edges}
        model_class, model_kwargs = load_architecture(cfg, in_dims=in_dims)

        multitask_full_graph_nn = model_class(**model_kwargs)

        # Test
        bg, num_nodes, num_edges, num_graphs, num_nodepairs = toy_test_data(
            in_dim=self.in_dim, in_dim_edges=self.in_dim_edges
        )
        batch_out = multitask_full_graph_nn.forward(bg)

        self.assertListEqual(
            list(batch_out["task_1"].shape),
            [num_nodes, task_1_kwargs["out_dim"]],
        )
        self.assertListEqual(
            list(batch_out["task_2"].shape),
            [num_edges, task_2_kwargs["out_dim"]],
        )
        self.assertListEqual(
            list(batch_out["task_3"].shape),
            [num_graphs, task_3_kwargs["out_dim"]],
        )
        self.assertListEqual(
            list(batch_out["task_4"].shape),
            [num_graphs, num_nodepairs, task_4_kwargs["out_dim"]],
        )

    def test_full_graph_multi_task_set_max_num_nodes(self):
        cfg = graphium.load_config(name="zinc_default_multitask_pyg")

        # Initialize the network
        in_dims = {"feat": self.in_dim, "edge_feat": self.in_dim_edges}
        model_class, model_kwargs = load_architecture(cfg, in_dims=in_dims)

        multitask_full_graph_nn: FullGraphMultiTaskNetwork = model_class(**model_kwargs)
        multitask_full_graph_nn.set_max_num_nodes_edges_per_graph(10, 4)

        # Probing
        self.assertEqual(
            multitask_full_graph_nn.task_heads.graph_output_nn["nodepair"].max_num_nodes_per_graph, 10
        )
        self.assertEqual(multitask_full_graph_nn.gnn.layers[2].max_num_edges_per_graph, 4)


if __name__ == "__main__":
    ut.main()


# tests/test_multitask_datamodule.py

import shutil
import tempfile
import unittest as ut

import pytest
from omegaconf import OmegaConf
import pandas as pd
import numpy as np
import graphium


class Test_Multitask_DataModule(ut.TestCase):
    def setUp(self):
        # Create a temporary directory
        self.tmp_test_dir = tempfile.mkdtemp()

    def tearDown(self):
        # Remove the directory after the test
        shutil.rmtree(self.tmp_test_dir)

    def test_multitask_fromsmiles_dm(
        self,
    ):  # TODO: I think we can remove this as it tests tiny_zinc which only contain graph level labels
        """Cover similar testing as for the original data module."""
        df = graphium.data.load_tiny_zinc()  # 100 molecules

        # Here we take the microzinc dataset and split the labels up into 'SA', 'logp' and 'score' in order to simulate having multiple single-task datasets
        df_micro_zinc_SA = df[["SMILES", "SA"]]
        df_micro_zinc_logp = df[["SMILES", "logp"]]
        df_micro_zinc_score = df[["SMILES", "score"]]

        # Setup the featurization. This will be the same across all tasks.
        featurization_args = {}
        featurization_args["atom_property_list_float"] = []  # ["weight", "valence"]
        featurization_args["atom_property_list_onehot"] = ["atomic-number", "degree"]
        featurization_args["edge_property_list"] = ["in-ring", "bond-type-onehot"]
        featurization_args["add_self_loop"] = False
        featurization_args["use_bonds_weights"] = False
        featurization_args["explicit_H"] = False

        # Config for multitask datamodule.

        # Per-task arguments.
        dm_task_args_SA = {}
        dm_task_args_SA["df"] = df_micro_zinc_SA
        dm_task_args_SA["task_level"] = "graph"
        dm_task_args_SA["smiles_col"] = "SMILES"
        dm_task_args_SA["label_cols"] = ["SA"]
        dm_task_args_SA["split_val"] = 0.2
        dm_task_args_SA["split_test"] = 0.2
        dm_task_args_SA["seed"] = 19
        dm_task_args_SA["splits_path"] = None  # This may not always be provided
        dm_task_args_SA["sample_size"] = None  # This may not always be provided
        dm_task_args_SA["idx_col"] = None  # This may not always be provided
        dm_task_args_SA["weights_col"] = None  # This may not always be provided
        dm_task_args_SA["weights_type"] = None  # This may not always be provided

        dm_task_args_logp = {}
        dm_task_args_logp["df"] = df_micro_zinc_logp
        dm_task_args_logp["task_level"] = "graph"
        dm_task_args_logp["smiles_col"] = "SMILES"
        dm_task_args_logp["label_cols"] = ["logp"]
        dm_task_args_logp["split_val"] = 0.2
        dm_task_args_logp["split_test"] = 0.2
        dm_task_args_logp["seed"] = 19
        dm_task_args_logp["splits_path"] = None  # This may not always be provided
        dm_task_args_logp["sample_size"] = None  # This may not always be provided
        dm_task_args_logp["idx_col"] = None  # This may not always be provided
        dm_task_args_logp["weights_col"] = None  # This may not always be provided
        dm_task_args_logp["weights_type"] = None  # This may not always be provided

        dm_task_args_score = {}
        dm_task_args_score["df"] = df_micro_zinc_score
        dm_task_args_score["task_level"] = "graph"
        dm_task_args_score["smiles_col"] = "SMILES"
        dm_task_args_score["label_cols"] = ["score"]
        dm_task_args_score["split_val"] = 0.2
        dm_task_args_score["split_test"] = 0.2
        dm_task_args_score["seed"] = 19
        dm_task_args_score["splits_path"] = None  # This may not always be provided
        dm_task_args_score["sample_size"] = None  # This may not always be provided
        dm_task_args_score["idx_col"] = None  # This may not always be provided
        dm_task_args_score["weights_col"] = None  # This may not always be provided
        dm_task_args_score["weights_type"] = None  # This may not always be provided

        dm_task_kwargs = {}
        dm_task_kwargs["SA"] = dm_task_args_SA
        dm_task_kwargs["logp"] = dm_task_args_logp
        dm_task_kwargs["score"] = dm_task_args_score

        dm_args = {}

        # Task-specific arguments for the datamodule
        dm_args["task_specific_args"] = dm_task_kwargs

        # Task-independent arguments
        dm_args["featurization"] = featurization_args
        dm_args["featurization_n_jobs"] = 16
        dm_args["featurization_progress"] = True
        dm_args["featurization_backend"] = "loky"
        dm_args["num_workers"] = 0
        dm_args["pin_memory"] = True
        dm_args["cache_data_path"] = None
        dm_args["batch_size_training"] = 16
        dm_args["batch_size_inference"] = 16

        # Create the data module
        dm = graphium.data.MultitaskFromSmilesDataModule(**dm_args)

        # self.assertEqual(50, dm.num_node_feats)    # Not implemeneted error
        # self.assertEqual(6, dm.num_edge_feats)

        dm.prepare_data()
        dm.setup()

        # self.assertEqual(len(dm), 100)                      # Should this have a fixed value for when it's initialized? MTL dataset only gets created after.
        self.assertEqual(len(dm.train_ds), 60)  # type: ignore
        self.assertEqual(len(dm.val_ds), 20)  # type: ignore
        self.assertEqual(len(dm.test_ds), 20)  # type: ignore
        # assert dm.num_node_feats == 50
        # assert dm.num_edge_feats == 6

        for dl in [dm.train_dataloader(), dm.val_dataloader(), dm.test_dataloader()]:
            it = iter(dl)
            batch = next(it)

            assert set(batch.keys()) == {"labels", "features"}

            # assert batch["labels"].shape == (16, 1)            # Single-task case
            assert batch["labels"]["graph_SA"].shape == (16, 1)
            assert batch["labels"]["graph_logp"].shape == (16, 1)
            assert batch["labels"]["graph_score"].shape == (16, 1)

    def test_multitask_fromsmiles_from_config(self):
        config = graphium.load_config(name="zinc_default_multitask_pyg")

        df = graphium.data.load_tiny_zinc()  # 100 molecules

        # Here we take the microzinc dataset and split the labels up into 'SA', 'logp' and 'score' in order to simulate having multiple single-task datasets
        df_micro_zinc_SA = df[["SMILES", "SA"]]
        df_micro_zinc_logp = df[["SMILES", "logp"]]
        df_micro_zinc_score = df[["SMILES", "score"]]

        # dm_args = dict(config.datamodule.args)
        dm_args = OmegaConf.to_container(config.datamodule.args, resolve=True)
        # dm_args["task_specific_args"]["SA"]["df"] = df
        dm_args["task_specific_args"]["SA"]["df"] = df_micro_zinc_SA
        dm_args["task_specific_args"]["logp"]["df"] = df_micro_zinc_logp
        dm_args["task_specific_args"]["score"]["df"] = df_micro_zinc_score

        dm_args["task_specific_args"]["SA"]["smiles_col"] = "SMILES"
        dm_args["task_specific_args"]["logp"]["smiles_col"] = "SMILES"
        dm_args["task_specific_args"]["score"]["smiles_col"] = "SMILES"

        dm_args["task_specific_args"]["SA"]["label_cols"] = ["SA"]
        dm_args["task_specific_args"]["logp"]["label_cols"] = ["logp"]
        dm_args["task_specific_args"]["score"]["label_cols"] = ["score"]

        dm_args["task_specific_args"]["SA"]["df_path"] = None
        dm_args["task_specific_args"]["logp"]["df_path"] = None
        dm_args["task_specific_args"]["score"]["df_path"] = None

        dm = graphium.data.MultitaskFromSmilesDataModule(**dm_args)

        # assert dm.num_node_feats == 50
        # assert dm.num_edge_feats == 6

        dm.prepare_data()
        dm.setup()

        # self.assertEqual(len(dm), 100)                      # Should this have a fixed value for when it's initialized? MTL dataset only gets created after.
        self.assertEqual(len(dm.train_ds), 60)  # type: ignore
        self.assertEqual(len(dm.val_ds), 20)  # type: ignore
        self.assertEqual(len(dm.test_ds), 20)  # type: ignore
        # assert dm.num_node_feats == 50
        # assert dm.num_edge_feats == 6

        for dl in [dm.train_dataloader(), dm.val_dataloader(), dm.test_dataloader()]:
            it = iter(dl)
            batch = next(it)

            assert set(batch.keys()) == {"labels", "features"}

            # assert batch["labels"].shape == (16, 1)            # Single-task case
            assert batch["labels"]["graph_SA"].shape == (16, 1)
            assert batch["labels"]["graph_logp"].shape == (16, 1)
            assert batch["labels"]["graph_score"].shape == (16, 1)

    def test_multitask_fromsmiles_from_config_csv(self):
        config = graphium.load_config(name="zinc_default_multitask_pyg")

        dm_args = OmegaConf.to_container(config.datamodule.args, resolve=True)
        dm = graphium.data.MultitaskFromSmilesDataModule(**dm_args)

        dm.prepare_data()
        dm.setup()

        # self.assertEqual(len(dm), 100)                      # Should this have a fixed value for when it's initialized? MTL dataset only gets created after.
        self.assertEqual(len(dm.train_ds), 60)  # type: ignore
        self.assertEqual(len(dm.val_ds), 20)  # type: ignore
        self.assertEqual(len(dm.test_ds), 20)  # type: ignore
        # assert dm.num_node_feats == 50
        # assert dm.num_edge_feats == 6

        for dl in [dm.train_dataloader(), dm.val_dataloader(), dm.test_dataloader()]:
            it = iter(dl)
            batch = next(it)

            assert set(batch.keys()) == {"labels", "features"}

            # assert batch["labels"].shape == (16, 1)            # Single-task case
            assert batch["labels"]["graph_SA"].shape == (16, 1)
            assert batch["labels"]["graph_logp"].shape == (16, 1)
            assert batch["labels"]["graph_score"].shape == (16, 1)

    def test_multitask_fromsmiles_from_config_parquet(self):
        config = graphium.load_config(name="fake_multilevel_multitask_pyg")

        dm_args = OmegaConf.to_container(config.datamodule.args, resolve=True)
        dm = graphium.data.MultitaskFromSmilesDataModule(**dm_args)

        dm.prepare_data()
        dm.setup()

        self.assertEqual(len(dm.train_ds), 1004)  # type: ignore

        dl = dm.train_dataloader()
        it = iter(dl)
        batch = next(it)

        assert set(batch.keys()) == {"labels", "features"}

        # assert batch["labels"].shape == (16, 1)            # Single-task case
        assert batch["labels"]["graph_SA"].shape == (16, 1)
        assert batch["labels"]["node_logp"].shape == (
            batch["features"].feat.size(0),
            2,
        )  # test node level
        assert batch["labels"]["edge_score"].shape == (
            batch["features"].edge_feat.size(0),
            2,
        )  # test edge level

    def test_multitask_with_missing_fromsmiles_from_config_parquet(self):
        config = graphium.load_config(name="fake_and_missing_multilevel_multitask_pyg")

        dm_args = OmegaConf.to_container(config.datamodule.args, resolve=True)
        dm = graphium.data.MultitaskFromSmilesDataModule(**dm_args)

        dm.prepare_data()
        dm.setup()

        self.assertEqual(len(dm.train_ds), 1004)  # type: ignore

        dl = dm.train_dataloader()
        it = iter(dl)
        batch = next(it)

        assert set(batch.keys()) == {"labels", "features"}

        # assert batch["labels"].shape == (16, 1)            # Single-task case
        assert batch["labels"]["graph_SA"].shape == (16, 1)
        assert batch["labels"]["node_logp"].shape == (
            batch["features"].feat.size(0),
            2,
        )  # test node level
        assert batch["labels"]["edge_score"].shape == (
            batch["features"].edge_feat.size(0),
            2,
        )  # test edge level

    def test_extract_graph_level_singletask(self):
        df = pd.read_parquet(f"tests/converted_fake_multilevel_data.parquet")
        num_graphs = len(df)
        label_cols = ["graph_label"]
        output = graphium.data.datamodule.extract_labels(df, "graph", label_cols)

        assert isinstance(output, np.ndarray)
        assert len(output.shape) == 2
        assert output.shape[0] == num_graphs
        assert output.shape[1] == 1

    def test_extract_graph_level_multitask(self):
        df = pd.read_parquet(f"tests/converted_fake_multilevel_data.parquet")
        num_graphs = len(df)
        label_cols = ["graph_label", "graph_label"]
        output = graphium.data.datamodule.extract_labels(df, "graph", label_cols)

        assert isinstance(output, np.ndarray)
        assert len(output.shape) == 2
        assert output.shape[0] == num_graphs
        assert output.shape[1] == len(label_cols)

    def test_extract_graph_level_multitask_missing_cols(self):
        df = pd.read_parquet(f"tests/converted_fake_multilevel_data.parquet")
        num_graphs = len(df)
        label_cols = ["graph_label", "graph_label"]

        drop_index = [2, 5, 21, 237, 192, 23, 127, 11]
        for replace in [1, 2]:
            for missing_col in label_cols[:replace]:
                df[missing_col].iloc[drop_index] = None

            output = graphium.data.datamodule.extract_labels(df, "graph", label_cols)

            assert isinstance(output, np.ndarray)
            assert len(output.shape) == 2
            assert output.shape[0] == num_graphs
            assert output.shape[1] == len(label_cols)

    def test_non_graph_level_extract_labels(self):
        df = pd.read_parquet(f"tests/converted_fake_multilevel_data.parquet")

        for level in ["node", "edge", "nodepair"]:
            label_cols = [f"{level}_label_{suffix}" for suffix in ["list", "np"]]
            output = graphium.data.datamodule.extract_labels(df, level, label_cols)

            assert isinstance(output, list)
            assert len(output[0].shape) == 2
            assert output[0].shape[1] == len(label_cols)

    def test_non_graph_level_extract_labels_missing_cols(self):
        df = pd.read_parquet(f"tests/converted_fake_multilevel_data.parquet")

        for level in ["node", "edge", "nodepair"]:
            label_cols = [f"{level}_label_{suffix}" for suffix in ["list", "np"]]
            drop_index = [2, 5, 21, 237, 192, 23, 127, 11]
            for replace in [1, 2]:
                for missing_col in label_cols[:replace]:
                    df.loc[drop_index, missing_col] = None

                output = graphium.data.datamodule.extract_labels(df, level, label_cols)

                for idx in drop_index:
                    assert len(output[idx].shape) == 2
                    assert output[idx].shape[1] == len(label_cols)

                    # Check that number of labels is adjusted correctly
                    if replace == 1:
                        non_missing_col = label_cols[1]
                        assert output[idx].shape[0] == len(df[non_missing_col][idx])

    def test_tdc_admet_benchmark_data_module(self):
        """
        Verifies that the ADMET-specific subclass of the MultiTaskDataModule works.
        Checks if all main endpoints can be run and if the split is correct.
        """

        try:
            from tdc.benchmark_group import admet_group
            from tdc.utils import retrieve_benchmark_names
        except ImportError:
            self.skipTest("PyTDC needs to be installed to run this test. Use `pip install PyTDC`.")
            raise

        # Make sure we can initialize the module and run the main endpoints
        data_module = graphium.data.ADMETBenchmarkDataModule()
        data_module.prepare_data()
        data_module.setup()

        for dl in [
            data_module.train_dataloader(),
            data_module.val_dataloader(),
            data_module.test_dataloader(),
        ]:
            batch = next(iter(dl))
            assert set(batch.keys()) == {"labels", "features"}

        # # Validate the split
        group = admet_group(path=self.tmp_test_dir)
        benchmark_names = retrieve_benchmark_names("admet_group")

        # For each of the endpoints...
        for name in benchmark_names:
            # Get the split from the benchmark group (ground truth)
            benchmark = group.get(name)
            train, val = group.get_train_valid_split(0, name)
            test = benchmark["test"]

            # Get the split from the data module
            params = data_module._get_task_specific_arguments(name, 0, self.tmp_test_dir)
            split = pd.read_csv(params.splits_path)
            data = params.df

            # Check that the split is the same
            for ground_truth, label in [(train, "train"), (val, "val"), (test, "test")]:
                y_true = ground_truth["Y"].values
                y_module = data.loc[split[label].dropna(), "Y"].values

                assert len(y_true) == len(y_module)
                assert np.allclose(np.sort(y_true), np.sort(y_module))


if __name__ == "__main__":
    ut.main()


# tests/test_mup.py

"""
Unit tests for the implementation of mup
"""

import unittest as ut
from copy import deepcopy
import torch.nn as nn
import torch
import yaml

from torch_geometric.data import Batch, Data

from graphium.nn.architectures import FeedForwardNN, FeedForwardPyg, FullGraphMultiTaskNetwork


def get_pyg_graphs(in_dim, in_dim_edges):
    edge_idx1 = torch.stack([torch.tensor([0, 1, 2, 3, 2]), torch.tensor([1, 2, 3, 0, 0])])
    edge_idx2 = torch.stack([torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])])
    x1 = torch.randn(edge_idx1.max() + 1, in_dim, dtype=torch.float32)
    e1 = torch.randn(edge_idx1.shape[-1], in_dim_edges, dtype=torch.float32)
    x2 = torch.randn(edge_idx2.max() + 1, in_dim, dtype=torch.float32)
    e2 = torch.randn(edge_idx2.shape[-1], in_dim_edges, dtype=torch.float32)
    g1 = Data(feat=x1, edge_index=edge_idx1, edge_feat=e1)
    g2 = Data(feat=x2, edge_index=edge_idx2, edge_feat=e2)
    bg = Batch.from_data_list([g1, g2])

    return bg


class test_mup(ut.TestCase):
    kwargs = dict(
        in_dim=12,
        out_dim=60,
        hidden_dims=8 * [84],
        depth=None,
        activation="LeakyReLU",
        last_activation="LeakyReLU",
        dropout=0.1,
        last_dropout=0.2,
        normalization="batch_norm",
        first_normalization="batch_norm",
        last_normalization="batch_norm",
        residual_type="simple",
        residual_skip_steps=2,
        name="testing",
        layer_type="fc",
        layer_kwargs=None,
    )

    def test_feedforwardnn_mup(self):
        kwargs = deepcopy(self.kwargs)
        model = FeedForwardNN(**kwargs, last_layer_is_readout=False)
        model_lastreadout = FeedForwardNN(**kwargs, last_layer_is_readout=True)
        base_1 = model.make_mup_base_kwargs(divide_factor=1)
        base_1_lastreadout = model_lastreadout.make_mup_base_kwargs(divide_factor=1)

        base_2 = model.make_mup_base_kwargs(divide_factor=2)
        base_2_lastreadout = model_lastreadout.make_mup_base_kwargs(divide_factor=2)
        kwargs_2 = deepcopy(base_1)
        kwargs_2.update(dict(out_dim=30, hidden_dims=8 * [42]))
        kwargs_2_lastreadout = deepcopy(base_1_lastreadout)
        kwargs_2_lastreadout.update(dict(hidden_dims=8 * [42]))

        # Check the kwargs matching
        for key in kwargs_2.keys():
            if isinstance(kwargs_2[key], nn.Module):
                # Can't match the random weights
                self.assertEqual(str(kwargs_2[key]), str(base_2[key]), msg=key)
            else:
                self.assertEqual(kwargs_2[key], base_2[key], msg=key)

        # Check the kwargs matching
        for key in kwargs_2_lastreadout.keys():
            if isinstance(kwargs_2_lastreadout[key], nn.Module):
                # Can't match the random weights
                self.assertEqual(str(kwargs_2_lastreadout[key]), str(base_2_lastreadout[key]), msg=key)
            else:
                self.assertEqual(kwargs_2_lastreadout[key], base_2_lastreadout[key], msg=key)

        # Test that the models with divide_factor=1 can be built run a forward pass
        in_features = torch.randn((10, kwargs["in_dim"]))
        model_1 = FeedForwardNN(**base_1)
        model_1.forward(deepcopy(in_features))
        model_1_lastreadout = FeedForwardNN(**base_1_lastreadout)
        model_1_lastreadout.forward(deepcopy(in_features))

        # Test that the models with divide_factor=2 can be built run a forward pass
        model_2 = FeedForwardNN(**base_2)
        model_2.forward(deepcopy(in_features))
        model_2_lastreadout = FeedForwardNN(**base_2_lastreadout)
        model_2_lastreadout.forward(deepcopy(in_features))

    def test_feedforwardgraph_mup(self):
        kwargs = deepcopy(self.kwargs)
        in_dim_edges = kwargs["in_dim"]
        kwargs.update(dict(layer_type="pyg:gine", in_dim_edges=in_dim_edges))
        model = FeedForwardPyg(**kwargs, last_layer_is_readout=False)
        model_lastreadout = FeedForwardPyg(**kwargs, last_layer_is_readout=True)
        base_1 = model.make_mup_base_kwargs(divide_factor=1)
        base_1_lastreadout = model_lastreadout.make_mup_base_kwargs(divide_factor=1)

        base_2 = model.make_mup_base_kwargs(divide_factor=2)
        base_2_lastreadout = model_lastreadout.make_mup_base_kwargs(divide_factor=2)
        kwargs_2 = deepcopy(base_1)
        kwargs_2.update(dict(out_dim=30, hidden_dims=8 * [42]))
        kwargs_2_lastreadout = deepcopy(base_1_lastreadout)
        kwargs_2_lastreadout.update(dict(hidden_dims=8 * [42]))

        # Check the kwargs matching
        for key in kwargs_2.keys():
            if isinstance(kwargs_2[key], nn.Module):
                # Can't match the random weights
                self.assertEqual(str(kwargs_2[key]), str(base_2[key]), msg=key)
            else:
                self.assertEqual(kwargs_2[key], base_2[key], msg=key)

        # Check the kwargs matching
        for key in kwargs_2_lastreadout.keys():
            if isinstance(kwargs_2_lastreadout[key], nn.Module):
                # Can't match the random weights
                self.assertEqual(str(kwargs_2_lastreadout[key]), str(base_2_lastreadout[key]), msg=key)
            else:
                self.assertEqual(kwargs_2_lastreadout[key], base_2_lastreadout[key], msg=key)

        # Test that the models with divide_factor=1 can be built run a forward pass
        in_features = get_pyg_graphs(in_dim=kwargs["in_dim"], in_dim_edges=in_dim_edges)
        model_1 = FeedForwardPyg(**base_1)
        model_1.forward(deepcopy(in_features))
        model_1_lastreadout = FeedForwardPyg(**base_1_lastreadout)
        model_1_lastreadout.forward(deepcopy(in_features))

        # Test that the models with divide_factor=2 can be built run a forward pass
        model_2 = FeedForwardPyg(**base_2)
        model_2.forward(deepcopy(in_features))
        model_2_lastreadout = FeedForwardPyg(**base_2_lastreadout)
        model_2_lastreadout.forward(deepcopy(in_features))

    def test_fullgraphmultitasknetwork(self):
        # Load the configuration file for the model
        CONFIG_FILE = "tests/config_test_ipu_dataloader.yaml"
        with open(CONFIG_FILE, "r") as f:
            cfg = yaml.safe_load(f)

        # Make fake graphs
        in_dim = 12
        in_dim_edges = 12
        pe_indims = {
            "laplacian_eigvec": 5,
            "laplacian_eigval": 5,
            "rw_return_probs": 2,
            "edge_rw_transition_probs": 2,
            "nodepair_rw_return_probs": 1,
            "electrostatic": 6,
            "edge_commute": 1,
            "nodepair_graphormer": 1,
            "positions_3d": 3,
        }
        in_features = get_pyg_graphs(in_dim=in_dim, in_dim_edges=in_dim_edges)
        for key, dim in pe_indims.items():
            if key.startswith("edge_"):
                in_features[key] = torch.randn(in_features.num_edges, dim)
            elif key.startswith("nodepair_"):
                in_features[key] = torch.randn(in_features.num_nodes, in_features.num_nodes, dim)
            else:
                in_features[key] = torch.randn(in_features.num_nodes, dim)

        # Load the model
        kwargs = {}
        for key, val in cfg["architecture"].items():
            if key in ["model_type", "mup_base_path"]:
                continue
            kwargs[key + "_kwargs"] = val
        kwargs["pre_nn_kwargs"]["in_dim"] = in_dim + kwargs["pe_encoders_kwargs"]["out_dim"]
        kwargs["pre_nn_edges_kwargs"]["in_dim"] = in_dim_edges + kwargs["pe_encoders_kwargs"]["edge_out_dim"]
        kwargs["pe_encoders_kwargs"]["in_dims"] = pe_indims

        model = FullGraphMultiTaskNetwork(**kwargs, last_layer_is_readout=True)

        kw_1 = model.make_mup_base_kwargs(divide_factor=1)
        kw_2 = model.make_mup_base_kwargs(divide_factor=2)

        # Check the parameter sizes
        for key, elem in kw_1.items():
            if not isinstance(elem, dict):
                continue
            for subkey, subelem in elem.items():
                if "dim" in subkey:
                    match = f"{key}:{subkey}"
                    if match in ["pre_nn_kwargs:in_dim"]:
                        self.assertNotEqual(subelem, kw_2[key][subkey], msg=match)
                    elif match in [
                        "pre_nn_kwargs:out_dim",
                        "pre_nn_edges_kwargs:out_dim",
                        "gnn_kwargs:in_dim",
                        "graph_output_nn_kwargs:in_dim",
                        "gnn_kwargs:out_dim",
                        "gnn_kwargs:in_dim_edges",
                        "pe_encoders_kwargs:out_dim",
                        "graph_output_nn_kwargs:out_dim",
                    ]:
                        # Divide by 2
                        self.assertEqual(round(subelem / 2), kw_2[key][subkey], msg=match)
                    elif match in [
                        "pre_nn_kwargs:hidden_dims",
                        "pre_nn_edges_kwargs:hidden_dims",
                        "gnn_kwargs:hidden_dims",
                        "graph_output_nn_kwargs:hidden_dims",
                        "gnn_kwargs:hidden_dims_edges",
                    ]:
                        # Arrays divide by 2
                        new_list = [round(e / 2) for e in subelem]
                        self.assertListEqual(new_list, kw_2[key][subkey], msg=match)
                elif subkey in ["homo", "alpha", "cv"]:
                    for subsubkey, subsubelem in subelem.items():
                        match = f"{key}:{subsubkey}"
                        if match in [
                            "task_heads_kwargs:out_dim",
                            "task_heads_kwargs:out_dim",
                            "task_heads_kwargs:out_dim",
                        ]:
                            # No divide
                            self.assertEqual(subsubelem, kw_2[key][subkey][subsubkey], msg=match)
                        elif match in [
                            "task_heads_kwargs:in_dim",
                            "task_heads_kwargs:in_dim",
                            "task_heads_kwargs:in_dim",
                        ]:
                            # Divide by 2
                            self.assertEqual(round(subsubelem / 2), kw_2[key][subkey][subsubkey], msg=match)
                        elif match in [
                            "task_heads_kwargs:hidden_dims",
                            "task_heads_kwargs:hidden_dims",
                            "task_heads_kwargs:hidden_dims",
                        ]:
                            # Divide by 2 a list
                            new_list = [round(e / 2) for e in subsubelem]
                            self.assertListEqual(new_list, kw_2[key][subkey][subsubkey], msg=match)

        # Test that the models with divide_factor=1 can be built run a forward pass
        kw_1["last_layer_is_readout"] = False
        model_1 = FullGraphMultiTaskNetwork(**kw_1)
        model_1.forward(deepcopy(in_features))
        kw_1["last_layer_is_readout"] = True
        model_1 = FullGraphMultiTaskNetwork(**kw_1)
        model_1.forward(deepcopy(in_features))

        # Test that the models with divide_factor=2 can be built run a forward pass
        kw_2["last_layer_is_readout"] = False
        model_2 = FullGraphMultiTaskNetwork(**kw_2)
        model_2.forward(deepcopy(in_features))
        kw_2["last_layer_is_readout"] = True
        model_2 = FullGraphMultiTaskNetwork(**kw_2)
        model_2.forward(deepcopy(in_features))


if __name__ == "__main__":
    ut.main()


# tests/test_packing.py

# General imports
import unittest as ut
import numpy as np

import torch
from torch_geometric.data import Data, Batch

# Current library imports
from graphium.utils.packing import (
    smart_packing,
    get_pack_sizes,
    fast_packing,
    hybrid_packing,
    node_to_pack_indices_mask,
)


def random_packing(num_nodes, batch_size):
    ipu_batch_size = int(len(num_nodes) / batch_size)
    indices = np.arange(len(num_nodes))
    np.random.shuffle(indices)
    indices = np.reshape(indices, (ipu_batch_size, batch_size)).tolist()
    return indices


class test_Packing(ut.TestCase):
    def test_smart_packing(self):
        np.random.seed(42)

        batch_sizes = [2, 4, 8, 16, 32, 64]
        ipu_batch_sizes = [2, 3, 4, 8, 16, 32, 64]

        for batch_size in batch_sizes:
            for ipu_batch_size in ipu_batch_sizes:
                err_msg = f"bz={batch_size}, ipu_bz={ipu_batch_size}"

                # Generate random batch size
                global_batch = batch_size * ipu_batch_size
                num_nodes = np.abs(np.random.gamma(2, 20, size=global_batch)).astype(int)

                # Use the smart packing
                packed_indices = smart_packing(num_nodes=num_nodes, batch_size=batch_size)
                pack_num_nodes = get_pack_sizes(packed_indices, num_nodes)

                # Use the random packing
                rand_packed_indices = random_packing(num_nodes=num_nodes, batch_size=batch_size)
                rand_pack_num_nodes = get_pack_sizes(rand_packed_indices, num_nodes)

                # Assert that the smart packing is better than the random packing
                self.assertLessEqual(max(pack_num_nodes), max(rand_pack_num_nodes), msg=err_msg)
                self.assertGreaterEqual(min(pack_num_nodes), min(rand_pack_num_nodes), msg=err_msg)

                # Assert that the total number of atoms is right
                self.assertEqual(sum(pack_num_nodes), sum(num_nodes), msg=err_msg)
                self.assertEqual(sum(rand_pack_num_nodes), sum(num_nodes), msg=err_msg)

                # Assert that all index are there
                self.assertListEqual(
                    np.sort(np.asarray(packed_indices).flatten()).tolist(), np.arange(len(num_nodes)).tolist()
                )
                self.assertListEqual(
                    np.sort(np.asarray(rand_packed_indices).flatten()).tolist(),
                    np.arange(len(num_nodes)).tolist(),
                )

    def test_fast_packing(self):
        np.random.seed(42)

        # Start at 4 for fast_packing for better statistical significance
        batch_sizes = [4, 8, 16, 32, 64]
        ipu_batch_sizes = [4, 8, 16, 32, 64]

        for batch_size in batch_sizes:
            for ipu_batch_size in ipu_batch_sizes:
                err_msg = f"bz={batch_size}, ipu_bz={ipu_batch_size}"

                # Generate random batch size
                global_batch = batch_size * ipu_batch_size
                num_nodes = np.abs(np.random.gamma(2, 20, size=global_batch)).astype(int)

                # Use the smart packing
                packed_indices = fast_packing(num_nodes=num_nodes, batch_size=batch_size)
                pack_num_nodes = get_pack_sizes(packed_indices, num_nodes)

                # Use the random packing
                rand_packed_indices = random_packing(num_nodes=num_nodes, batch_size=batch_size)
                rand_pack_num_nodes = get_pack_sizes(rand_packed_indices, num_nodes)

                # Assert that the smart packing is better than the random packing
                self.assertLessEqual(max(pack_num_nodes), max(rand_pack_num_nodes), msg=err_msg)
                self.assertGreaterEqual(min(pack_num_nodes), min(rand_pack_num_nodes), msg=err_msg)

                # Assert that the total number of atoms is right
                self.assertEqual(sum(pack_num_nodes), sum(num_nodes), msg=err_msg)
                self.assertEqual(sum(rand_pack_num_nodes), sum(num_nodes), msg=err_msg)

                # Assert that all index are there
                self.assertListEqual(
                    np.sort(np.asarray(packed_indices).flatten()).tolist(), np.arange(len(num_nodes)).tolist()
                )
                self.assertListEqual(
                    np.sort(np.asarray(rand_packed_indices).flatten()).tolist(),
                    np.arange(len(num_nodes)).tolist(),
                )

    def test_hybrid_packing(self):
        np.random.seed(42)

        batch_sizes = [2, 4, 8, 16, 32, 64]
        ipu_batch_sizes = [2, 3, 4, 8, 16, 32, 64]

        for batch_size in batch_sizes:
            for ipu_batch_size in ipu_batch_sizes:
                err_msg = f"bz={batch_size}, ipu_bz={ipu_batch_size}"

                # Generate random batch size
                global_batch = batch_size * ipu_batch_size
                num_nodes = np.abs(np.random.gamma(2, 20, size=global_batch)).astype(int)

                # Use the smart packing
                packed_indices = hybrid_packing(num_nodes=num_nodes, batch_size=batch_size)
                pack_num_nodes = get_pack_sizes(packed_indices, num_nodes)

                # Use the random packing
                rand_packed_indices = random_packing(num_nodes=num_nodes, batch_size=batch_size)
                rand_pack_num_nodes = get_pack_sizes(rand_packed_indices, num_nodes)

                # Assert that the smart packing is better than the random packing
                self.assertLessEqual(max(pack_num_nodes), max(rand_pack_num_nodes), msg=err_msg)
                self.assertGreaterEqual(min(pack_num_nodes), min(rand_pack_num_nodes), msg=err_msg)

                # Assert that the total number of atoms is right
                self.assertEqual(sum(pack_num_nodes), sum(num_nodes), msg=err_msg)
                self.assertEqual(sum(rand_pack_num_nodes), sum(num_nodes), msg=err_msg)

                # Assert that all index are there
                self.assertListEqual(
                    np.sort(np.asarray(packed_indices).flatten()).tolist(), np.arange(len(num_nodes)).tolist()
                )
                self.assertListEqual(
                    np.sort(np.asarray(rand_packed_indices).flatten()).tolist(),
                    np.arange(len(num_nodes)).tolist(),
                )

    def test_node_to_pack_indices_mask(self):
        # Create a dummy batch
        in_dim = 7
        in_dim_edges = 11
        max_num_nodes_per_graph = 20
        batch_size_per_pack = 5

        torch.manual_seed(42)

        # Create a dummy batch of graphs
        batch, all_num_nodes = [], []
        for ii in range(100):
            num_nodes = torch.randint(1, max_num_nodes_per_graph, (1,)).item()
            all_num_nodes.append(num_nodes)
            num_edges = abs(round(2.2 * num_nodes) + torch.randint(-2, 2, (1,)).item()) + 1
            x = torch.randn(num_nodes, in_dim, dtype=torch.float32)
            edge_idx = torch.randint(0, num_nodes, (2, num_edges))
            e = torch.randn(edge_idx.shape[-1], in_dim_edges, dtype=torch.float32)
            g = Data(h=x, edge_index=edge_idx, edge_attr=e)
            batch.append(g)
        batch = Batch.from_data_list(batch)

        # Get the packing
        packed_graph_idx = fast_packing(all_num_nodes, batch_size_per_pack)
        pack_sizes = get_pack_sizes(packed_graph_idx, all_num_nodes)
        max_pack_size = max(pack_sizes)
        num_packs = len(pack_sizes)

        # Get the node to pack indices and the mask
        pack_from_node_idx, pack_attn_mask = node_to_pack_indices_mask(packed_graph_idx, all_num_nodes)

        # Assert that the nodes to pack indices are correct
        h = torch.arange(batch.num_nodes, dtype=torch.float32)
        packed_shape = [num_packs, max_pack_size]
        h_packed = torch.zeros(packed_shape)
        h_packed[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]] = h
        h_packed_unique = torch.sort(torch.unique(h_packed))[0]
        np.testing.assert_array_equal(h_packed_unique, torch.arange(batch.num_nodes))
        self.assertEqual(h_packed.sum(), h.sum())

        # Test again with additional h dimension
        h = batch.h
        packed_shape = [num_packs, max_pack_size] + list(h.shape[1:])
        h_packed = torch.zeros(packed_shape)
        h_packed[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]] = h
        h_packed_unique = torch.sort(torch.unique(h_packed))[0]
        h_packed_unique = h_packed_unique[h_packed_unique != 0]
        np.testing.assert_array_almost_equal(h_packed_unique, torch.unique(h))
        self.assertAlmostEqual(h_packed.sum().item(), h.sum().item(), places=3)

        # Assert that the mask is correct by counting the number of False values (the sum of squared number of nodes per pack)
        num_false = (~pack_attn_mask).sum([1, 2])
        num_expected = torch.as_tensor(
            [sum([all_num_nodes[graph_idx] ** 2 for graph_idx in pack]) for pack in packed_graph_idx]
        )
        np.testing.assert_array_equal(num_false, num_expected)

        # Assert that the mask is correct by counting the number of elements in each row and column
        num_expected = []
        for pack in packed_graph_idx:
            pack_num_expected = []
            for graph_idx in pack:
                num_nodes = all_num_nodes[graph_idx]
                for ii in range(num_nodes):
                    pack_num_expected.append(num_nodes)
            pack_num_expected.extend([0] * (max_pack_size - len(pack_num_expected)))
            num_expected.append(pack_num_expected)
        num_expected = torch.as_tensor(num_expected)
        num_false_row = (~pack_attn_mask).sum([2])
        num_false_col = (~pack_attn_mask).sum([1])
        np.testing.assert_array_equal(num_false_row, num_expected)
        np.testing.assert_array_equal(num_false_col, num_expected)


if __name__ == "__main__":
    ut.main()


# tests/test_pe_nodepair.py

"""
Unit tests for the positional encodings in graphium/features/*
"""

import numpy as np
import networkx as nx
import unittest as ut

from graphium.features.electrostatic import compute_electrostatic_interactions
from graphium.features.commute import compute_commute_distances
from graphium.features.graphormer import compute_graphormer_distances


class test_positional_encodings(ut.TestCase):
    # Test graphs
    adj_dict = {}
    max_dict = {}

    # 6-ring
    adj = np.asarray(
        [
            [0, 1, 0, 0, 0, 1],
            [1, 0, 1, 0, 0, 0],
            [0, 1, 0, 1, 0, 0],
            [0, 0, 1, 0, 1, 0],
            [0, 0, 0, 1, 0, 1],
            [1, 0, 0, 0, 1, 0],
        ]
    )
    adj_dict["6-ring"] = adj
    max_dict["6-ring"] = 3

    # 5-path
    G = nx.path_graph(5)
    adj = nx.to_numpy_array(G)
    adj_dict["5-path"] = adj
    max_dict["5-path"] = 4

    # 4-clique
    adj = 1 - np.eye(4)
    adj_dict["4-clique"] = adj
    max_dict["4-clique"] = 1

    # 4-barbell
    H = nx.barbell_graph(4, 0)
    adj = nx.to_numpy_array(H)
    adj_dict["4-barbell"] = adj
    max_dict["4-barbell"] = 3

    def test_dimensions(self):
        for _, adj in self.adj_dict.items():
            pe, _, _ = compute_electrostatic_interactions(adj, cache={})
            self.assertEqual(pe.shape, adj.shape)

            pe, _, _ = compute_graphormer_distances(adj, adj.shape[0], cache={})
            self.assertEqual(pe.shape, adj.shape)

            pe, _, _ = compute_commute_distances(adj, adj.shape[0], cache={})
            self.assertEqual(pe.shape, adj.shape)

    def test_symmetry(self):
        for _, adj in self.adj_dict.items():
            pe, _, _ = compute_graphormer_distances(adj, adj.shape[0], cache={})
            np.testing.assert_array_almost_equal(pe, pe.T)

            pe, _, _ = compute_commute_distances(adj, adj.shape[0], cache={})
            np.testing.assert_array_almost_equal(pe, pe.T)

    def test_max_dist(self):
        for key, adj in self.adj_dict.items():
            pe, _, _ = compute_graphormer_distances(adj, adj.shape[0], cache={})
            np.testing.assert_array_almost_equal(pe.max(), self.max_dict[key])


if __name__ == "__main__":
    ut.main()


# tests/test_pe_rw.py

"""
Unit tests for the positional encodings in graphium/features/*
"""

import numpy as np
import networkx as nx
import unittest as ut

from graphium.features.rw import compute_rwse


class test_pe_spectral(ut.TestCase):
    def test_caching_and_outputs(self):
        # 4-barbell
        G = nx.barbell_graph(4, 0)
        adj = nx.to_numpy_array(G)
        num_nodes = adj.shape[0]
        cache = {}

        ksteps1 = [4, 6]
        ksteps2 = [2]
        ksteps3 = [6, 7]

        pe1, _, cache = compute_rwse(
            adj.astype(np.float32), ksteps1, num_nodes, cache, pos_type="rw_transition_probs"
        )

        pe2, _, cache = compute_rwse(
            adj.astype(np.float32), ksteps2, num_nodes, cache, pos_type="rw_return_probs"
        )

        pe3, _, cache = compute_rwse(
            adj.astype(np.float32), ksteps3, num_nodes, cache, pos_type="rw_return_probs"
        )

        self.assertTrue(all([k in cache["ksteps"] for k in ksteps1 + ksteps2 + ksteps3]))
        self.assertTrue(pe1.shape, np.zeros((num_nodes, num_nodes, len(ksteps1))))
        self.assertTrue(pe2.shape, np.zeros((num_nodes, len(ksteps2))))
        self.assertTrue(pe3.shape, np.zeros((num_nodes, len(ksteps3))))

        for i in range(len(ksteps1)):
            np.testing.assert_array_almost_equal(pe1[..., i].sum(1), np.ones(num_nodes))

        self.assertGreaterEqual(pe1.min(), 0.0)
        self.assertLessEqual(pe1.max(), 1.0)

        self.assertGreaterEqual(pe2.min(), 0.0)
        self.assertLessEqual(pe2.max(), 1.0)

        self.assertGreaterEqual(pe3.min(), 0.0)
        self.assertLessEqual(pe3.max(), 1.0)


if __name__ == "__main__":
    ut.main()


# tests/test_pe_spectral.py

"""
Unit tests for the positional encodings in graphium/features/*
"""

import numpy as np
import networkx as nx
import unittest as ut

from graphium.features.spectral import compute_laplacian_pe


class test_pe_spectral(ut.TestCase):
    # 2 disconnected 3 cliques
    adj1 = np.zeros((6, 6))
    adj_3clq = 1 - np.eye(3)
    adj1[:3, :3] = adj_3clq
    adj1[3:, 3:] = adj_3clq

    # 3-clique
    adj2 = 1 - np.eye(6)

    def test_for_connected_vs_disconnected_graph(self):
        num_pos = 3

        # test if pe works identically on connected vs disconnected graphs
        eigvals_pe1, _, _, cache = compute_laplacian_pe(self.adj1, num_pos, cache={})
        eigvals_pe1 = np.real(eigvals_pe1).astype(np.float32)
        _, eigvecs_pe1, _, _ = compute_laplacian_pe(self.adj1, num_pos, cache=cache)

        # We expect to cache 4 objects in when running the functon for the first time
        self.assertEqual(len(cache.keys()), 4)

        eigvals_pe2, _, _, _ = compute_laplacian_pe(self.adj2, num_pos, cache={})
        eigvals_pe2 = np.real(eigvals_pe2).astype(np.float32)
        _, eigvecs_pe2, _, _ = compute_laplacian_pe(self.adj2, num_pos, cache={})

        np.testing.assert_array_almost_equal(2 * eigvals_pe1, eigvals_pe2)
        self.assertListEqual(list(eigvals_pe2.shape), [self.adj2.shape[0], num_pos])
        self.assertListEqual(list(eigvecs_pe2.shape), [self.adj2.shape[0], num_pos])


if __name__ == "__main__":
    ut.main()


# tests/test_pos_transfer_funcs.py

"""
Unit tests for the positional encodings in graphium/features/*
"""

import numpy as np
import networkx as nx
import unittest as ut

from graphium.features.spectral import compute_laplacian_pe
from graphium.features.transfer_pos_level import (
    node_to_edge,
    node_to_nodepair,
    edge_to_nodepair,
    nodepair_to_node,
    nodepair_to_edge,
    graph_to_node,
)


class test_pos_transfer_funcs(ut.TestCase):
    # 4-barbell
    G = nx.barbell_graph(4, 0)
    adj = nx.to_numpy_array(G)
    num_nodes, num_feat = 8, 5
    node_pe = np.random.rand(num_nodes, num_feat)

    def test_different_pathways_from_node_to_edge(self):
        edge_pe1, _ = node_to_edge(self.node_pe, self.adj, {})
        nodepair_pe1 = node_to_nodepair(self.node_pe, self.num_nodes)
        edge_pe2, _ = nodepair_to_edge(nodepair_pe1, self.adj, {})
        nodepair_pe2, _ = edge_to_nodepair(edge_pe1, self.adj, self.num_nodes, {})
        edge_pe3, _ = nodepair_to_edge(nodepair_pe2, self.adj, {})
        np.testing.assert_array_almost_equal(edge_pe1, edge_pe2)
        np.testing.assert_array_almost_equal(edge_pe1, edge_pe3)


if __name__ == "__main__":
    ut.main()


# tests/test_positional_encoders.py

"""
Unit tests for the different datasets of graphium/features/featurizer.py
"""

import numpy as np
import unittest as ut
from copy import deepcopy
from rdkit import Chem
import datamol as dm
import torch
from scipy.sparse import coo_matrix

from graphium.features.featurizer import GraphDict
from graphium.features.positional_encoding import graph_positional_encoder
from graphium.nn.encoders import laplace_pos_encoder, mlp_encoder, signnet_pos_encoder

# TODO: Test the MLP_encoder and signnet_pos_encoder


class test_positional_encoder(ut.TestCase):
    smiles = [
        "C",
        "CC",
        "CC.CCCC",
        "CC(C)CC1=CC=C(C=C1)C(C)C(=O)O",
        "OCCc1c(C)[n+](cs1)Cc2cnc(C)nc2N",
        "O1C=C[C@H]([C@H]1O2)c3c2cc(OC)c4c3OC(=O)C5=C4CCC(=O)5",
    ]
    mols = [dm.to_mol(s) for s in smiles]
    adjs = [Chem.rdmolops.GetAdjacencyMatrix(mol) for mol in mols]

    def test_laplacian_eigvec_eigval(self):
        for ii, adj in enumerate(deepcopy(self.adjs)):
            for num_pos in [1, 2, 4]:  # Can't test too much eigs because of multiplicities
                for disconnected_comp in [True, False]:
                    err_msg = f"adj_id={ii}, num_pos={num_pos}, disconnected_comp={disconnected_comp}"

                    # returns a dictionary of computed pe
                    pos_kwargs = {
                        "pos_type": "laplacian_eigvec",
                        "num_pos": num_pos,
                        "disconnected_comp": disconnected_comp,
                        "pos_level": "node",
                    }
                    num_nodes = adj.shape[0]
                    eigvecs, cache = graph_positional_encoder(adj, num_nodes, pos_kwargs=pos_kwargs)
                    pos_kwargs["pos_type"] = "laplacian_eigval"
                    eigvals, cache = graph_positional_encoder(adj, num_nodes, pos_kwargs=pos_kwargs)

                    self.assertEqual(list(eigvecs.shape), [adj.shape[0], num_pos], msg=err_msg)
                    self.assertEqual(list(eigvals.shape), [adj.shape[0], num_pos], msg=err_msg)

                    # Compute eigvals and eigvecs
                    lap = np.diag(np.sum(adj, axis=1)) - adj
                    true_eigvals, true_eigvecs = np.linalg.eig(lap)
                    sort_idx = np.argsort(true_eigvals)
                    true_eigvals, true_eigvecs = true_eigvals[sort_idx], true_eigvecs[:, sort_idx]
                    true_eigvecs = true_eigvecs / (np.sum(true_eigvecs**2, axis=0, keepdims=True) + 1e-8)

                    true_num_pos = min(num_pos, len(true_eigvals))
                    true_eigvals, true_eigvecs = true_eigvals[:true_num_pos], true_eigvecs[:, :true_num_pos]

                    if not ("." in self.smiles[ii]):
                        np.testing.assert_array_almost_equal(
                            np.abs(true_eigvecs),
                            np.abs(eigvecs[:, :true_num_pos]),
                            decimal=6,
                            err_msg=err_msg,
                        )
                        self.assertAlmostEqual(np.sum(true_eigvecs[:, 1:]), 0, places=6, msg=err_msg)
                        np.testing.assert_array_almost_equal(
                            true_eigvals, eigvals[0, :true_num_pos], decimal=6, err_msg=err_msg
                        )

    # didn't actually check the exact computation result because the code was adapted
    def test_rwse(self):
        for ii, adj in enumerate(deepcopy(self.adjs)):
            for ksteps in [1, 2, 4]:
                err_msg = f"adj_id={ii}, ksteps={ksteps}"

                num_nodes = adj.shape[0]
                pos_kwargs = {"pos_type": "rw_return_probs", "ksteps": ksteps, "pos_level": "node"}
                rwse_embed, cache = graph_positional_encoder(adj, num_nodes, pos_kwargs=pos_kwargs)
                self.assertEqual(list(rwse_embed.shape), [num_nodes, ksteps], msg=err_msg)

    # TODO: work in progress

    """
    continue debugging here, see how to adapt the laplace_pos_encoder
    code running now, question is where to add the laplace_pos_encoder
    """

    def test_laplacian_eigvec_with_encoder(self):
        for ii, adj in enumerate(deepcopy(self.adjs)):
            for num_pos in [2, 4, 8]:  # Can't test too much eigs because of multiplicities
                for disconnected_comp in [True, False]:
                    for model_type in ["Transformer", "DeepSet", "MLP"]:
                        err_msg = f"adj_id={ii}, num_pos={num_pos}, disconnected_comp={disconnected_comp}"

                        # returns a dictionary of computed pe
                        pos_kwargs = {
                            "pos_type": "laplacian_eigvec",
                            "num_pos": num_pos,
                            "disconnected_comp": disconnected_comp,
                            "pos_level": "node",
                        }
                        num_nodes = adj.shape[0]
                        eigvecs, cache = graph_positional_encoder(adj, num_nodes, pos_kwargs=pos_kwargs)
                        pos_kwargs["pos_type"] = "laplacian_eigval"
                        eigvals, cache = graph_positional_encoder(adj, num_nodes, pos_kwargs=pos_kwargs)

                        input_keys = ["laplacian_eigvec", "laplacian_eigval"]
                        in_dim = num_pos
                        hidden_dim = 64
                        out_dim = 64
                        num_layers = 1

                        eigvecs = torch.from_numpy(eigvecs)
                        eigvals = torch.from_numpy(eigvals)

                        g = GraphDict(
                            {
                                "adj": coo_matrix(adj),
                                "data": {"laplacian_eigval": eigvals, "laplacian_eigvec": eigvecs},
                            }
                        )
                        batch = g.make_pyg_graph()

                        encoder = laplace_pos_encoder.LapPENodeEncoder(
                            input_keys=input_keys,
                            output_keys=["node"],
                            in_dim=in_dim,  # Size of Laplace PE embedding
                            hidden_dim=hidden_dim,
                            out_dim=out_dim,
                            model_type=model_type,  # 'Transformer' or 'DeepSet'
                            num_layers=num_layers,
                            num_layers_post=2,  # Num. layers to apply after pooling
                            dropout=0.1,
                            first_normalization=None,
                        )

                        hidden_embed = encoder(batch, key_prefix=None)
                        assert "node" in hidden_embed.keys()
                        self.assertEqual(list(hidden_embed["node"].shape), [num_nodes, out_dim], msg=err_msg)


if __name__ == "__main__":
    ut.main()


# tests/test_positional_encodings.py

"""
Unit tests for the positional encodings in graphium/features/*
"""

import numpy as np
import networkx as nx
import unittest as ut

# from graphium.features.spectral import compute_laplacian_positional_eigvecs # TODO: add tests
# from graphium.features.rw import compute_rwse # TODO: add tests
from graphium.features.electrostatic import compute_electrostatic_interactions
from graphium.features.commute import compute_commute_distances
from graphium.features.graphormer import compute_graphormer_distances


class test_positional_encodings(ut.TestCase):
    # Test graphs
    adj_dict = {}
    max_dict = {}

    # 6-ring
    adj = np.asarray(
        [
            [0, 1, 0, 0, 0, 1],
            [1, 0, 1, 0, 0, 0],
            [0, 1, 0, 1, 0, 0],
            [0, 0, 1, 0, 1, 0],
            [0, 0, 0, 1, 0, 1],
            [1, 0, 0, 0, 1, 0],
        ]
    )
    adj_dict["6-ring"] = adj
    max_dict["6-ring"] = 3

    # 5-path
    G = nx.path_graph(5)
    adj = nx.to_numpy_array(G)
    adj_dict["5-path"] = adj
    max_dict["5-path"] = 4

    # 4-clique
    adj = 1 - np.eye(4)
    adj_dict["4-clique"] = adj
    max_dict["4-clique"] = 1

    # 4-barbell
    H = nx.barbell_graph(4, 0)
    adj = nx.to_numpy_array(H)
    adj_dict["4-barbell"] = adj
    max_dict["4-barbell"] = 3

    def test_dimensions(self):
        for _, adj in self.adj_dict.items():
            pe, _, _ = compute_electrostatic_interactions(adj, cache={})
            self.assertEqual(pe.shape, adj.shape)

            pe, _, _ = compute_graphormer_distances(adj, adj.shape[0], cache={})
            self.assertEqual(pe.shape, adj.shape)

            pe, _, _ = compute_commute_distances(adj, adj.shape[0], cache={})
            self.assertEqual(pe.shape, adj.shape)

    def test_symmetry(self):
        for _, adj in self.adj_dict.items():
            pe, _, _ = compute_graphormer_distances(adj, adj.shape[0], cache={})
            np.testing.assert_array_almost_equal(pe, pe.T)

            pe, _, _ = compute_commute_distances(adj, adj.shape[0], cache={})
            np.testing.assert_array_almost_equal(pe, pe.T)

    def test_max_dist(self):
        for key, adj in self.adj_dict.items():
            pe, _, _ = compute_graphormer_distances(adj, adj.shape[0], cache={})
            np.testing.assert_array_almost_equal(pe.max(), self.max_dict[key])


if __name__ == "__main__":
    ut.main()


# tests/test_predictor.py

"""
Unit tests for the file graphium/trainer/predictor.py
"""

import torch
from torch.nn import BCELoss, MSELoss
import unittest as ut

from graphium.trainer.predictor_options import EvalOptions


class test_Predictor(ut.TestCase):
    def test_parse_loss_fun(self):
        losses = ["bce", "mse", "mae", "l1", BCELoss(), MSELoss()]
        preds = torch.rand(10, 5)
        target = (torch.rand(10, 5) > 0.5).to(preds.dtype)
        for this_loss in losses:
            loss_fun = EvalOptions.parse_loss_fun(this_loss)
            loss = loss_fun(preds, target)


if __name__ == "__main__":
    ut.main()


# tests/test_pyg_layers.py

"""
Unit tests for the different layers of graphium/nn/pyg_layers/...

The layers are not thoroughly tested due to the difficulty of testing them
"""

import numpy as np
import torch
import unittest as ut
from torch_geometric.data import Data, Batch
from copy import deepcopy
import pytest

from graphium.nn.pyg_layers import (
    GINConvPyg,
    GINEConvPyg,
    MPNNPlusPyg,
    GatedGCNPyg,
    PNAMessagePassingPyg,
    GPSLayerPyg,
    VirtualNodePyg,
    DimeNetPyg,
)

from graphium.nn.pyg_layers.utils import (
    PreprocessPositions,
    GaussianLayer,
)


class test_Pyg_Layers(ut.TestCase):
    in_dim = 21
    out_dim = 11
    in_dim_edges = 13
    out_dim_edges = 17

    edge_idx1 = torch.stack([torch.tensor([0, 1, 2, 3, 2]), torch.tensor([1, 2, 3, 0, 0])])
    edge_idx2 = torch.stack([torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])])
    x1 = torch.randn(edge_idx1.max() + 1, in_dim, dtype=torch.float32)
    e1 = torch.randn(edge_idx1.shape[-1], in_dim_edges, dtype=torch.float32)
    x2 = torch.randn(edge_idx2.max() + 1, in_dim, dtype=torch.float32)
    e2 = torch.randn(edge_idx2.shape[-1], in_dim_edges, dtype=torch.float32)
    # edge_idx1, e1 = add_self_loops(edge_idx1, e1)
    # edge_idx2, e2 = add_self_loops(edge_idx2, e2)
    g1 = Data(feat=x1, edge_index=edge_idx1, edge_feat=e1)
    g2 = Data(feat=x2, edge_index=edge_idx2, edge_feat=e2)
    bg = Batch.from_data_list([g1, g2])

    kwargs = {
        "activation": "relu",
        "dropout": 0.1,
        "normalization": "batch_norm",
        "droppath_rate": 0.1,
        "layer_idx": 1,
        "layer_depth": 10,
    }

    def test_gpslayer(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        kwargs = deepcopy(self.kwargs)
        kwargs.pop("droppath_rate")
        kwargs["droppath_rate_attn"] = 0.2
        kwargs["droppath_rate_ffn"] = 0.3
        kwargs["mpnn_kwargs"] = {"droppath_rate_ffn": 0.4}

        layer = GPSLayerPyg(in_dim=self.in_dim, out_dim=self.out_dim, **kwargs)

        # Check the re-implementation of abstract methods
        self.assertTrue(layer.layer_supports_edges)
        self.assertTrue(layer.layer_inputs_edges)
        self.assertFalse(layer.layer_outputs_edges)
        self.assertEqual(layer.out_dim_factor, 1)

        # Apply layer. Should crash due to different dim for nodes vs edges
        with self.assertRaises(ValueError):
            bg = layer.forward(bg)

        # Create new edge attributes with same dim and check that it works
        bg.edge_feat = torch.zeros((bg.edge_feat.shape[0], self.in_dim), dtype=torch.float32)
        bg = layer.forward(bg)
        self.assertEqual(bg.feat.shape[0], feat_in.shape[0])
        self.assertEqual(bg.feat.shape[1], self.out_dim * layer.out_dim_factor)

    def test_ginlayer(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        layer = GINConvPyg(in_dim=self.in_dim, out_dim=self.out_dim, **self.kwargs)

        # Check the re-implementation of abstract methods
        self.assertFalse(layer.layer_supports_edges)
        self.assertFalse(layer.layer_inputs_edges)
        self.assertFalse(layer.layer_outputs_edges)
        self.assertEqual(layer.out_dim_factor, 1)

        # Apply layer
        bg = layer.forward(bg)
        self.assertEqual(bg.feat.shape[0], feat_in.shape[0])
        self.assertEqual(bg.feat.shape[1], self.out_dim * layer.out_dim_factor)

    def test_ginelayer(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        layer = GINEConvPyg(in_dim=self.in_dim, out_dim=self.out_dim, **self.kwargs)

        # Check the re-implementation of abstract methods
        self.assertTrue(layer.layer_supports_edges)
        self.assertTrue(layer.layer_inputs_edges)
        self.assertFalse(layer.layer_outputs_edges)
        self.assertEqual(layer.out_dim_factor, 1)

        # Apply layer. Should crash due to different dim for nodes vs edges
        with self.assertRaises(ValueError):
            bg = layer.forward(bg)

        # Create new edge attributes with same dim and check that it works
        bg.edge_feat = torch.zeros((bg.edge_feat.shape[0], self.in_dim), dtype=torch.float32)
        bg = layer.forward(bg)
        self.assertEqual(bg.feat.shape[0], feat_in.shape[0])
        self.assertEqual(bg.feat.shape[1], self.out_dim * layer.out_dim_factor)

    def test_mpnnlayer(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        # need in_dim = out_dim for skip connection
        # mpnn layer accept different dimension for node and edge features
        layer = MPNNPlusPyg(
            in_dim=self.in_dim,
            out_dim=self.in_dim,
            use_edges=True,
            in_dim_edges=self.in_dim_edges,
            out_dim_edges=self.in_dim_edges,
            **self.kwargs,
        )

        # Check the re-implementation of abstract methods
        self.assertTrue(layer.layer_supports_edges)
        self.assertTrue(layer.layer_inputs_edges)
        self.assertTrue(layer.layer_outputs_edges)
        self.assertEqual(layer.out_dim_factor, 1)

        # Create new edge attributes with same dim and check that it works
        bg.edge_feat = torch.zeros((bg.edge_feat.shape[0], self.in_dim_edges), dtype=torch.float32)
        bg = layer.forward(bg)
        self.assertEqual(bg.feat.shape[0], feat_in.shape[0])
        self.assertEqual(bg.feat.shape[1], self.in_dim * layer.out_dim_factor)

    def test_gatedgcnlayer(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        e_in = bg.edge_feat
        layer = GatedGCNPyg(
            in_dim=self.in_dim,
            out_dim=self.out_dim,
            in_dim_edges=self.in_dim_edges,
            out_dim_edges=self.out_dim_edges,
            **self.kwargs,
        )

        # Check the re-implementation of abstract methods
        self.assertTrue(layer.layer_supports_edges)
        self.assertTrue(layer.layer_inputs_edges)
        self.assertTrue(layer.layer_outputs_edges)
        self.assertIsInstance(layer.layer_supports_edges, bool)
        self.assertIsInstance(layer.layer_inputs_edges, bool)
        self.assertIsInstance(layer.layer_outputs_edges, bool)
        self.assertEqual(layer.out_dim_factor, 1)

        # Apply layer with edges
        bg2 = layer.forward(bg)
        self.assertEqual(bg2.feat.shape[0], feat_in.shape[0])
        self.assertEqual(bg2.feat.shape[1], self.out_dim * layer.out_dim_factor)

    def test_pnamessagepassinglayer(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        aggregators = ["mean", "max", "min", "std", "sum"]
        scalers = ["identity", "amplification", "attenuation"]

        layer = PNAMessagePassingPyg(
            in_dim=self.in_dim, out_dim=self.out_dim, aggregators=aggregators, scalers=scalers, **self.kwargs
        )

        # Check the re-implementation of abstract methods
        self.assertTrue(layer.layer_supports_edges)
        self.assertIsInstance(layer.layer_supports_edges, bool)
        # self.assertFalse(layer.layer_inputs_edges)
        self.assertFalse(layer.layer_outputs_edges)
        self.assertEqual(layer.out_dim_factor, 1)

        # Apply layer
        bg2 = layer.forward(bg)
        self.assertEqual(bg2.feat.shape[0], feat_in.shape[0])
        self.assertEqual(bg2.feat.shape[1], self.out_dim * layer.out_dim_factor)

        # Now try with edges
        bg = deepcopy(self.bg)
        layer = PNAMessagePassingPyg(
            in_dim=self.in_dim,
            out_dim=self.out_dim,
            aggregators=aggregators,
            scalers=scalers,
            in_dim_edges=self.in_dim_edges,
            **self.kwargs,
        )

        # Check the re-implementation of abstract methods
        self.assertTrue(layer.layer_supports_edges)
        # self.assertTrue(layer.layer_inputs_edges)
        self.assertIsInstance(layer.layer_supports_edges, bool)
        # self.assertIsInstance(layer.layer_inputs_edges, bool)
        self.assertFalse(layer.layer_outputs_edges)
        self.assertEqual(layer.out_dim_factor, 1)

        bg2 = layer.forward(bg)
        self.assertEqual(bg2.feat.shape[0], feat_in.shape[0])
        self.assertEqual(bg2.feat.shape[1], self.out_dim * layer.out_dim_factor)
        self.assertTrue((bg2.edge_feat == self.bg.edge_feat).all)

    @pytest.mark.skip_ipu
    def test_dimenetlayer(self):
        from graphium.nn.encoders.bessel_pos_encoder import BesselSphericalPosEncoder

        bg = deepcopy(self.bg)
        # dummy input pos
        bg.pos = torch.randn((bg.feat.shape[0], 3), dtype=torch.float32)

        # position 3d encoder
        pos_enc = BesselSphericalPosEncoder(
            input_keys=["pos"],
            output_keys=[
                "node_feat",
                "edge_feat",
                "edge_rbf",
                "triplet_sbf",
                "radius_edge_index",
            ],  # The keys to return
            in_dim=3,
            out_dim=self.in_dim,
            out_dim_edges=self.in_dim_edges,
            num_output_layers=2,
            num_layers=2,
            num_spherical=4,
            num_radial=32,
        )

        enc_output = pos_enc(bg, None)  # forward requires: pos, edge_index, batch

        bg.feat = bg.feat + enc_output["node_feat"]  # [num_nodes, in_dim]
        bg.edge_feat = enc_output["edge_feat"]  # [num_edges, out_dim_edges]
        bg.radius_edge_index = enc_output["radius_edge_index"]  # [2, num_edges]
        bg.edge_rbf = enc_output["edge_rbf"]  # [num_edges, num_radial]
        bg.triplet_sbf = enc_output["triplet_sbf"]  # [num_triplets, num_spherical * num_radial]

        kwargs = deepcopy(self.kwargs)
        kwargs["num_bilinear"] = 32
        kwargs["num_spherical"] = 4
        kwargs["num_radial"] = 32

        layer = DimeNetPyg(
            in_dim=self.in_dim,
            out_dim=self.out_dim,
            in_dim_edges=self.in_dim_edges,
            out_dim_edges=self.out_dim_edges,
            **kwargs,
        )

        # Check the re-implementation of abstract methods
        self.assertTrue(layer.layer_supports_edges)
        self.assertTrue(layer.layer_inputs_edges)
        self.assertTrue(layer.layer_outputs_edges)
        self.assertEqual(layer.out_dim_factor, 1)

        # Apply layer with encoded feats as input
        bg2 = layer.forward(bg)
        # output sanity check
        self.assertEqual(bg2.edge_feat.shape, bg.edge_feat.shape)
        self.assertEqual(bg2.edge_feat.shape[1], self.out_dim_edges)
        self.assertEqual(bg2.feat.shape[1], self.out_dim)
        # not change rbf/sbf embedding
        self.assertTrue((bg2.edge_rbf == bg.edge_rbf).all)
        self.assertTrue((bg2.triplet_sbf == bg.triplet_sbf).all)

    def test_preprocess3Dfeaturelayer(self):
        bg = deepcopy(self.bg)
        num_heads = 2
        num_kernel = 2
        in_dim = 3
        bg.positions_3d = torch.zeros(bg.feat.size()[0], in_dim)
        layer = PreprocessPositions(
            num_heads=num_heads,
            embed_dim=self.out_dim,
            num_kernel=num_kernel,
            in_dim=in_dim,
            first_normalization="layer_norm",
        )
        # bias: [batch, num_heads, nodes, nodes]
        # node_feature: [total_nodes, embed_dim]
        bias, node_feature = layer.forward(
            bg, max_num_nodes_per_graph=4, on_ipu=False, positions_3d_key="positions_3d"
        )
        self.assertEqual(bias.size(), torch.Size([2, num_heads, 4, 4]))
        self.assertFalse(np.isnan(bias.detach().numpy()).any())
        self.assertEqual(node_feature.size(), torch.Size([7, self.out_dim]))

    def test_gaussianlayer(self):
        num_kernels = 3
        input = torch.zeros(2, 4, 4)
        layer = GaussianLayer(num_kernels=num_kernels)
        # tensor_with_kernel: [batch, nodes, nodes, num_kernel]
        tensor_with_kernel = layer.forward(input)
        self.assertEqual(tensor_with_kernel.size(), torch.Size([2, 4, 4, 3]))

    def test_pooling_virtual_node(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        e_in = bg.edge_feat
        vn_h = 0

        vn_types = ["sum", "mean"]
        for vn_type, use_edges, expected_v_node in zip(vn_types, [False, True], [(2, 21), (2, 34)]):
            with self.subTest(
                vn_type=vn_type,
                use_edges=use_edges,
                expected_v_node=expected_v_node,
            ):
                vn_h = 0.0
                print(vn_h, self.out_dim, feat_in.size())
                layer = VirtualNodePyg(
                    in_dim=self.in_dim,
                    out_dim=self.in_dim,
                    in_dim_edges=self.in_dim_edges,
                    out_dim_edges=self.in_dim_edges,
                    vn_type=vn_type,
                    use_edges=use_edges,
                    **self.kwargs,
                )

                feat_out, vn_out, e_out = layer.forward(bg, feat_in, vn_h, edge_feat=e_in)
                assert vn_out.shape == expected_v_node
                assert feat_out.shape == (7, 21)
                assert e_out.shape == (9, 13)
                if use_edges is False:
                    # i.e. that the node features have been updated
                    assert torch.equal(feat_out, feat_in) == False
                    # And that the edge features have not
                    assert torch.equal(e_out, e_in) == True
                else:
                    assert torch.equal(feat_out, feat_in) == False
                    assert torch.equal(e_out, e_in) == False


if __name__ == "__main__":
    ut.main()


# tests/test_residual_connections.py

"""
Unit tests for the file residual_connections.py
"""

import numpy as np
import torch
import unittest as ut

from graphium.nn.residual_connections import (
    ResidualConnectionConcat,
    ResidualConnectionDenseNet,
    ResidualConnectionNone,
    ResidualConnectionSimple,
    ResidualConnectionWeighted,
    ResidualConnectionRandom,
)


class test_ResidualConnectionNone(ut.TestCase):
    def test_get_true_out_dims_none(self):
        full_dims = [4, 6, 8, 10, 12]
        in_dims, out_dims = full_dims[:-1], full_dims[1:]
        rc = ResidualConnectionNone(skip_steps=1)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = out_dims[:-1]

        self.assertListEqual(expected_out_dims, true_out_dims)

    def test_forward_none(self):
        rc = ResidualConnectionNone(skip_steps=1)
        num_loops = 10
        shape = (3, 11)
        h_original = [torch.rand(shape) for _ in range(num_loops)]

        h_prev = None
        for ii in range(num_loops):
            feat, h_prev = rc.forward(h_original[ii], h_prev, step_idx=ii)
            np.testing.assert_array_equal(feat.numpy(), h_original[ii].numpy(), err_msg=f"ii={ii}")
            self.assertIsNone(h_prev)


class test_ResidualConnectionSimple(ut.TestCase):
    def test_get_true_out_dims_simple(self):
        full_dims = [4, 6, 8, 10, 12]
        in_dims, out_dims = full_dims[:-1], full_dims[1:]
        rc = ResidualConnectionSimple(skip_steps=1)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = out_dims[:-1]

        self.assertListEqual(expected_out_dims, true_out_dims)

    def test_forward_simple(self):
        for skip_steps in [1, 2, 3]:
            rc = ResidualConnectionSimple(skip_steps=skip_steps)
            num_loops = 10
            shape = (3, 11)
            h_original = [torch.ones(shape) * (ii + 1) for ii in range(num_loops)]

            h_prev = None
            for ii in range(num_loops):
                feat, h_prev = rc.forward(h_original[ii], h_prev, step_idx=ii)

                if ((ii % skip_steps) == 0) and (ii > 0):
                    h_expected = (
                        torch.sum(torch.stack(h_original[0 : ii + 1 : skip_steps], dim=0), dim=0)
                    ).numpy()
                    h_expected_prev = h_expected
                else:
                    h_expected = h_original[ii].numpy()
                if ii == 0:
                    h_expected_prev = h_expected

                np.testing.assert_array_equal(
                    feat.numpy(), h_expected, err_msg=f"Error at: skip_steps={skip_steps}, ii={ii}"
                )
                np.testing.assert_array_equal(
                    h_prev.numpy(), h_expected_prev, err_msg=f"Error at: skip_steps={skip_steps}, ii={ii}"
                )


class test_ResidualConnectionRandom(ut.TestCase):
    def test_get_true_out_dims_random(self):
        full_dims = [4, 6, 8, 10, 12]
        in_dims, out_dims = full_dims[:-1], full_dims[1:]
        rc = ResidualConnectionRandom(skip_steps=1, out_dims=out_dims)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = out_dims[:-1]

        self.assertListEqual(expected_out_dims, true_out_dims)

        rc = ResidualConnectionRandom(skip_steps=1, num_layers=len(out_dims))
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = out_dims[:-1]

        self.assertListEqual(expected_out_dims, true_out_dims)

        with self.assertRaises(ValueError):
            rc = ResidualConnectionRandom(skip_steps=1, out_dims=None, num_layers=None)

    def test_forward_random(self):
        for skip_steps in [1, 2, 3]:
            num_loops = 10

            if skip_steps > 1:
                with self.assertRaises(ValueError):
                    rc = ResidualConnectionRandom(skip_steps=skip_steps, num_layers=num_loops)
                continue

            rc = ResidualConnectionRandom(skip_steps=skip_steps, num_layers=num_loops + 1)
            shape = (3, 11)
            h_original = [torch.ones(shape) * (ii + 1) for ii in range(num_loops)]

            h_prev = None
            # Not really testing the expected values due to randomness, just testing if it runs
            for ii in range(num_loops):
                print(ii)
                feat, h_prev = rc.forward(h_original[ii], h_prev, step_idx=ii)


class test_ResidualConnectionWeighted(ut.TestCase):
    def test_get_true_out_dims_weighted(self):
        full_dims = [4, 6, 8, 10, 12]
        in_dims, out_dims = full_dims[:-1], full_dims[1:]
        rc = ResidualConnectionWeighted(skip_steps=1, out_dims=full_dims[1:])
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = out_dims[:-1]

        self.assertListEqual(expected_out_dims, true_out_dims)

    def test_forward_weighted(self):
        for skip_steps in [1, 2, 3]:
            num_loops = 10
            shape = (3, 11)
            full_dims = [shape[1]] * (num_loops + 1)
            rc = ResidualConnectionWeighted(
                skip_steps=skip_steps,
                out_dims=full_dims[1:],
                activation="none",
                normalization="none",
                bias=False,
            )

            h_original = [torch.ones(shape) * (ii + 1) for ii in range(num_loops)]
            h_forward = []

            h_prev = None
            step_counter = 0
            for ii in range(num_loops):
                h_prev_backup = h_prev
                feat, h_prev = rc.forward(h_original[ii], h_prev, step_idx=ii)

                if ((ii % skip_steps) == 0) and (ii > 0):
                    h_forward.append(rc.residual_list[step_counter].forward(h_prev_backup))
                    h_expected = (h_forward[-1] + h_original[ii]).detach().numpy()
                    h_expected_prev = h_expected
                    step_counter += 1
                else:
                    h_expected = h_original[ii].detach().numpy()
                if ii == 0:
                    h_expected_prev = h_expected

                np.testing.assert_array_equal(
                    feat.detach().numpy(), h_expected, err_msg=f"Error at: skip_steps={skip_steps}, ii={ii}"
                )
                np.testing.assert_array_equal(
                    h_prev.detach().numpy(),
                    h_expected_prev,
                    err_msg=f"Error at: skip_steps={skip_steps}, ii={ii}",
                )


class test_ResidualConnectionConcat(ut.TestCase):
    def test_get_true_out_dims_concat(self):
        full_dims = [4, 6, 8, 10, 12, 14, 16, 18, 20]
        in_dims, out_dims = full_dims[:-1], full_dims[1:]

        # skip_steps=1
        rc = ResidualConnectionConcat(skip_steps=1)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = [6, 14, 18, 22, 26, 30, 34]
        self.assertListEqual(expected_out_dims, true_out_dims)

        # skip_steps=2
        rc = ResidualConnectionConcat(skip_steps=2)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = [6, 8, 16, 12, 24, 16, 32]
        self.assertListEqual(expected_out_dims, true_out_dims)

        # skip_steps=3
        rc = ResidualConnectionConcat(skip_steps=3)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = [6, 8, 10, 18, 14, 16, 30]
        self.assertListEqual(expected_out_dims, true_out_dims)

    def test_forward_concat(self):
        for skip_steps in [1, 2, 3]:
            rc = ResidualConnectionConcat(skip_steps=skip_steps)
            num_loops = 10
            shape = (3, 11)
            h_original = [torch.ones(shape) * (ii + 1) for ii in range(num_loops)]

            h_prev = None
            for ii in range(num_loops):
                feat, h_prev = rc.forward(h_original[ii], h_prev, step_idx=ii)

                if ((ii % skip_steps) == 0) and (ii > 0):
                    h_expected = (
                        torch.cat(h_original[ii - skip_steps : ii + 1 : skip_steps][::-1], dim=-1)
                    ).numpy()
                    h_expected_prev = h_original[ii].numpy()
                else:
                    h_expected = h_original[ii].numpy()
                if ii == 0:
                    h_expected_prev = h_expected

                np.testing.assert_array_equal(
                    feat.numpy(), h_expected, err_msg=f"Error at: skip_steps={skip_steps}, ii={ii}"
                )
                np.testing.assert_array_equal(
                    h_prev.numpy(), h_expected_prev, err_msg=f"Error at: skip_steps={skip_steps}, ii={ii}"
                )


class test_ResidualConnectionDenseNet(ut.TestCase):
    def test_get_true_out_dims_densenet(self):
        full_dims = [4, 6, 8, 10, 12, 14, 16, 18, 20]
        in_dims, out_dims = full_dims[:-1], full_dims[1:]

        # skip_steps=1
        rc = ResidualConnectionDenseNet(skip_steps=1)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = np.cumsum(out_dims).tolist()[:-1]
        self.assertListEqual(expected_out_dims, true_out_dims)

        # skip_steps=2
        rc = ResidualConnectionDenseNet(skip_steps=2)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = [6, 8, 10 + 6, 12, 14 + 10 + 6, 16, 18 + 14 + 10 + 6]
        self.assertListEqual(expected_out_dims, true_out_dims)

        # skip_steps=3
        rc = ResidualConnectionDenseNet(skip_steps=3)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = [6, 8, 10, 12 + 6, 14, 16, 18 + 12 + 6]
        self.assertListEqual(expected_out_dims, true_out_dims)

    def test_forward_densenet(self):
        for skip_steps in [1, 2, 3]:
            rc = ResidualConnectionDenseNet(skip_steps=skip_steps)
            num_loops = 10
            shape = (3, 11)
            h_original = [torch.ones(shape) * (ii + 1) for ii in range(num_loops)]

            h_prev = None
            for ii in range(num_loops):
                feat, h_prev = rc.forward(h_original[ii], h_prev, step_idx=ii)

                if ((ii % skip_steps) == 0) and (ii > 0):
                    h_expected = (torch.cat(h_original[0 : ii + 1 : skip_steps][::-1], dim=-1)).numpy()
                    h_expected_prev = h_expected
                else:
                    h_expected = h_original[ii].numpy()
                if ii == 0:
                    h_expected_prev = h_expected

                np.testing.assert_array_equal(
                    feat.numpy(), h_expected, err_msg=f"Error at: skip_steps={skip_steps}, ii={ii}"
                )
                np.testing.assert_array_equal(
                    h_prev.numpy(), h_expected_prev, err_msg=f"Error at: skip_steps={skip_steps}, ii={ii}"
                )


if __name__ == "__main__":
    ut.main()


# tests/test_utils.py

"""
Unit tests for the metrics and wrappers of graphium/utils/...
"""

from graphium.utils.tensor import nan_mad, nan_mean, nan_std, nan_var, nan_median
from graphium.utils.safe_run import SafeRun
import torch
import numpy as np
import scipy as sp
import unittest as ut
import gzip
from graphium.utils.read_file import file_opener


class test_nan_statistics(ut.TestCase):
    torch.manual_seed(42)

    dims = [
        None,
        (0),
        (1),
        (2),
        (-1),
        (-2),
        (-3),
        (0, 1),
        (0, 2),
    ]
    # Create tensor
    sz = (10, 6, 7)
    tensor = torch.randn(sz, dtype=torch.float32) ** 2 + 3
    is_nan = torch.rand(sz) > 0.4
    tensor[is_nan] = float("nan")

    def test_nan_mean(self):
        for keepdim in [False, True]:
            for dim in self.dims:
                err_msg = f"Error for :\n dim = {dim}\n keepdim = {keepdim}"

                tensor = self.tensor.clone()

                # Prepare the arguments for numpy vs torch
                if dim is not None:
                    torch_kwargs = {"dim": dim, "keepdim": keepdim}
                    numpy_kwargs = {"axis": dim, "keepdims": keepdim}
                else:
                    torch_kwargs = {}
                    numpy_kwargs = {}

                # Compare the nan-mean
                torch_mean = nan_mean(tensor, **torch_kwargs)
                numpy_mean = np.nanmean(tensor.numpy(), **numpy_kwargs)
                np.testing.assert_almost_equal(torch_mean.numpy(), numpy_mean, decimal=6, err_msg=err_msg)

    def test_nan_std_var(self):
        for unbiased in [True, False]:
            for keepdim in [False, True]:
                for dim in self.dims:
                    err_msg = f"Error for :\n\tdim = {dim}\n\tkeepdim = {keepdim}\n\tunbiased = {unbiased}"

                    tensor = self.tensor.clone()

                    # Prepare the arguments for numpy vs torch
                    if dim is not None:
                        torch_kwargs = {"dim": dim, "keepdim": keepdim, "unbiased": unbiased}
                        numpy_kwargs = {"axis": dim, "keepdims": keepdim, "ddof": float(unbiased)}
                    else:
                        torch_kwargs = {"unbiased": unbiased}
                        numpy_kwargs = {"ddof": float(unbiased)}

                    # Compare the std
                    torch_std = nan_std(tensor, **torch_kwargs)
                    numpy_std = np.nanstd(tensor.numpy(), **numpy_kwargs)
                    np.testing.assert_almost_equal(torch_std.numpy(), numpy_std, decimal=6, err_msg=err_msg)

                    # Compare the variance
                    torch_var = nan_var(tensor, **torch_kwargs)
                    numpy_var = np.nanvar(tensor.numpy(), **numpy_kwargs)
                    np.testing.assert_almost_equal(torch_var.numpy(), numpy_var, decimal=6, err_msg=err_msg)

    def test_nan_median(self):
        for keepdim in [False, True]:
            # Cannot test
            for dim in self.dims:  # in [d for d in self.dims if not isinstance(d, Tuple)]:
                err_msg = f"Error for :\n dim = {dim}\n keepdim = {keepdim}"

                tensor = torch.randn(
                    (7, 9, 11, 13)
                )  # Need odd number of values to properly compare torch to numpy

                # Prepare the arguments for numpy vs torch
                if dim is not None:
                    torch_kwargs = {"dim": dim, "keepdim": keepdim}
                    numpy_kwargs = {"axis": dim, "keepdims": keepdim}
                else:
                    torch_kwargs = {}
                    numpy_kwargs = {}

                # Compare the nan-median
                torch_med = nan_median(tensor, **torch_kwargs)
                numpy_med = np.nanmedian(tensor.numpy(), **numpy_kwargs)
                torch_sum = torch.nansum(tensor, **torch_kwargs)
                np.testing.assert_almost_equal(torch_med.numpy(), numpy_med, decimal=4, err_msg=err_msg)
                self.assertListEqual(list(torch_med.shape), list(torch_sum.shape))

    def test_nan_mad(self):
        for normal in [False, True]:
            # Cannot test
            for dim in self.dims:  # in [d for d in self.dims if not isinstance(d, Tuple)]:
                err_msg = f"Error for :\n dim = {dim}\n normal = {normal}"

                tensor = torch.randn(
                    (7, 9, 11, 13)
                )  # Need odd number of values to properly compare torch to numpy

                # Prepare the arguments for numpy vs torch
                if dim is not None:
                    torch_kwargs = {"dim": dim, "keepdim": False, "normal": normal}
                    numpy_kwargs = {"axis": dim, "nan_policy": "omit", "scale": 1 / 1.4826 if normal else 1.0}
                else:
                    torch_kwargs = {"normal": normal}
                    numpy_kwargs = {"axis": dim, "nan_policy": "omit", "scale": 1 / 1.4826 if normal else 1.0}

                # Compare the nan-median
                torch_mad = nan_mad(tensor, **torch_kwargs)
                numpy_mad = sp.stats.median_abs_deviation(tensor.numpy(), **numpy_kwargs)
                np.testing.assert_almost_equal(torch_mad.numpy(), numpy_mad, decimal=4, err_msg=err_msg)


def test_file_opener(tmp_path):
    # Create a temporary file
    txt_file = tmp_path / "test.txt"
    txt_file.write_text("Hello, World!")

    # Test opening file in read mode
    with file_opener(txt_file, "r") as f:
        assert f.read() == "Hello, World!"

    # Test opening file in write mode
    with file_opener(txt_file, "w") as f:
        f.write("New text")

    with file_opener(txt_file, "r") as f:
        assert f.read() == "New text"

    # Create a temporary gzip file
    gzip_file = tmp_path / "test.txt.gz"
    with gzip.open(gzip_file, "wt") as f:
        f.write("Hello, Gzip!")

    # Test opening gzip file in read mode
    with file_opener(gzip_file, "r") as f:
        assert f.read() == "Hello, Gzip!"


class test_SafeRun(ut.TestCase):
    def test_safe_run(self):
        # Error is caught
        with SafeRun(name="bob", raise_error=False, verbose=0):
            raise ValueError("This is an error")

        # Error is caught
        with SafeRun(name="bob", raise_error=False, verbose=0):
            2 + None

        # Error is not caught
        with self.assertRaises(ValueError):
            with SafeRun(name="bob", raise_error=True, verbose=0):
                raise ValueError("This is an error")

        # Error is not caught
        with self.assertRaises(TypeError):
            with SafeRun(name="bob", raise_error=True, verbose=0):
                2 + None

        # No error. Runs correctly
        with SafeRun(name="bob", raise_error=True, verbose=0):
            print("This is not an error")

        # No error. Runs correctly
        with SafeRun(name="bob", raise_error=False, verbose=0):
            print("This is not an error")


if __name__ == "__main__":
    ut.main()


# graphium/cli/__init__.py

from .main import main_cli
from .data import data_cli


# graphium/cli/data.py

import click

from loguru import logger

import graphium

from .main import main_cli


@main_cli.group(name="data", help="Graphium datasets.")
def data_cli():
    pass


@data_cli.command(name="download", help="Download a Graphium dataset.")
@click.option(
    "-n",
    "--name",
    type=str,
    required=True,
    help="Name of the graphium dataset to download.",
)
@click.option(
    "-o",
    "--output",
    type=str,
    required=True,
    help="Where to download the Graphium dataset.",
)
@click.option(
    "--progress",
    type=bool,
    is_flag=True,
    default=False,
    required=False,
    help="Whether to extract the dataset if it's a zip file.",
)
def download(name, output, progress):
    args = {}
    args["name"] = name
    args["output_path"] = output
    args["extract_zip"] = True
    args["progress"] = progress

    logger.info(f"Download dataset '{name}' into {output}.")

    fpath = graphium.data.utils.download_graphium_dataset(**args)

    logger.info(f"Dataset available at {fpath}.")


@data_cli.command(name="list", help="List available Graphium dataset.")
def list():
    logger.info("Graphium datasets:")
    logger.info(graphium.data.utils.list_graphium_datasets())


# graphium/cli/main.py

import click


@click.group()
@click.version_option()
def main_cli():
    pass


if __name__ == "__main__":
    main_cli()


# graphium/config/__init__.py

from ._load import load_config

from ._loader import load_architecture
from ._loader import load_datamodule
from ._loader import load_metrics
from ._loader import load_predictor
from ._loader import load_trainer


# graphium/config/_load.py

import importlib.resources

import omegaconf


def load_config(name: str):
    """Load a default config file by its name.

    Args:
        name: name of the config to load.
    """

    with importlib.resources.open_text("graphium.config", f"{name}.yaml") as f:
        config = omegaconf.OmegaConf.load(f)

    return config


# graphium/config/_loader.py

from typing import Dict, Mapping, Tuple, Type, Union, Any, Optional, Callable

# Misc
import os
import omegaconf
from copy import deepcopy
from loguru import logger
import yaml
import joblib
import pathlib
import warnings

# Torch
import torch
import mup

# Lightning
from lightning import Trainer
from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint
from lightning.pytorch.loggers import WandbLogger, Logger

# Graphium
from graphium.utils.mup import set_base_shapes
from graphium.ipu.ipu_dataloader import IPUDataloaderOptions
from graphium.trainer.metrics import MetricWrapper
from graphium.nn.architectures import FullGraphMultiTaskNetwork
from graphium.nn.utils import MupMixin
from graphium.trainer.predictor import PredictorModule
from graphium.utils.spaces import DATAMODULE_DICT
from graphium.ipu.ipu_utils import import_poptorch, load_ipu_options
from graphium.data.datamodule import MultitaskFromSmilesDataModule, BaseDataModule
from graphium.utils.command_line_utils import update_config, get_anchors_and_aliases


def get_accelerator(
    config_acc: Union[omegaconf.DictConfig, Dict[str, Any]],
) -> str:
    """
    Get the accelerator from the config file, and ensure that they are
    consistant.
    """

    # Get the accelerator type
    accelerator_type = config_acc["type"]

    # Get the GPU info
    if (accelerator_type == "gpu") and (not torch.cuda.is_available()):
        raise ValueError(f"GPUs selected, but GPUs are not available on this device")

    # Get the IPU info
    if accelerator_type == "ipu":
        poptorch = import_poptorch()
        if poptorch is None:
            raise ValueError("IPUs selected, but PopTorch is not available")
        if not poptorch.ipuHardwareIsAvailable():
            raise ValueError(
                "IPUs selected, but no IPU is available/visible on this device. "
                "If you do have IPUs, please check that the IPUOF_VIPU_API_PARTITION_ID and "
                "IPUOF_VIPU_API_HOST environment variables are set."
            )

    # Fall on cpu at the end
    if accelerator_type is None:
        accelerator_type = "cpu"
    return accelerator_type


def _get_ipu_opts(config: Union[omegaconf.DictConfig, Dict[str, Any]]) -> Tuple[str, str]:
    r"""
    Get the paths of the IPU-specific config files from the main YAML config
    """

    accelerator_options = config["accelerator"]
    accelerator_type = accelerator_options["type"]

    if accelerator_type != "ipu":
        return None, None

    ipu_opts = accelerator_options["ipu_config"]
    ipu_inference_opts = accelerator_options.get("ipu_inference_config", None)

    return ipu_opts, ipu_inference_opts


def load_datamodule(
    config: Union[omegaconf.DictConfig, Dict[str, Any]], accelerator_type: str
) -> BaseDataModule:
    """
    Load the datamodule from the specified configurations at the key
    `datamodule: args`.
    If the accelerator is IPU, load the IPU options as well.

    Parameters:
        config: The config file, with key `datamodule: args`
        accelerator_type: The accelerator type, e.g. "cpu", "gpu", "ipu"
    Returns:
        datamodule: The datamodule used to process and load the data
    """

    cfg_data = config["datamodule"]["args"]

    # Instanciate the datamodule
    module_class = DATAMODULE_DICT[config["datamodule"]["module_type"]]

    if accelerator_type != "ipu":
        datamodule = module_class(
            **config["datamodule"]["args"],
        )
        return datamodule

    # IPU specific adjustments
    else:
        ipu_opts, ipu_inference_opts = _get_ipu_opts(config)

        # Default empty values for the IPU configurations
        ipu_training_opts = None

        ipu_dataloader_training_opts = cfg_data.pop("ipu_dataloader_training_opts", {})
        ipu_dataloader_inference_opts = cfg_data.pop("ipu_dataloader_inference_opts", {})
        ipu_training_opts, ipu_inference_opts = load_ipu_options(
            ipu_opts=ipu_opts,
            seed=config["constants"]["seed"],
            model_name=config["constants"]["name"],
            gradient_accumulation=config["trainer"]["trainer"].get("accumulate_grad_batches", None),
            ipu_inference_opts=ipu_inference_opts,
            precision=config["trainer"]["trainer"].get("precision"),
        )
        # Define the Dataloader options for the IPU on the training sets
        bz_train = cfg_data["batch_size_training"]
        ipu_dataloader_training_opts = IPUDataloaderOptions(
            batch_size=bz_train, **ipu_dataloader_training_opts
        )
        ipu_dataloader_training_opts.set_kwargs()

        # Define the Dataloader options for the IPU on the inference sets
        bz_test = cfg_data["batch_size_inference"]
        ipu_dataloader_inference_opts = IPUDataloaderOptions(
            batch_size=bz_test, **ipu_dataloader_inference_opts
        )
        ipu_dataloader_inference_opts.set_kwargs()

        datamodule = module_class(
            ipu_training_opts=ipu_training_opts,
            ipu_inference_opts=ipu_inference_opts,
            ipu_dataloader_training_opts=ipu_dataloader_training_opts,
            ipu_dataloader_inference_opts=ipu_dataloader_inference_opts,
            **config["datamodule"]["args"],
        )

        return datamodule


def load_metrics(config: Union[omegaconf.DictConfig, Dict[str, Any]]) -> Dict[str, MetricWrapper]:
    """
    Loading the metrics to be tracked.
    Parameters:
        config: The config file, with key `metrics`
    Returns:
        metrics: A dictionary of all the metrics
    """

    task_metrics = {}
    cfg_metrics = config.get("metrics", None)
    if cfg_metrics is None:
        return task_metrics
    cfg_metrics = {key: deepcopy(value) for key, value in config["metrics"].items()}
    # Wrap every metric in the class `MetricWrapper` to standardize them
    for task in cfg_metrics:
        task_metrics[task] = {}
        if cfg_metrics[task] is None:
            cfg_metrics[task] = []
        for this_metric in cfg_metrics[task]:
            name = this_metric.pop("name")
            task_metrics[task][name] = MetricWrapper(**this_metric)
    return task_metrics


def load_architecture(
    config: Union[omegaconf.DictConfig, Dict[str, Any]],
    in_dims: Dict[str, int],
) -> Union[FullGraphMultiTaskNetwork, torch.nn.Module]:
    """
    Loading the architecture used for training.
    Parameters:
        config: The config file, with key `architecture`
        in_dims: Dictionary of the input dimensions for various
    Returns:
        architecture: The datamodule used to process and load the data
    """

    if isinstance(config, dict):
        config = omegaconf.OmegaConf.create(config)
    cfg_arch = config["architecture"]

    kwargs = {}

    # Select the architecture
    model_type = cfg_arch["model_type"].lower()
    if model_type == "fullgraphmultitasknetwork":
        model_class = FullGraphMultiTaskNetwork
    else:
        raise ValueError(f"Unsupported model_type=`{model_type}`")

    # Prepare the various kwargs
    pe_encoders_kwargs = (
        dict(cfg_arch["pe_encoders"]) if cfg_arch.get("pe_encoders", None) is not None else None
    )

    pre_nn_kwargs = dict(cfg_arch["pre_nn"]) if cfg_arch["pre_nn"] is not None else None
    pre_nn_edges_kwargs = dict(cfg_arch["pre_nn_edges"]) if cfg_arch["pre_nn_edges"] is not None else None
    gnn_kwargs = dict(cfg_arch["gnn"])
    graph_output_nn_kwargs = (
        dict(cfg_arch["graph_output_nn"]) if cfg_arch["graph_output_nn"] is not None else None
    )
    task_heads_kwargs = (
        cfg_arch["task_heads"] if cfg_arch["task_heads"] is not None else None
    )  # This is of type ListConfig containing TaskHeadParams

    # Initialize the input dimension for the positional encoders
    if pe_encoders_kwargs is not None:
        pe_encoders_kwargs = dict(pe_encoders_kwargs)
        for encoder in pe_encoders_kwargs["encoders"]:
            pe_encoders_kwargs["encoders"][encoder] = dict(pe_encoders_kwargs["encoders"][encoder])
        pe_encoders_kwargs.setdefault(
            "in_dims", in_dims
        )  # set the input dimensions of all pe with info from the data-module
    pe_out_dim = 0 if pe_encoders_kwargs is None else pe_encoders_kwargs.get("out_dim", None)
    edge_pe_out_dim = 0 if pe_encoders_kwargs is None else pe_encoders_kwargs.get("edge_out_dim", None)

    # Set the default `node` input dimension for the pre-processing neural net and graph neural net
    in_dim = in_dims["feat"]
    if pe_out_dim is not None:
        in_dim += pe_out_dim
    if pre_nn_kwargs is not None:
        pre_nn_kwargs = dict(pre_nn_kwargs)
        pre_nn_kwargs.setdefault("in_dim", in_dim)
    else:
        gnn_kwargs.setdefault("in_dim", in_dim)

    # Set the default `edge` input dimension for the pre-processing neural net and graph neural net
    edge_in_dim = in_dims["edge_feat"]
    if edge_pe_out_dim is not None:
        edge_in_dim += edge_pe_out_dim
    if pre_nn_edges_kwargs is not None:
        pre_nn_edges_kwargs = dict(pre_nn_edges_kwargs)
        pre_nn_edges_kwargs.setdefault("in_dim", edge_in_dim)
    else:
        gnn_kwargs.setdefault("in_dim", edge_in_dim)

    # Set the parameters for the full network
    task_heads_kwargs = omegaconf.OmegaConf.to_object(task_heads_kwargs)

    # Set all the input arguments for the model
    model_kwargs = dict(
        gnn_kwargs=gnn_kwargs,
        pre_nn_kwargs=pre_nn_kwargs,
        pre_nn_edges_kwargs=pre_nn_edges_kwargs,
        pe_encoders_kwargs=pe_encoders_kwargs,
        graph_output_nn_kwargs=graph_output_nn_kwargs,
        task_heads_kwargs=task_heads_kwargs,
    )

    return model_class, model_kwargs


def load_predictor(
    config: Union[omegaconf.DictConfig, Dict[str, Any]],
    model_class: Type[torch.nn.Module],
    model_kwargs: Dict[str, Any],
    metrics: Dict[str, MetricWrapper],
    accelerator_type: str,
    task_norms: Optional[Dict[Callable, Any]] = None,
) -> PredictorModule:
    """
    Defining the predictor module, which handles the training logic from `lightning.LighningModule`
    Parameters:
        model_class: The torch Module containing the main forward function
        accelerator_type: The accelerator type, e.g. "cpu", "gpu", "ipu"
    Returns:
        predictor: The predictor module
    """

    if accelerator_type == "ipu":
        from graphium.ipu.ipu_wrapper import PredictorModuleIPU

        predictor_class = PredictorModuleIPU
    else:
        predictor_class = PredictorModule

    cfg_pred = dict(deepcopy(config["predictor"]))
    predictor = predictor_class(
        model_class=model_class,
        model_kwargs=model_kwargs,
        metrics=metrics,
        task_norms=task_norms,
        **cfg_pred,
    )

    mup_scale_factor = config["architecture"].pop("mup_scale_factor", None)

    if mup_scale_factor is not None and mup_scale_factor != 1:
        unscaled_model = predictor.model
        scaled_model_kwargs = unscaled_model.scale_kwargs(scale_factor=mup_scale_factor)
        del predictor
        predictor = predictor_class(
            model_class=model_class,
            model_kwargs=scaled_model_kwargs,
            metrics=metrics,
            **cfg_pred,
        )

    # mup base shapes
    mup_base_path = config["architecture"].pop("mup_base_path", None)
    predictor = load_mup(mup_base_path, predictor)

    return predictor


def load_mup(mup_base_path: str, predictor: PredictorModule) -> PredictorModule:
    """
    Load the base shapes for the mup, based either on a `.ckpt` or `.yaml` file.
    If `.yaml`, it should be generated by `mup.save_base_shapes`
    """
    model = predictor.model

    if not isinstance(model, MupMixin):
        raise TypeError("load_mup can only be applied to models that use the MupMixin")

    if mup_base_path is None:
        base = model.__class__(**model.make_mup_base_kwargs(divide_factor=2))
    elif mup_base_path.endswith(".ckpt"):
        base = predictor.__class__.load_from_checkpoint(mup_base_path, map_location="cpu")
    elif mup_base_path.endswith(".yaml"):
        base = mup_base_path
    else:
        raise ValueError(f"Unrecognized file type {mup_base_path}")
    predictor.model = set_base_shapes(predictor.model, base, rescale_params=False)
    return predictor


def load_trainer(
    config: Union[omegaconf.DictConfig, Dict[str, Any]],
    run_name: str,
    accelerator_type: str,
    date_time_suffix: str = "",
) -> Trainer:
    """
    Defining the pytorch-lightning Trainer module.
    Parameters:
        config: The config file, with key `trainer`
        run_name: The name of the current run. To be used for logging.
        accelerator_type: The accelerator type, e.g. "cpu", "gpu", "ipu"
        date_time_suffix: The date and time of the current run. To be used for logging.
    Returns:
        trainer: the trainer module
    """
    cfg_trainer = deepcopy(config["trainer"])

    # Define the IPU plugin if required
    strategy = "auto"
    if accelerator_type == "ipu":
        ipu_opts, ipu_inference_opts = _get_ipu_opts(config)

        training_opts, inference_opts = load_ipu_options(
            ipu_opts=ipu_opts,
            ipu_inference_opts=ipu_inference_opts,
            seed=config["constants"]["seed"],
            model_name=config["constants"]["name"],
            gradient_accumulation=config["trainer"]["trainer"].get("accumulate_grad_batches", None),
            precision=config["trainer"]["trainer"].get("precision"),
        )

        from lightning_graphcore import IPUStrategy

        strategy = IPUStrategy(training_opts=training_opts, inference_opts=inference_opts)

    # Get devices
    devices = cfg_trainer["trainer"].pop("devices", 1)
    if accelerator_type == "ipu":
        devices = 1  # number of IPUs used is defined in the ipu options files

    # Remove the gradient accumulation from IPUs, since it's handled by the device
    if accelerator_type == "ipu":
        cfg_trainer["trainer"].pop("accumulate_grad_batches", None)

    # Define the early stopping parameters
    trainer_kwargs = {}
    callbacks = []
    if "early_stopping" in cfg_trainer.keys():
        callbacks.append(EarlyStopping(**cfg_trainer["early_stopping"]))

    # Define the early model checkpoing parameters
    if "model_checkpoint" in cfg_trainer.keys():
        cfg_trainer["model_checkpoint"]["dirpath"] += str(cfg_trainer["seed"]) + "/"
        callbacks.append(ModelCheckpoint(**cfg_trainer["model_checkpoint"]))

    # Define the logger parameters
    logger = cfg_trainer.pop("logger", None)
    if logger is not None:
        name = logger.pop("name", run_name)
        if len(date_time_suffix) > 0:
            name += f"_{date_time_suffix}"
        trainer_kwargs["logger"] = WandbLogger(name=name, **logger)

    trainer_kwargs["callbacks"] = callbacks
    trainer = Trainer(
        detect_anomaly=True,
        strategy=strategy,
        accelerator=accelerator_type,
        devices=devices,
        **cfg_trainer["trainer"],
        **trainer_kwargs,
    )
    return trainer


def save_params_to_wandb(
    logger: Logger,
    config: Union[omegaconf.DictConfig, Dict[str, Any]],
    predictor: PredictorModule,
    datamodule: MultitaskFromSmilesDataModule,
):
    """
    Save a few stuff to weights-and-biases WandB
    Parameters:
        logger: The object used to log the training. Usually WandbLogger
        config: The config file, with key `trainer`
        predictor: The predictor used to handle the train/val/test steps logic
        datamodule: The datamodule used to load the data into training
    """

    # Get the wandb runner and directory
    wandb_run = logger.experiment
    if wandb_run is None:
        wandb_run = ""
    wandb_dir = wandb_run.dir

    # Save the mup base model to WandB as a yaml file
    mup.save_base_shapes(predictor.model, os.path.join(wandb_dir, "mup_base_params.yaml"))

    # Save the full configs as a YAML file
    with open(os.path.join(wandb_dir, "full_configs.yaml"), "w") as file:
        yaml.dump(config, file)

    # Save the featurizer into wandb
    featurizer_path = os.path.join(wandb_dir, "featurizer.pickle")
    joblib.dump(datamodule.smiles_transformer, featurizer_path)

    # Save the featurizer and configs into wandb
    if wandb_run is not None:
        wandb_run.save("*.yaml")
        wandb_run.save("*.pickle")


def load_accelerator(config: Union[omegaconf.DictConfig, Dict[str, Any]]) -> Tuple[Dict[str, Any], str]:
    config = deepcopy(config)
    config_acc = config.get("accelerator", {})

    # Merge the accelerator config with the main config
    config_override = config_acc.get("config_override", {})
    merge_dicts(config, config_override)
    accelerator_type = get_accelerator(config_acc)

    if accelerator_type == "gpu":
        precision = config_acc.get("float32_matmul_precision", None)
        if precision is not None:
            torch.set_float32_matmul_precision(precision)

    return config, accelerator_type


def load_config_override(
    config: Union[omegaconf.DictConfig, Dict[str, Any]], main_dir: Optional[Union[str, os.PathLike]] = None
) -> Dict[str, Any]:
    config = deepcopy(config)
    config_override_path = config["constants"].get("config_override", None)
    if config_override_path is not None:
        if main_dir is not None:
            config_override_path = os.path.join(main_dir, config_override_path)
        with open(config_override_path, "r") as f:
            cfg_override = yaml.safe_load(f)
        config = merge_dicts(cfg_override, config, on_exist="overwrite")
    return config


def load_yaml_config(
    config_path: Union[str, os.PathLike],
    main_dir: Optional[Union[str, os.PathLike]] = None,
    unknown_args=None,
) -> Dict[str, Any]:
    """
    Load a YAML config file and return it as a dictionary.
    Also returns the anchors `&` and aliases `*` of the YAML file.
    Then, update the config with the unknown arguments.
    Finally, update the config with the config override file specified in `constants.config_override`.

    Parameters:
        config_path: The path to the YAML config file
        main_dir: The main directory of the project. If specified, the config override file will be loaded from this directory
        unknown_args: The unknown arguments to update the config with, taken from `argparse.parse_known_args`

    Returns:
        config: The config dictionary

    """
    if main_dir is not None:
        config_path = os.path.join(main_dir, config_path)

    with open(config_path, "r") as f:
        config = yaml.safe_load(f)
        refs = get_anchors_and_aliases(config_path)
        if unknown_args is not None:
            config = update_config(config, unknown_args, refs)
    config = load_config_override(config, main_dir)  # This goes here to avoid overriding the hparam search

    return config


def merge_dicts(
    dict_a: Dict[str, Any], dict_b: Dict[str, Any], previous_dict_path: str = "", on_exist: str = "raise"
) -> None:
    """
    Recursively merges dict_b into dict_a. If a key is missing from dict_a,
    it is added from dict_b. If a key exists in both, an error is raised.
    `dict_a` is modified in-place.

    Parameters:
        dict_a: The dictionary to merge into. Modified in-place.
        dict_b: The dictionary to merge from.
        previous_dict_path: The key path of the parent dictionary,
        used to track the recursive calls.
        on_exist: What to do if a key already exists in dict_a. Options are "raise", "overwrite", "ignore".

    Raises:
        ValueError: If a key path already exists in dict_a.

    """
    assert on_exist in [
        "raise",
        "overwrite",
        "ignore",
    ], f"on_exist must be one of ['raise', 'overwrite', 'ignore'], got {on_exist}"

    for key, value_b in dict_b.items():
        if key not in dict_a:
            dict_a[key] = value_b
        else:
            value_a = dict_a[key]
            if previous_dict_path == "":
                previous_dict_path = key
            else:
                previous_dict_path = f"{previous_dict_path}/{key}"
            if isinstance(value_a, dict) and isinstance(value_b, dict):
                merge_dicts(value_a, value_b, previous_dict_path=previous_dict_path, on_exist=on_exist)
            else:
                if value_a != value_b:
                    if on_exist == "raise":
                        raise ValueError(f"Dict path already exists: {previous_dict_path}")
                    elif on_exist == "overwrite":
                        dict_a[key] = value_b
                    elif on_exist == "ignore":
                        pass
    return dict_a


# graphium/config/config_convert.py

import omegaconf


def recursive_config_reformating(configs):
    r"""
    For a given configuration file, convert all `DictConfig` to `dict`,
    all `ListConfig` to `list`, and all `byte` to `str`.

    This helps avoid errors when dumping a yaml file.
    """

    if isinstance(configs, omegaconf.DictConfig):
        configs = dict(configs)
    elif isinstance(configs, omegaconf.ListConfig):
        configs = list(configs)

    if isinstance(configs, dict):
        for k, v in configs.items():
            if isinstance(v, bytes):
                configs[k] = str(v)
            else:
                configs[k] = recursive_config_reformating(v)
    elif isinstance(configs, list):
        for k, v in enumerate(configs):
            if isinstance(v, bytes):
                configs[k] = str(v)
            else:
                configs[k] = recursive_config_reformating(v)

    return configs


# graphium/data/__init__.py

from .utils import load_micro_zinc
from .utils import load_tiny_zinc

from .collate import graphium_collate_fn

from .datamodule import GraphOGBDataModule
from .datamodule import MultitaskFromSmilesDataModule
from .datamodule import ADMETBenchmarkDataModule
from .datamodule import FakeDataModule

from .dataset import SingleTaskDataset
from .dataset import MultitaskDataset
from .dataset import FakeDataset


# graphium/data/collate.py

from collections.abc import Mapping, Sequence

# from pprint import pprint
import torch
from numpy import ndarray
from scipy.sparse import spmatrix
from torch.utils.data.dataloader import default_collate
from typing import Union, List, Optional, Dict, Type, Any, Iterable
from torch_geometric.data import Data, Batch

from graphium.features import GraphDict, to_dense_array
from graphium.utils.packing import fast_packing, get_pack_sizes, node_to_pack_indices_mask
from loguru import logger
from graphium.data.utils import get_keys


def graphium_collate_fn(
    elements: Union[List[Any], Dict[str, List[Any]]],
    labels_size_dict: Optional[Dict[str, Any]] = None,
    labels_dtype_dict: Optional[Dict[str, Any]] = None,
    mask_nan: Union[str, float, Type[None]] = "raise",
    do_not_collate_keys: List[str] = [],
    batch_size_per_pack: Optional[int] = None,
) -> Union[Any, Dict[str, Any]]:
    """This collate function is identical to the default
    pytorch collate function but add support for `pyg.data.Data` to batch graphs.

    Beside pyg graph collate, other objects are processed the same way
    as the original torch collate function. See https://pytorch.org/docs/stable/data.html#dataloader-collate-fn
    for more details.

    Note:
        If graphium needs to manipulate other tricky-to-batch objects. Support
        for them should be added to this single collate function.

    Parameters:

        elements:
            The elements to batch. See `torch.utils.data.dataloader.default_collate`.

        labels_size_dict:
            (Note): This is an attribute of the `MultitaskDataset`.
            A dictionary of the form Dict[tasks, sizes] which has task names as keys
            and the size of the label tensor as value. The size of the tensor corresponds to how many
            labels/values there are to predict for that task.

        labels_dtype_dict:
            (Note): This is an attribute of the `MultitaskDataset`.
            A dictionary of the form Dict[tasks, dtypes] which has task names as keys
            and the dtype of the label tensor as value. This is necessary to ensure the missing labels are added with NaNs of the right dtype

        mask_nan:
            Deal with the NaN/Inf when calling the function `make_pyg_graph`.
            Some values become `Inf` when changing data type. This allows to deal
            with that.

            - "raise": Raise an error when there is a nan or inf in the featurization
            - "warn": Raise a warning when there is a nan or inf in the featurization
            - "None": DEFAULT. Don't do anything
            - "Floating value": Replace nans or inf by the specified value

        do_not_batch_keys:
            Keys to ignore for the collate

        batch_size_per_pack: The number of graphs to pack together.
            This is useful for using packing with the Transformer.
            If None, no packing is done.
            Otherwise, indices are generated to map the nodes to the pack they belong to under the key `"pack_from_node_idx"`,
            with an additional mask to indicate which nodes are from the same graph under the key `"pack_attn_mask"`.

    Returns:
        The batched elements. See `torch.utils.data.dataloader.default_collate`.
    """

    elem = elements[0]
    if isinstance(elem, Mapping):
        batch = {}
        for key in elem:
            # Multitask setting: We have to pad the missing labels
            if key == "labels":
                labels = [d[key] for d in elements]
                batch[key] = collate_labels(labels, labels_size_dict, labels_dtype_dict)

            # If the features are a dictionary containing GraphDict elements,
            # Convert to pyg graphs and use the pyg batching.
            elif isinstance(elem[key], GraphDict):
                pyg_graphs = [d[key].make_pyg_graph(mask_nan=mask_nan) for d in elements]
                batch[key] = collage_pyg_graph(pyg_graphs)

            # If a PyG Graph is provided, use the PyG batching
            elif isinstance(elem[key], Data):
                pyg_graphs = [d[key] for d in elements]
                batch[key] = collage_pyg_graph(pyg_graphs, batch_size_per_pack=batch_size_per_pack)

            # Ignore the collate for specific keys
            elif key in do_not_collate_keys:
                batch[key] = [d[key] for d in elements]
            # Otherwise, use the default torch batching
            else:
                batch[key] = default_collate([d[key] for d in elements])
        return batch
    elif isinstance(elements, Sequence) and isinstance(elem, Sequence):
        temp_elements = [{ii: sub_elem for ii, sub_elem in enumerate(elem)} for elem in elements]
        batch = graphium_collate_fn(temp_elements)
        return list(batch.values())
    elif isinstance(elements, Sequence) and not isinstance(elem, Sequence):
        temp_elements = [{"temp_key": elem} for elem in elements]
        batch = graphium_collate_fn(temp_elements)
        return batch["temp_key"]
    else:
        return default_collate(elements)


def collage_pyg_graph(pyg_graphs: Iterable[Union[Data, Dict]], batch_size_per_pack: Optional[int] = None):
    """
    Function to collate pytorch geometric graphs.
    Convert all numpy types to torch
    Convert edge indices to int64

    Parameters:
        pyg_graphs: Iterable of PyG graphs
        batch_size_per_pack: The number of graphs to pack together.
            This is useful for using packing with the Transformer,
    """

    # Calculate maximum number of nodes per graph in current batch
    num_nodes_list = []
    for pyg_graph in pyg_graphs:
        num_nodes_list.append(pyg_graph["num_nodes"])
    max_num_nodes_per_graph = max(num_nodes_list)

    pyg_batch = []
    for pyg_graph in pyg_graphs:
        for pyg_key in get_keys(pyg_graph):
            tensor = pyg_graph[pyg_key]

            # Convert numpy/scipy to Pytorch
            if isinstance(tensor, (ndarray, spmatrix)):
                tensor = torch.as_tensor(to_dense_array(tensor, tensor.dtype))

            # pad nodepair-level positional encodings
            if pyg_key.startswith("nodepair_"):
                pyg_graph[pyg_key] = pad_nodepairs(tensor, pyg_graph["num_nodes"], max_num_nodes_per_graph)
            else:
                pyg_graph[pyg_key] = tensor

        # Convert edge index to int64
        pyg_graph.edge_index = pyg_graph.edge_index.to(torch.int64)
        pyg_batch.append(pyg_graph)

    # Apply the packing at the mini-batch level. This is useful for using packing with the Transformer,
    # especially in the case of the large graphs being much larger than the small graphs.
    # CAREFUL!!! This changes the order of the graphs in the batch, without changing the order of the labels or other objects.
    # An error is raised temporarily.
    if batch_size_per_pack is not None:
        raise NotImplementedError(
            "Packing is not yet functional, as it changes the order of the graphs in the batch without changing the label order"
        )
        num_nodes = [g.num_nodes for g in pyg_batch]
        packed_graph_idx = fast_packing(num_nodes, batch_size_per_pack)

        # Get the node to pack indices and the mask
        pack_from_node_idx, pack_attn_mask = node_to_pack_indices_mask(packed_graph_idx, num_nodes)
        for pyg_graph in pyg_batch:
            pyg_graph.pack_from_node_idx = pack_from_node_idx
            pyg_graph.pack_attn_mask = pack_attn_mask

    return Batch.from_data_list(pyg_batch)


def pad_to_expected_label_size(labels: torch.Tensor, label_size: List[int]):
    """Determine difference of ``labels`` shape to expected shape `label_size` and pad
    with ``torch.nan`` accordingly.
    """
    if label_size == list(labels.shape):
        return labels

    missing_dims = len(label_size) - len(labels.shape)
    for _ in range(missing_dims):
        labels.unsqueeze(-1)

    pad_sizes = [(0, expected - actual) for expected, actual in zip(label_size, labels.shape)]
    pad_sizes = [item for before_after in pad_sizes for item in before_after]
    pad_sizes.reverse()

    if any([s < 0 for s in pad_sizes]):
        logger.warning(f"More labels available than expected. Will remove data to fit expected size.")

    return torch.nn.functional.pad(labels, pad_sizes, value=torch.nan)


def collate_pyg_graph_labels(pyg_labels: List[Data]):
    """
    Function to collate pytorch geometric labels.
    Convert all numpy types to torch

    Parameters:
        pyg_labels: Iterable of PyG label Data objects
    """
    pyg_batch = []
    for pyg_label in pyg_labels:
        for pyg_key in set(get_keys(pyg_label)) - set(["x", "edge_index"]):
            tensor = pyg_label[pyg_key]
            # Convert numpy/scipy to Pytorch
            if isinstance(tensor, (ndarray, spmatrix)):
                tensor = torch.as_tensor(to_dense_array(tensor, tensor.dtype))

            pyg_label[pyg_key] = tensor

        pyg_batch.append(pyg_label)

    return Batch.from_data_list(pyg_batch)


def get_expected_label_size(label_data: Data, task: str, label_size: List[int]):
    """Determines expected label size based on the specfic graph properties
    and the number of targets in the task-dataset.
    """
    if task.startswith("graph_"):
        num_labels = 1
    elif task.startswith("node_"):
        num_labels = label_data.x.size(0)
    elif task.startswith("edge_"):
        num_labels = label_data.edge_index.size(1)
    elif task.startswith("nodepair_"):
        raise NotImplementedError()
    return [num_labels] + label_size


def collate_labels(
    labels: List[Data],
    labels_size_dict: Optional[Dict[str, Any]] = None,
    labels_dtype_dict: Optional[Dict[str, Any]] = None,
):
    """Collate labels for multitask learning.

    Parameters:
        labels: List of labels
        labels_size_dict: Dict of the form Dict[tasks, sizes] which has task names as keys
            and the size of the label tensor as value. The size of the tensor corresponds to how many
            labels/values there are to predict for that task.
        labels_dtype_dict:
            (Note): This is an attribute of the `MultitaskDataset`.
            A dictionary of the form Dict[tasks, dtypes] which has task names as keys
            and the dtype of the label tensor as value. This is necessary to ensure the missing labels are added with NaNs of the right dtype

    Returns:
        A dictionary of the form Dict[tasks, labels] where tasks is the name of the task and labels
        is a tensor of shape (batch_size, *labels_size_dict[task]).
    """
    if labels_size_dict is not None:
        for this_label in labels:
            for task in labels_size_dict.keys():
                labels_size_dict[task] = list(labels_size_dict[task])
                if len(labels_size_dict[task]) >= 2:
                    labels_size_dict[task] = labels_size_dict[task][1:]
                elif not task.startswith("graph_"):
                    labels_size_dict[task] = [1]
            label_keys_set = set(get_keys(this_label))
            empty_task_labels = set(labels_size_dict.keys()) - label_keys_set
            for task in empty_task_labels:
                labels_size_dict[task] = get_expected_label_size(this_label, task, labels_size_dict[task])
                dtype = labels_dtype_dict[task]
                this_label[task] = torch.full([*labels_size_dict[task]], torch.nan, dtype=dtype)

            for task in label_keys_set - set(["x", "edge_index"]) - empty_task_labels:
                labels_size_dict[task] = get_expected_label_size(this_label, task, labels_size_dict[task])

                if not isinstance(this_label[task], (torch.Tensor)):
                    this_label[task] = torch.as_tensor(this_label[task])

                # Ensure explicit task dimension also for single task labels
                if len(this_label[task].shape) == 1:
                    # Distinguish whether target dim or entity dim is missing
                    if labels_size_dict[task][0] == this_label[task].shape[0]:
                        # num graphs/nodes/edges/nodepairs already matching
                        this_label[task] = this_label[task].unsqueeze(1)
                    else:
                        # data lost unless entity dim is supposed to be 1
                        if labels_size_dict[task][0] == 1:
                            this_label[task] = this_label[task].unsqueeze(0)
                        else:
                            raise ValueError(
                                f"Labels for {labels_size_dict[task][0]} nodes/edges/nodepairs expected, got 1."
                            )

                this_label[task] = pad_to_expected_label_size(this_label[task], labels_size_dict[task])

    return collate_pyg_graph_labels(labels)


def pad_nodepairs(pe: torch.Tensor, num_nodes: int, max_num_nodes_per_graph: int):
    """
    This function zero-pads nodepair-level positional encodings to conform with the batching logic.

    Parameters:
        pe (torch.Tensor, [num_nodes, num_nodes, num_feat]): Nodepair pe
        num_nodes (int): Number of nodes of processed graph
        max_num_nodes_per_graph (int): Maximum number of nodes among graphs in current batch

    Returns:
        padded_pe (torch.Tensor, [num_nodes, max_num_nodes_per_graph, num_feat]): padded nodepair pe tensor
    """
    padded_pe = torch.zeros((num_nodes, max_num_nodes_per_graph, pe.size(-1)), dtype=pe.dtype)
    padded_pe[:, :num_nodes] = pe[:, :num_nodes]
    # Above, pe[:, :num_nodes] in the rhs is needed to "overwrite" zero-padding from previous epoch

    return padded_pe


# graphium/data/datamodule.py

import tempfile
from contextlib import redirect_stderr, redirect_stdout
from typing import Type, List, Dict, Union, Any, Callable, Optional, Tuple, Iterable, Literal

import os
from functools import partial
import importlib.resources
import zipfile
from copy import deepcopy
import time
import gc

import platformdirs
import re
from graphium.data.utils import get_keys

from loguru import logger
import fsspec
import omegaconf

import pandas as pd
import numpy as np
import datamol as dm
from tqdm import tqdm
import os.path as osp
from fastparquet import ParquetFile

from sklearn.model_selection import train_test_split

import lightning
from lightning.pytorch.trainer.states import RunningStage

import torch
from torch.utils.data.dataloader import DataLoader, Dataset
from torch.utils.data import Subset

from graphium.utils import fs
from graphium.features import (
    mol_to_graph_dict,
    GraphDict,
    mol_to_pyggraph,
)

from graphium.data.sampler import DatasetSubSampler
from graphium.data.utils import graphium_package_path, found_size_mismatch
from graphium.utils.arg_checker import check_arg_iterator
from graphium.utils.hashing import get_md5_hash
from graphium.data.smiles_transform import (
    did_featurization_fail,
    BatchingSmilesTransform,
    smiles_to_unique_mol_ids,
)
from graphium.data.collate import graphium_collate_fn
import graphium.data.dataset as Datasets
from graphium.data.normalization import LabelNormalization
from graphium.data.multilevel_utils import extract_labels

torch.multiprocessing.set_sharing_strategy("file_system")


PCQM4M_meta = {
    "num tasks": 1,
    "eval metric": "mae",
    "download_name": "pcqm4m_kddcup2021",
    "url": "https://dgl-data.s3-accelerate.amazonaws.com/dataset/OGB-LSC/pcqm4m_kddcup2021.zip",  # TODO: Allow PyG
    "data type": "mol",
    "has_node_attr": True,
    "has_edge_attr": True,
    "task type": "regression",
    "num classes": -1,
    "split": "scaffold",
    "additional node files": "None",
    "additional edge files": "None",
    "binary": False,
    "version": 1,
}

PCQM4Mv2_meta = deepcopy(PCQM4M_meta)
PCQM4Mv2_meta.update(
    {
        "download_name": "pcqm4m-v2",
        "url": "https://dgl-data.s3-accelerate.amazonaws.com/dataset/OGB-LSC/pcqm4m-v2.zip",  # TODO: Allow PyG
        "version": 2,
    }
)


class BaseDataModule(lightning.LightningDataModule):
    def __init__(
        self,
        batch_size_training: int = 16,
        batch_size_inference: int = 16,
        batch_size_per_pack: Optional[int] = None,
        num_workers: int = 0,
        pin_memory: bool = True,
        persistent_workers: bool = False,
        collate_fn: Optional[Callable] = None,
    ):
        """
        base dataset module for all datasets (to be inherented)

        Parameters:
            batch_size_training: batch size for training
            batch_size_inference: batch size for inference
            num_workers: number of workers for data loading
            pin_memory: whether to pin memory
            persistent_workers: whether to use persistent workers
            collate_fn: collate function for batching
        """
        super().__init__()

        self.batch_size_training = batch_size_training
        self.batch_size_inference = batch_size_inference
        self.batch_size_per_pack = batch_size_per_pack
        if self.batch_size_per_pack is not None:
            # Check that batch_size_per_pack is a divisor of batch_size_training and batch_size_inference
            assert (
                self.batch_size_training % self.batch_size_per_pack == 0
            ), f"batch_size_training must be a multiple of batch_size_per_pack, provided batch_size_training={self.batch_size_training}, batch_size_per_pack={self.batch_size_per_pack}"
            assert (
                self.batch_size_inference % self.batch_size_per_pack == 0
            ), f"batch_size_inference must be a multiple of batch_size_per_pack, provided batch_size_inference={self.batch_size_inference}, batch_size_per_pack={self.batch_size_per_pack}"

        self.num_workers = num_workers
        self.pin_memory = pin_memory
        self.persistent_workers = persistent_workers

        self.collate_fn = self.get_collate_fn(collate_fn)

        self.train_ds = None
        self.val_ds = None
        self.test_ds = None
        self._predict_ds = None

        self._data_is_prepared = False

    def prepare_data(self):
        raise NotImplementedError()

    def setup(self):
        raise NotImplementedError()

    def train_dataloader(self, **kwargs):
        """
        return the training dataloader
        """
        return self.get_dataloader(
            dataset=self.train_ds,  # type: ignore
            shuffle=True,
            stage=RunningStage.TRAINING,
            **kwargs,
        )

    def val_dataloader(self, **kwargs):
        r"""
        return the validation dataloader
        """
        return self.get_dataloader(
            dataset=self.val_ds,  # type: ignore
            shuffle=False,
            stage=RunningStage.VALIDATING,
            **kwargs,
        )

    def test_dataloader(self, **kwargs):
        r"""
        return the test dataloader
        """
        return self.get_dataloader(
            dataset=self.test_ds,  # type: ignore
            shuffle=False,
            stage=RunningStage.TESTING,
            **kwargs,
        )

    def predict_dataloader(self, **kwargs):
        """
        return the dataloader for prediction
        """
        return self.get_dataloader(
            dataset=self.predict_ds,  # type: ignore
            shuffle=False,
            stage=RunningStage.PREDICTING,
            **kwargs,
        )

    def get_collate_fn(self, collate_fn):
        if collate_fn is None:
            # Some values become `inf` when changing data type. `mask_nan` deals with that
            collate_fn = partial(
                graphium_collate_fn, mask_nan=0, batch_size_per_pack=self.batch_size_per_pack
            )
            collate_fn.__name__ = graphium_collate_fn.__name__

        return collate_fn

    @property
    def is_prepared(self):
        raise NotImplementedError()

    @property
    def is_setup(self):
        raise NotImplementedError()

    @property
    def num_node_feats(self):
        raise NotImplementedError()

    @property
    def num_edge_feats(self):
        raise NotImplementedError()

    @property
    def predict_ds(self):
        """Get the dataset used for the prediction"""
        if self._predict_ds is None:
            return self.test_ds
        else:
            return self._predict_ds

    @property
    def get_num_workers(self):
        """
        get the number of workers to use
        """
        if self.num_workers == -1:
            num_workers = os.cpu_count()
            num_workers = num_workers if num_workers is not None else 0
        else:
            num_workers = self.num_workers
        return num_workers

    @predict_ds.setter
    def predict_ds(self, value):
        """Set the dataset for the prediction"""
        self._predict_ds = value

    def get_fake_graph(self):
        raise NotImplementedError()

    # Private methods

    @staticmethod
    def _read_csv(
        path: str,
        **kwargs,
    ) -> pd.DataFrame:
        """
        private method for reading a csv file
        Parameters:
            path: path to the csv file
            kwargs: keyword arguments for pd.read_csv
        Returns:
            pd.DataFrame: the panda dataframe storing molecules
        """

        path = str(path)

        if path.endswith((".csv", ".csv.gz", ".csv.zip", ".csv.bz2")):
            sep = ","
        elif path.endswith((".tsv", ".tsv.gz", ".tsv.zip", ".tsv.bz2")):
            sep = "\t"
        else:
            raise ValueError(f"unsupported file `{path}`")
        kwargs.setdefault("sep", sep)

        if path.startswith("graphium://"):
            path = graphium_package_path(path)

        df = pd.read_csv(path, **kwargs)
        return df

    @staticmethod
    def _get_data_file_type(path):
        # Extract the extension
        name, ext = os.path.splitext(path)

        # support compressed files
        _, ext2 = os.path.splitext(name)
        if ext2 != "":
            ext = f"{ext2}{ext}"

        if ext.endswith((".parquet")):  # Support parquet files. Compression is implicit
            return "parquet"
        elif ".sdf" in ext:  # support compressed sdf files
            return "sdf"
        elif ".csv" in ext:  # support compressed csv files
            return "csv"
        elif ".tsv" in ext:  # support compressed tsv files
            return "tsv"
        elif ".pkl" in ext:  # support compressed pickle files
            return "pkl"
        elif ext.endswith(".pt"):  # Pytorch tensor files, used for storing index splits
            return "pt"
        else:
            raise ValueError(f"unsupported file `{path}`")

    @staticmethod
    def _get_table_columns(path: str) -> List[str]:
        """
        Get the columns of a table without reading all the data.
        Might be slow to decompress the file if the file is compressed.

        Parameters:
            path: path to the table file

        Returns:
            List[str]: the column names
        """

        datafile_type = BaseDataModule._get_data_file_type(path)

        if datafile_type == "parquet":
            # Read the schema of a parquet file
            file = ParquetFile(path)
            schema = file.pandas_metadata["columns"]
            column_names = [s["name"] for s in schema if s["name"] is not None]
        elif datafile_type == "sdf":
            df = BaseDataModule._read_sdf(path, max_num_mols=5, discard_invalud=True, n_jobs=1)
            column_names = df.columns
        elif datafile_type in ["csv", "tsv"]:
            # Read the schema of a csv / tsv file
            df = BaseDataModule._read_csv(path, nrows=5)
            column_names = df.columns
        return column_names

    @staticmethod
    def _read_parquet(path, **kwargs):
        kwargs.pop("dtype", None)  # Only useful for csv
        column_names = BaseDataModule._get_table_columns(path)

        # Change the 'usecols' parameter to 'columns'
        columns = kwargs.pop("columns", None)
        if "usecols" in kwargs.keys():
            assert columns is None, "Ambiguous value of `columns`"
            columns = kwargs.pop("usecols")
        if columns is None:
            columns = column_names
        for column in columns:
            assert (
                column in column_names
            ), f"Column `{column}` is not in the parquet file with columns {column_names}"

        # Read the parquet file per column, and convert the data to float16 to reduce memory consumption
        all_series = {}
        progress = tqdm(columns)
        for col in progress:
            # Read single column
            progress.set_description(f"Reading parquet column `{col}`")
            this_series = pd.read_parquet(path, columns=[col], engine="fastparquet", **kwargs)[col]

            # Check if the data is float
            first_elem = this_series.values[0]
            is_float = False
            if isinstance(first_elem, (list, tuple)):
                is_float = isinstance(first_elem[0], float) or (first_elem[0].dtype.kind == "f")
            elif isinstance(first_elem, np.ndarray):
                is_float = isinstance(first_elem, float) or (first_elem.dtype.kind == "f")

            # Convert floats to float16
            if is_float:
                if isinstance(first_elem, (np.ndarray, list)):
                    this_series.update(
                        pd.Series([np.asarray(elem).astype(np.float16) for elem in this_series])
                    )
                else:
                    this_series = this_series.astype(np.float16)

            all_series[col] = this_series
            gc.collect()  # Reset memory after each column

        # Merge columns into a dataframe
        df = pd.concat(all_series, axis=1)

        return df

    @staticmethod
    def _read_sdf(path: str, mol_col_name: str = "_rdkit_molecule_obj", **kwargs):
        r"""
        read a given sdf file into a pandas dataframe
        uses datamol.read_sdf to read the sdf file
        Parameters:
            path: path to the sdf file
            mol_col_name: name of the column containing the molecule object
            kwargs: arguments to pass to datamol.read_sdf
        Returns:
            pandas dataframe containing the molecules and their conformer coordinates from the sdf file

        Note: please change mol_col_name to be the column name corresoonding to the molecule object in your sdf file
        """

        # Set default arguments for reading the SDF
        kwargs.setdefault("smiles_column", "smiles")
        kwargs.setdefault("sanitize", False)
        kwargs.setdefault("include_private", True)
        kwargs.setdefault("include_computed", True)
        kwargs.setdefault("remove_hs", True)
        kwargs.setdefault("n_jobs", -1)
        kwargs.setdefault("max_num_mols", kwargs.pop("sample_size", None))
        kwargs.setdefault("discard_invalid", False)

        # Get the interesting columns
        mol_cols = mol_col_name  # adjust this to be the column name in your sdf file
        kwargs.setdefault("mol_column", mol_cols)
        usecols = kwargs.pop("usecols", None)
        dtype = kwargs.pop("dtype", None)
        smiles_col = kwargs["smiles_column"]

        # Read the SDF
        df = dm.read_sdf(path, as_df=True, **kwargs)

        # Keep only the columns needed
        if usecols is not None:
            df = df[usecols + [mol_cols]]

        # Convert the dtypes
        if dtype is not None:
            label_columns = list(set(usecols) - set([mol_cols, smiles_col]))
            dtype_mapper = {col: dtype for col in label_columns}
            df.astype(dtype=dtype_mapper, copy=False)

        return df

    @staticmethod
    def _glob(path: str) -> List[str]:
        """
        glob a given path
        Parameters:
            path: path to glob
        Returns:
            List[str]: list of paths
        """
        files = dm.fs.glob(path)
        files = [f.replace("file://", "") for f in files]
        return files

    def _read_table(self, path: str, **kwargs) -> pd.DataFrame:
        """
        a general read file function which determines if which function to use, either _read_csv or _read_parquet
        Parameters:
            path: path to the file to read
            kwargs: keyword arguments for pd.read_csv or pd.read_parquet
        Returns:
            pd.DataFrame: the panda dataframe storing molecules
        """
        files = self._glob(path)
        if len(files) == 0:
            raise FileNotFoundError(f"No such file or directory `{path}`")

        if len(files) > 1:
            files = tqdm(sorted(files), desc=f"Reading files at `{path}`")
        dfs = []
        for file in files:
            file_type = self._get_data_file_type(file)
            if file_type == "parquet":
                df = self._read_parquet(file, **kwargs)
            elif file_type == "sdf":  # support compressed sdf files
                df = self._read_sdf(file, **kwargs)
            elif file_type in ["csv", "tsv"]:  # support compressed csv and tsv files
                df = self._read_csv(file, **kwargs)
            else:
                raise ValueError(f"unsupported file `{file}`")
            dfs.append(df)

        return pd.concat(dfs, ignore_index=True)

    def get_dataloader_kwargs(self, stage: RunningStage, shuffle: bool, **kwargs) -> Dict[str, Any]:
        """
        Get the options for the dataloader depending on the current stage.

        Parameters:
            stage: Whether in Training, Validating, Testing, Sanity-checking, Predicting, or Tuning phase.
            shuffle: set to ``True`` to have the data reshuffled at every epoch.

        Returns:
            Arguments to pass to the `DataLoader` during initialization
        """
        loader_kwargs = {}

        # Get batch size and IPU options for training set
        # if stage in [RunningStage.TRAINING, RunningStage.TUNING]:
        if stage in [RunningStage.TRAINING]:
            loader_kwargs["batch_size"] = self.batch_size_training

        # Get batch size and IPU options for validation / testing sets
        elif stage in [RunningStage.VALIDATING, RunningStage.TESTING, RunningStage.PREDICTING]:
            loader_kwargs["batch_size"] = self.batch_size_inference
        else:
            raise ValueError(f"Wrong value for `stage`. Provided `{stage}`")

        # Set default parameters
        loader_kwargs["shuffle"] = shuffle
        loader_kwargs["collate_fn"] = self.collate_fn
        loader_kwargs["num_workers"] = self.get_num_workers
        loader_kwargs["pin_memory"] = self.pin_memory
        loader_kwargs["persistent_workers"] = self.persistent_workers

        # Update from provided parameters
        loader_kwargs.update(**kwargs)

        return loader_kwargs

    def get_dataloader(self, dataset: Dataset, shuffle: bool, stage: RunningStage) -> DataLoader:
        """
        Get the dataloader for a given dataset

        Parameters:
            dataset: The dataset from which to load the data
            shuffle: set to ``True`` to have the data reshuffled at every epoch.
            stage: Whether in Training, Validating, Testing, Sanity-checking, Predicting, or Tuning phase.

        Returns:
            The dataloader to sample from
        """
        kwargs = self.get_dataloader_kwargs(stage=stage, shuffle=shuffle)
        return self._dataloader(dataset=dataset, shuffle=shuffle, stage=stage, **kwargs)

    def _dataloader(self, dataset: Dataset, **kwargs) -> DataLoader:
        r"""
        Get a dataloader for a given dataset
        Parameters:
            dataset: The dataset from which to load the data
            kwargs: keyword arguments for DataLoader
        Returns:
            The dataloader to sample from
        """

        loader = DataLoader(
            dataset=dataset,
            **kwargs,
        )

        return loader

    def get_max_num_nodes_datamodule(self, stages: Optional[List[str]] = None) -> int:
        """
        Get the maximum number of nodes across all datasets from the datamodule

        Parameters:
            datamodule: The datamodule from which to extract the maximum number of nodes
            stages: The stages from which to extract the max num nodes.
                Possible values are ["train", "val", "test", "predict"].
                If None, all stages are considered.

        Returns:
            max_num_nodes: The maximum number of nodes across all datasets from the datamodule
        """

        allowed_stages = ["train", "val", "test", "predict"]
        if stages is None:
            stages = allowed_stages
        for stage in stages:
            assert stage in allowed_stages, f"stage value `{stage}` not allowed."

        max_num_nodes = 0
        # Max number of nodes in the training dataset
        if (self.train_ds is not None) and ("train" in stages):
            logger.info("Max num nodes being calcuated train")
            max_num_nodes = max(max_num_nodes, self.train_ds.max_num_nodes_per_graph)

        # Max number of nodes in the validation dataset
        if (
            (self.val_ds is not None)
            and ("val" in stages)
            and (self.val_ds.max_num_nodes_per_graph is not None)
        ):
            logger.info("Max num nodes being calcuated val")
            max_num_nodes = max(max_num_nodes, self.val_ds.max_num_nodes_per_graph)

        # Max number of nodes in the test dataset
        if (
            (self.test_ds is not None)
            and ("test" in stages)
            and (self.test_ds.max_num_nodes_per_graph is not None)
        ):
            logger.info("Max num nodes being calcuated test")
            max_num_nodes = max(max_num_nodes, self.test_ds.max_num_nodes_per_graph)

        # Max number of nodes in the predict dataset
        if (
            (self.predict_ds is not None)
            and ("predict" in stages)
            and (self.predict_ds.max_num_nodes_per_graph is not None)
        ):
            max_num_nodes = max(max_num_nodes, self.predict_ds.max_num_nodes_per_graph)

        return max_num_nodes

    def get_max_num_edges_datamodule(self, stages: Optional[List[str]] = None) -> int:
        """
        Get the maximum number of edges across all datasets from the datamodule

        Parameters:
            datamodule: The datamodule from which to extract the maximum number of nodes
            stages: The stages from which to extract the max num nodes.
                Possible values are ["train", "val", "test", "predict"].
                If None, all stages are considered.

        Returns:
            max_num_edges: The maximum number of edges across all datasets from the datamodule
        """

        allowed_stages = ["train", "val", "test", "predict"]
        if stages is None:
            stages = allowed_stages
        for stage in stages:
            assert stage in allowed_stages, f"stage value `{stage}` not allowed."

        max_num_edges = 0
        # Max number of nodes/edges in the training dataset
        if (
            (self.train_ds is not None)
            and ("train" in stages)
            and (self.train_ds.max_num_edges_per_graph is not None)
        ):
            max_num_edges = max(max_num_edges, self.train_ds.max_num_edges_per_graph)

        # Max number of nodes/edges in the validation dataset
        if (
            (self.val_ds is not None)
            and ("val" in stages)
            and (self.val_ds.max_num_edges_per_graph is not None)
        ):
            max_num_edges = max(max_num_edges, self.val_ds.max_num_edges_per_graph)

        # Max number of nodes/edges in the test dataset
        if (
            (self.test_ds is not None)
            and ("test" in stages)
            and (self.test_ds.max_num_edges_per_graph is not None)
        ):
            max_num_edges = max(max_num_edges, self.test_ds.max_num_edges_per_graph)

        # Max number of nodes/edges in the predict dataset
        if (
            (self.predict_ds is not None)
            and ("predict" in stages)
            and (self.predict_ds.max_num_edges_per_graph is not None)
        ):
            max_num_edges = max(max_num_edges, self.predict_ds.max_num_edges_per_graph)

        return max_num_edges


class DatasetProcessingParams:
    def __init__(
        self,
        task_level: Optional[str] = None,
        df: Optional[pd.DataFrame] = None,
        df_path: Optional[Union[str, os.PathLike]] = None,
        smiles_col: Optional[str] = None,
        label_cols: List[str] = None,
        weights_col: Optional[str] = None,  # Not needed
        weights_type: Optional[str] = None,  # Not needed
        idx_col: Optional[str] = None,
        mol_ids_col: Optional[str] = None,
        sample_size: Union[int, float, Type[None]] = None,
        split_val: float = 0.2,
        split_test: float = 0.2,
        seed: int = None,
        epoch_sampling_fraction: float = 1.0,
        splits_path: Optional[Union[str, os.PathLike]] = None,
        split_names: Optional[List[str]] = ["train", "val", "test"],
        label_normalization: Optional[Union[Dict[str, Any], omegaconf.DictConfig]] = None,
    ):
        """
        object to store the parameters for the dataset processing
        Parameters:
            task_level: The task level, wether it is graph, node, edge or nodepair
            df: The dataframe containing the data
            df_path: The path to the dataframe containing the data
            smiles_col: The column name of the smiles
            label_cols: The column names of the labels
            weights_col: The column name of the weights
            weights_type: The type of weights
            idx_col: The column name of the indices
            mol_ids_col: The column name of the molecule ids
            sample_size: The size of the sample
            split_val: The fraction of the data to use for validation
            split_test: The fraction of the data to use for testing
            seed: The seed to use for the splits and subsampling
            splits_path: The path to the splits
        """

        if df is None and df_path is None:
            raise ValueError("Either `df` or `df_path` must be provided")
        if epoch_sampling_fraction <= 0 or epoch_sampling_fraction > 1:
            raise ValueError("The value of epoch_sampling_fraction must be in the range of (0, 1].")

        self.df = df
        self.task_level = task_level
        self.df_path = df_path
        self.smiles_col = smiles_col
        self.label_cols = label_cols
        self.weights_col = weights_col
        self.weights_type = weights_type
        self.idx_col = idx_col
        self.mol_ids_col = mol_ids_col
        self.sample_size = sample_size
        self.split_val = split_val
        self.split_test = split_test
        self.seed = seed
        self.splits_path = splits_path
        self.split_names = split_names
        self.label_normalization = label_normalization
        self.epoch_sampling_fraction = epoch_sampling_fraction


class IPUDataModuleModifier:
    def __init__(
        self,
        ipu_inference_opts: Optional["poptorch.Options"] = None,
        ipu_training_opts: Optional["poptorch.Options"] = None,
        ipu_dataloader_training_opts: Optional["IPUDataloaderOptions"] = None,
        ipu_dataloader_inference_opts: Optional["IPUDataloaderOptions"] = None,
        *args,
        **kwargs,
    ) -> None:
        r"""
        wrapper functions from the a `DataModule` to support IPU and IPU options To be used in dual inheritance, for example:
        ```
        IPUDataModule(BaseDataModule, IPUDataModuleModifier):
            def __init__(self, **kwargs):
                BaseDataModule.__init__(self, **kwargs)
                IPUDataModuleModifier.__init__(self, **kwargs)
        ```

        Parameters:
            ipu_inference_opts: Options for the IPU in inference mode. Ignore if not using IPUs
            ipu_training_opts: Options for the IPU in training mode. Ignore if not using IPUs
            ipu_dataloader_kwargs_train_val: Options for the dataloader for the IPU. Ignore if not using IPUs
            ipu_dataloader_kwargs_test: Options for the dataloader for the IPU. Ignore if not using IPUs
            args: Arguments for the `DataModule`
            kwargs: Keyword arguments for the `DataModule`
        """
        self.ipu_inference_opts = ipu_inference_opts
        self.ipu_training_opts = ipu_training_opts
        self.ipu_dataloader_training_opts = ipu_dataloader_training_opts
        self.ipu_dataloader_inference_opts = ipu_dataloader_inference_opts

    def _dataloader(self, dataset: Dataset, **kwargs) -> "poptorch.DataLoader":
        """
        Get a poptorch dataloader for a given dataset
        Parameters:
            dataset: The dataset to use
            kwargs: Keyword arguments for the dataloader
        Returns:
            The poptorch dataloader
        """

        # Use regular Dataloader if no IPUs
        if ("ipu_options" not in kwargs.keys()) or (kwargs["ipu_options"] is None):
            raise ValueError(f"No IPU options provided.")

        # Initialize the IPU dataloader
        from graphium.ipu.ipu_dataloader import create_ipu_dataloader

        loader = create_ipu_dataloader(
            dataset=dataset,
            **kwargs,
        )

        return loader


class MultitaskFromSmilesDataModule(BaseDataModule, IPUDataModuleModifier):
    def __init__(
        self,
        task_specific_args: Union[DatasetProcessingParams, Dict[str, Any]],
        cache_data_path: Optional[Union[str, os.PathLike]] = None,
        processed_graph_data_path: Optional[Union[str, os.PathLike]] = None,
        featurization: Optional[Union[Dict[str, Any], omegaconf.DictConfig]] = None,
        batch_size_training: int = 16,
        batch_size_inference: int = 16,
        batch_size_per_pack: Optional[int] = None,
        num_workers: int = 0,
        pin_memory: bool = True,
        persistent_workers: bool = False,
        featurization_n_jobs: int = -1,
        featurization_progress: bool = False,
        featurization_backend: str = "loky",
        featurization_batch_size: int = 1000,
        collate_fn: Optional[Callable] = None,
        prepare_dict_or_graph: str = "pyg:graph",
        **kwargs,
    ):
        """
        only for parameters beginning with task_*, we have a dictionary where the key is the task name
        and the value is specified below.
        Parameters:
            task_df: (value) a dataframe
            task_df_path: (value) a path to a dataframe to load (CSV file). `df` takes precedence over
                `df_path`.
            task_smiles_col: (value) Name of the SMILES column. If set to `None`, it will look for
                a column with the word "smile" (case insensitive) in it.
                If no such column is found, an error will be raised.
            task_label_cols: (value) Name of the columns to use as labels, with different options.

                - `list`: A list of all column names to use
                - `None`: All the columns are used except the SMILES one.
                - `str`: The name of the single column to use
                - `*str`: A string starting by a `*` means all columns whose name
                  ends with the specified `str`
                - `str*`: A string ending by a `*` means all columns whose name
                  starts with the specified `str`

            task_weights_col: (value) Name of the column to use as sample weights. If `None`, no
                weights are used. This parameter cannot be used together with `weights_type`.
            task_weights_type: (value) The type of weights to use. This parameter cannot be used together with `weights_col`.
                **It only supports multi-label binary classification.**

                Supported types:

                - `None`: No weights are used.
                - `"sample_balanced"`: A weight is assigned to each sample inversely
                    proportional to the number of positive value. If there are multiple
                    labels, the product of the weights is used.
                - `"sample_label_balanced"`: Similar to the `"sample_balanced"` weights,
                    but the weights are applied to each element individually, without
                    computing the product of the weights for a given sample.

            task_idx_col: (value) Name of the columns to use as indices. Unused if set to None.
            task_sample_size: (value)

                - `int`: The maximum number of elements to take from the dataset.
                - `float`: Value between 0 and 1 representing the fraction of the dataset to consider
                - `None`: all elements are considered.
            task_split_val: (value) Ratio for the validation split.
            task_split_test: (value) Ratio for the test split.
            task_seed: (value) Seed to use for the random split and subsampling. More complex splitting strategy
                should be implemented.
            task_splits_path: (value) A path a CSV file containing indices for the splits. The file must contains
                3 columns "train", "val" and "test". It takes precedence over `split_val` and `split_test`.

            cache_data_path: path where to save or reload the cached data. The path can be
                remote (S3, GS, etc).
            featurization: args to apply to the SMILES to Graph featurizer.
            batch_size_training: batch size for training and val dataset.
            batch_size_inference: batch size for test dataset.
            num_workers: Number of workers for the dataloader. Use -1 to use all available
                cores.
            pin_memory: Whether to pin on paginated CPU memory for the dataloader.
            featurization_n_jobs: Number of cores to use for the featurization.
            featurization_progress: whether to show a progress bar during featurization.
            featurization_backend: The backend to use for the molecular featurization.

                - "multiprocessing": Found to cause less memory issues.
                - "loky": joblib's Default. Found to cause memory leaks.
                - "threading": Found to be slow.
            featurization_batch_size: Batch size to use for the featurization.

            collate_fn: A custom torch collate function. Default is to `graphium.data.graphium_collate_fn`
            prepare_dict_or_graph: Whether to preprocess all molecules as Graph dict or PyG graphs.
                Possible options:

                - "pyg:dict": Process molecules as a `dict`. It's faster and requires less RAM during
                  pre-processing. It is slower during training with with `num_workers=0` since
                  pyg `Data` will be created during data-loading, but faster with large
                  `num_workers`, and less likely to cause memory issues with the parallelization.
                - "pyg:graph": Process molecules as `pyg.data.Data`.
        """
        BaseDataModule.__init__(
            self,
            batch_size_training=batch_size_training,
            batch_size_inference=batch_size_inference,
            batch_size_per_pack=batch_size_per_pack,
            num_workers=num_workers,
            pin_memory=pin_memory,
            persistent_workers=persistent_workers,
            collate_fn=collate_fn,
        )
        IPUDataModuleModifier.__init__(self, **kwargs)

        self.task_specific_args = task_specific_args

        self.task_dataset_processing_params = {}
        for task, ds_args in task_specific_args.items():
            if not isinstance(ds_args, DatasetProcessingParams):
                # This is needed as long as not all classes have been migrated
                # to use the new `DatasetProcessingParams` class
                ds_args = DatasetProcessingParams(**ds_args)

            key = self._get_task_key(ds_args.task_level, task)
            self.task_dataset_processing_params[key] = ds_args

        self.sampler_task_dict = {
            task: self.task_dataset_processing_params[task].epoch_sampling_fraction
            for task in self.task_dataset_processing_params.keys()
        }

        self.featurization_n_jobs = featurization_n_jobs
        self.featurization_progress = featurization_progress
        self.featurization_backend = featurization_backend
        self.featurization_batch_size = featurization_batch_size

        self.task_train_indices = None
        self.task_val_indices = None
        self.task_test_indices = None

        self.single_task_datasets = None
        self.train_singletask_datasets = None
        self.val_singletask_datasets = None
        self.test_singletask_datasets = None

        self.train_ds = None
        self.val_ds = None
        self.test_ds = None

        self.cache_data_path = cache_data_path
        self.processed_graph_data_path = processed_graph_data_path

        self.load_from_file = processed_graph_data_path is not None

        self.task_norms = {}

        if featurization is None:
            featurization = {}

        # Whether to transform the smiles into a pyg `Data` graph or a dictionary compatible with pyg
        if prepare_dict_or_graph == "pyg:dict":
            self.smiles_transformer = partial(mol_to_graph_dict, **featurization)
        elif prepare_dict_or_graph == "pyg:graph":
            self.smiles_transformer = partial(mol_to_pyggraph, **featurization)
        else:
            raise ValueError(
                f"`prepare_dict_or_graph` should be either 'pyg:dict' or 'pyg:graph', Provided: `{prepare_dict_or_graph}`"
            )
        self.data_hash = self.get_data_hash()

    def _get_task_key(self, task_level: str, task: str):
        task_prefix = f"{task_level}_"
        if not task.startswith(task_prefix):
            task = task_prefix + task
        return task

    def prepare_data(self):
        """Called only from a single process in distributed settings. Steps:

        - If each cache is set and exists, reload from cache and return. Otherwise,
        - For each single-task dataset:
            - Load its dataframe from a path (if provided)
            - Subsample the dataframe
            - Extract the smiles, labels from the dataframe
        - In the previous step, we were also able to get the unique smiles, which we use to compute the features
        - For each single-task dataframe and associated data (smiles, labels, etc.):
            - Filter out the data corresponding to molecules which failed featurization.
            - Create a corresponding SingletaskDataset
            - Split the SingletaskDataset according to the task-specific splits for train, val and test
        """

        def has_atoms_after_h_removal(smiles):
            # Remove all 'H' characters from the SMILES
            smiles_without_h = re.sub("H", "", smiles)
            # Check if any letters are remaining in the modified string
            has_atoms = bool(re.search("[a-zA-Z]", smiles_without_h))
            if has_atoms == False:
                logger.info(f"Removed Hydrogen molecule: {smiles}")
            return has_atoms

        if self._data_is_prepared:
            logger.info("Data is already prepared. Skipping the preparation")
            return

        if self.load_from_file:
            if self._ready_to_load_all_from_file():
                self.get_label_statistics(self.processed_graph_data_path, self.data_hash, dataset=None)
                self._data_is_prepared = True
                return

        else:
            # If a path for data caching is provided, try to load from the path.
            # If successful, skip the data preparation.
            # For next task: load the single graph files for train, val and test data
            cache_data_exists = self.load_data_from_cache()
            # need to check if cache exist properly
            if cache_data_exists:
                self.get_label_statistics(self.cache_data_path, self.data_hash, dataset=None)
                self._data_is_prepared = True
                return

        """Load all single-task dataframes."""
        task_df = {}
        for task, args in self.task_dataset_processing_params.items():
            if args.label_normalization is None:
                args.label_normalization = {}
            label_normalization = LabelNormalization(**args.label_normalization)
            logger.info(f"Reading data for task '{task}'")
            if args.df is None:
                # Only load the useful columns, as some datasets can be very large when loading all columns.
                label_cols = self._parse_label_cols(
                    df=None, df_path=args.df_path, label_cols=args.label_cols, smiles_col=args.smiles_col
                )
                usecols = (
                    check_arg_iterator(args.smiles_col, enforce_type=list)
                    + label_cols
                    + check_arg_iterator(args.idx_col, enforce_type=list)
                    + check_arg_iterator(args.weights_col, enforce_type=list)
                )
                label_dtype = {col: np.float32 for col in label_cols}
                task_df[task] = self._read_table(args.df_path, usecols=usecols, dtype=label_dtype)

            else:
                label_cols = self._parse_label_cols(
                    df=args.df, df_path=None, label_cols=args.label_cols, smiles_col=args.smiles_col
                )
                task_df[task] = args.df
            task_df[task] = task_df[task]
            args.label_cols = label_cols
            self.task_norms[task] = label_normalization
        logger.info("Done reading datasets")

        """Subsample the data frames and extract the necessary data to create SingleTaskDatasets for each task (smiles, labels, extras)."""
        task_dataset_args = {}
        for task in task_df.keys():
            task_dataset_args[task] = {}

        for task, df in task_df.items():
            # Subsample all the dataframes
            sample_size = self.task_dataset_processing_params[task].sample_size
            df = self._sub_sample_df(df, sample_size, self.task_dataset_processing_params[task].seed)

            logger.info(f"Prepare single-task dataset for task '{task}' with {len(df)} data points.")

            logger.info("Filtering the molecules for Hydrogen")
            logger.info(f"Looking at column {df.columns[0]}")
            logger.info("Filtering done")
            # Extract smiles, labels, extras
            args = self.task_dataset_processing_params[task]
            smiles, labels, sample_idx, extras = self._extract_smiles_labels(
                df,
                task_level=args.task_level,
                smiles_col=args.smiles_col,
                label_cols=args.label_cols,
                idx_col=args.idx_col,
                weights_col=args.weights_col,
                weights_type=args.weights_type,
            )

            # Store the relevant information for each task's dataset
            task_dataset_args[task]["smiles"] = smiles
            task_dataset_args[task]["labels"] = labels
            task_dataset_args[task]["sample_idx"] = sample_idx
            task_dataset_args[task]["extras"] = extras

        """Convert SMILES to features (graphs, fingerprints, etc.) for the unique molecules found."""
        all_smiles = []
        all_tasks = []
        idx_per_task = {}
        total_len = 0
        for task, dataset_args in task_dataset_args.items():
            all_smiles.extend(dataset_args["smiles"])
            num_smiles = len(dataset_args["smiles"])
            idx_per_task[task] = (total_len, total_len + num_smiles)
            total_len += num_smiles
            for count in range(len(dataset_args["smiles"])):
                all_tasks.append(task)
        # Get all unique mol ids
        all_unique_mol_ids = smiles_to_unique_mol_ids(
            all_smiles,
            n_jobs=self.featurization_n_jobs,
            featurization_batch_size=self.featurization_batch_size,
            backend=self.featurization_backend,
        )
        _, unique_ids_idx, unique_ids_inv = np.unique(
            all_unique_mol_ids, return_index=True, return_inverse=True
        )

        smiles_to_featurize = [all_smiles[ii] for ii in unique_ids_idx]

        # Convert SMILES to features
        features, _ = self._featurize_molecules(smiles_to_featurize)

        # Store the features (including Nones, which will be filtered in the next step)
        for task in task_dataset_args.keys():
            task_dataset_args[task]["features"] = []
            task_dataset_args[task]["idx_none"] = []
        # Create a list of features matching up with the original smiles
        all_features = [features[unique_idx] for unique_idx in unique_ids_inv]

        # Add the features to the task-specific data
        for all_idx, task in enumerate(all_tasks):
            task_dataset_args[task]["features"].append(all_features[all_idx])

        """Filter data based on molecules which failed featurization. Create single task datasets as well."""
        self.single_task_datasets = {}
        for task, args in task_dataset_args.items():
            # Find out which molecule failed featurization, and filter them out
            idx_none = []
            for idx, (feat, labels, smiles) in enumerate(
                zip(args["features"], args["labels"], args["smiles"])
            ):
                if did_featurization_fail(feat) or found_size_mismatch(task, feat, labels, smiles):
                    idx_none.append(idx)
            this_unique_ids = all_unique_mol_ids[idx_per_task[task][0] : idx_per_task[task][1]]
            df, features, smiles, labels, sample_idx, extras, this_unique_ids = self._filter_none_molecules(
                idx_none,
                task_df[task],
                args["features"],
                args["smiles"],
                args["labels"],
                args["sample_idx"],
                args["extras"],
                this_unique_ids,
            )
            task_dataset_args[task]["smiles"] = smiles
            task_dataset_args[task]["labels"] = labels
            task_dataset_args[task]["features"] = features
            task_dataset_args[task]["sample_idx"] = sample_idx
            task_dataset_args[task]["extras"] = extras

            # We have the necessary components to create single-task datasets.
            self.single_task_datasets[task] = Datasets.SingleTaskDataset(
                features=task_dataset_args[task]["features"],
                labels=task_dataset_args[task]["labels"],
                smiles=task_dataset_args[task]["smiles"],
                unique_ids=this_unique_ids,
                indices=task_dataset_args[task]["sample_idx"],
                **task_dataset_args[task]["extras"],
            )

        """We split the data up to create train, val and test datasets"""
        self.task_train_indices = {}
        self.task_val_indices = {}
        self.task_test_indices = {}

        for task, df in task_df.items():
            train_indices, val_indices, test_indices = self._get_split_indices(
                len(df),
                split_val=self.task_dataset_processing_params[task].split_val,
                split_test=self.task_dataset_processing_params[task].split_test,
                split_seed=self.task_dataset_processing_params[task].seed,
                splits_path=self.task_dataset_processing_params[task].splits_path,
                split_names=self.task_dataset_processing_params[task].split_names,
                sample_idx=task_dataset_args[task]["sample_idx"],
            )
            self.task_train_indices[task] = train_indices
            self.task_val_indices[task] = val_indices
            self.task_test_indices[task] = test_indices

        (
            self.train_singletask_datasets,
            self.val_singletask_datasets,
            self.test_singletask_datasets,
        ) = self.get_subsets_of_datasets(
            self.single_task_datasets, self.task_train_indices, self.task_val_indices, self.task_test_indices
        )

        if self.load_from_file:
            self._save_data_to_files()

        # When a cache path is provided but no cache is found, save to cache
        elif (self.cache_data_path is not None) and (not cache_data_exists):
            self.save_data_to_cache()

        self._data_is_prepared = True

    def setup(
        self,
        stage: str = None,
        save_smiles_and_ids: bool = False,
    ):
        """
        Prepare the torch dataset. Called on every GPUs. Setting state here is ok.
        Parameters:
            stage (str): Either 'fit', 'test', or None.
        """

        # Can possibly get rid of setup because a single dataset will have molecules exclusively in train, val or test
        # Produce the label sizes to update the collate function
        labels_size = {}
        labels_dtype = {}
        if stage == "fit" or stage is None:
            if self.load_from_file:
                processed_train_data_path = self._path_to_load_from_file("train")
                assert self._data_ready_at_path(
                    processed_train_data_path
                ), "Loading from file + setup() called but training data not ready"
                processed_val_data_path = self._path_to_load_from_file("val")
                assert self._data_ready_at_path(
                    processed_val_data_path
                ), "Loading from file + setup() called but validation data not ready"
            else:
                processed_train_data_path = None
                processed_val_data_path = None

            self.train_ds = self._make_multitask_dataset("train", save_smiles_and_ids=save_smiles_and_ids)
            self.val_ds = self._make_multitask_dataset("val", save_smiles_and_ids=save_smiles_and_ids)
            logger.info(self.train_ds)
            logger.info(self.val_ds)
            labels_size.update(
                self.train_ds.labels_size
            )  # Make sure that all task label sizes are contained in here. Maybe do the update outside these if statements.
            labels_size.update(self.val_ds.labels_size)
            labels_dtype.update(self.train_ds.labels_dtype)
            labels_dtype.update(self.val_ds.labels_dtype)

        if stage == "test" or stage is None:
            if self.load_from_file:
                processed_test_data_path = self._path_to_load_from_file("test")
                assert self._data_ready_at_path(
                    processed_test_data_path
                ), "Loading from file + setup() called but test data not ready"
            else:
                processed_test_data_path = None
            self.test_ds = self._make_multitask_dataset("test", save_smiles_and_ids=save_smiles_and_ids)
            logger.info(self.test_ds)

            labels_size.update(self.test_ds.labels_size)
            labels_dtype.update(self.test_ds.labels_dtype)

        default_labels_size_dict = self.collate_fn.keywords.get("labels_size_dict", None)

        if default_labels_size_dict is None:
            self.collate_fn.keywords["labels_size_dict"] = labels_size

        default_labels_dtype_dict = self.collate_fn.keywords.get("labels_dtype_dict", None)

        if default_labels_dtype_dict is None:
            self.collate_fn.keywords["labels_dtype_dict"] = labels_dtype

    def _make_multitask_dataset(
        self,
        stage: Literal["train", "val", "test"],
        save_smiles_and_ids: bool,
        load_from_file: Optional[bool] = None,
    ) -> Datasets.MultitaskDataset:
        """
        Create a MultitaskDataset for the given stage using single task datasets
        The single task datasets must exist before this can be used

        Parameters:
            stage: Stage to create multitask dataset for
            save_smiles_and_ids: Whether to save SMILES strings and unique IDs
            data_path: path to load from if loading from file
            load_from_file: whether to load from file. If `None`, defers to `self.load_from_file`
        """

        allowed_stages = ["train", "val", "test"]
        assert stage in allowed_stages, f"Multitask dataset stage `{stage}` not in {allowed_stages}"

        if stage == "train":
            singletask_datasets = self.train_singletask_datasets
            about = "training set"
        elif stage == "val":
            singletask_datasets = self.val_singletask_datasets
            about = "validation set"
        elif stage == "test":
            singletask_datasets = self.test_singletask_datasets
            about = "test set"
        else:
            raise ValueError(f"Unknown stage {stage}")

        if load_from_file is None:
            load_from_file = self.load_from_file

        # assert singletask_datasets is not None, "Single task datasets must exist to make multitask dataset"
        if singletask_datasets is None:
            assert load_from_file
            assert self._data_ready_at_path(
                self._path_to_load_from_file(stage)
            ), "Trying to create multitask dataset without single-task datasets but data not ready"
            files_ready = True
        else:
            files_ready = False

        multitask_dataset = Datasets.MultitaskDataset(
            singletask_datasets,
            n_jobs=self.featurization_n_jobs,
            backend=self.featurization_backend,
            featurization_batch_size=self.featurization_batch_size,
            progress=self.featurization_progress,
            about=about,
            save_smiles_and_ids=save_smiles_and_ids,
            data_path=self._path_to_load_from_file(stage) if load_from_file else None,
            load_from_file=load_from_file,
            files_ready=files_ready,
        )  # type: ignore

        # calculate statistics for the train split and used for all splits normalization
        if stage == "train":
            self.get_label_statistics(
                self.processed_graph_data_path, self.data_hash, multitask_dataset, train=True
            )
        if not load_from_file:
            self.normalize_label(multitask_dataset, stage)

        return multitask_dataset

    def _ready_to_load_all_from_file(self) -> bool:
        """
        Check if the data for all stages is ready to be loaded from files
        """

        paths = [self._path_to_load_from_file(stage) for stage in ["train", "val", "test"]]
        ready = all(self._data_ready_at_path(path) for path in paths)

        return ready

    def _path_to_load_from_file(self, stage: Literal["train", "val", "test"]) -> Optional[str]:
        """
        Get path from which to load the data from files
        """
        if self.processed_graph_data_path is None:
            return None
        return osp.join(self.processed_graph_data_path, f"{stage}_{self.data_hash}")

    def _data_ready_at_path(self, path: str) -> bool:
        """
        Check if data can be loaded from this path
        """
        can_load_from_file = osp.exists(path) and self.get_folder_size(path) > 0

        return can_load_from_file

    def _save_data_to_files(self) -> None:
        """
        Save data to files so that they can be loaded from file during training/validation/test
        """

        stages = ["train", "val", "test"]

        # At the moment, we need to merge the `SingleTaskDataset`'s into `MultitaskDataset`s in order to save to file
        #     This is because the combined labels need to be stored together. We can investigate not doing this if this is a problem
        temp_datasets = {
            stage: self._make_multitask_dataset(stage, save_smiles_and_ids=False, load_from_file=False)
            for stage in stages
        }
        for stage in stages:
            self.save_featurized_data(temp_datasets[stage], self._path_to_load_from_file(stage))
            temp_datasets[stage].save_metadata(self._path_to_load_from_file(stage))
        # self.train_ds, self.val_ds, self.test_ds will be created during `setup()`
        del temp_datasets

    def get_folder_size(self, path):
        # check if the data items are actually saved into the folders
        return sum(os.path.getsize(osp.join(path, f)) for f in os.listdir(path))

    def calculate_statistics(self, dataset: Datasets.MultitaskDataset, train: bool = False):
        """
        Calculate the statistics of the labels for each task, and overwrites the `self.task_norms` attribute.

        Parameters:
            dataset: the dataset to calculate the statistics from
            train: whether the dataset is the training set

        """
        if self.task_norms and train:
            for task in dataset.labels_size.keys():
                # if the label type is graph_*, we need to stack them as the tensor shape is (num_labels, )
                if task.startswith("graph"):
                    labels = np.stack(
                        np.array([datum["labels"][task] for datum in dataset if task in datum["labels"]]),
                        axis=0,
                    )
                # for other tasks with node_ and edge_, the label shape is [num_nodes/num_edges, num_labels]
                # we can concatenate them directly
                else:
                    labels = np.concatenate(
                        [datum["labels"][task] for datum in dataset if task in datum["labels"]], axis=0
                    )

                self.task_norms[task].calculate_statistics(labels)

    def get_label_statistics(
        self,
        data_path: Union[str, os.PathLike],
        data_hash: str,
        dataset: Datasets.MultitaskDataset,
        train: bool = False,
    ):
        """
        Get the label statistics from the dataset, and save them to file, if needed.
        `self.task_norms` will be modified in-place with the label statistics.

        Parameters:
            data_path: the path to save and load the label statistics to. If None, no saving and loading will be done.
            data_hash: the hash of the dataset generated by `get_data_hash()`
            dataset: the dataset to calculate the statistics from
            train: whether the dataset is the training set

        """
        if data_path is None:
            self.calculate_statistics(dataset, train=train)
        else:
            path_with_hash = os.path.join(data_path, data_hash)
            os.makedirs(path_with_hash, exist_ok=True)
            filename = os.path.join(path_with_hash, "task_norms.pkl")
            if self.task_norms and train and not os.path.isfile(filename):
                self.calculate_statistics(dataset, train=train)
                torch.save(self.task_norms, filename, pickle_protocol=4)
            # if any of the above three condition does not satisfy, we load from file.
            else:
                self.task_norms = torch.load(filename)

    def normalize_label(self, dataset: Datasets.MultitaskDataset, stage) -> Datasets.MultitaskDataset:
        """
        Normalize the labels in the dataset using the statistics in `self.task_norms`.

        Parameters:
            dataset: the dataset to normalize the labels from

        Returns:
            the dataset with normalized labels
        """
        for task in dataset.labels_size.keys():
            # we normalize the dataset if (it is train split) or (it is val/test splits and normalize_val_test is set to true)
            if (stage == "train") or (stage in ["val", "test"] and self.task_norms[task].normalize_val_test):
                for i in range(len(dataset)):
                    if task in dataset[i]["labels"]:
                        dataset[i]["labels"][task] = self.task_norms[task].normalize(
                            dataset[i]["labels"][task]
                        )
        return dataset

    def save_featurized_data(self, dataset: Datasets.MultitaskDataset, processed_data_path):
        os.makedirs(processed_data_path)  # In case the len(dataset) is 0
        for i in range(0, len(dataset), 1000):
            os.makedirs(os.path.join(processed_data_path, format(i // 1000, "04d")), exist_ok=True)
        process_params = [(index, datum, processed_data_path) for index, datum in enumerate(dataset)]

        # Check if "about" is in the Dataset object
        about = ""
        if hasattr(dataset, "about"):
            about = dataset.about
        for param in tqdm(process_params, desc=f"Saving featurized data {about}"):
            self.process_func(param)
        return

    def process_func(self, param):
        index, datum, folder = param
        filename = os.path.join(folder, format(index // 1000, "04d"), format(index, "07d") + ".pkl")
        torch.save(
            {"graph_with_features": datum["features"], "labels": datum["labels"]},
            filename,
            pickle_protocol=4,
        )
        return

    def get_dataloader_kwargs(self, stage: RunningStage, shuffle: bool, **kwargs) -> Dict[str, Any]:
        """
        Get the options for the dataloader depending on the current stage.

        Parameters:
            stage: Whether in Training, Validating, Testing, Sanity-checking, Predicting, or Tuning phase.
            shuffle: set to ``True`` to have the data reshuffled at every epoch.

        Returns:
            Arguments to pass to the `DataLoader` during initialization
        """
        loader_kwargs = super().get_dataloader_kwargs(stage=stage, shuffle=shuffle, **kwargs)

        # Get batch size and IPU options for training set
        # if stage in [RunningStage.TRAINING, RunningStage.TUNING]:
        if stage in [RunningStage.TRAINING]:
            loader_kwargs["ipu_dataloader_options"] = self.ipu_dataloader_training_opts
            loader_kwargs["ipu_options"] = self.ipu_training_opts

        # Get batch size and IPU options for validation / testing sets
        elif stage in [RunningStage.VALIDATING, RunningStage.TESTING, RunningStage.PREDICTING]:
            loader_kwargs["ipu_dataloader_options"] = self.ipu_dataloader_inference_opts
            loader_kwargs["ipu_options"] = self.ipu_inference_opts
        else:
            raise ValueError(f"Wrong value for `stage`. Provided `{stage}`")

        # Remove the IPU options if not available
        if loader_kwargs["ipu_options"] is None:
            loader_kwargs.pop("ipu_options")
            if loader_kwargs["ipu_dataloader_options"] is not None:
                logger.warning(
                    "`ipu_dataloader_options` will be ignored since it is provided without `ipu_options`."
                )
            loader_kwargs.pop("ipu_dataloader_options")
        return loader_kwargs

    def get_dataloader(
        self, dataset: Dataset, shuffle: bool, stage: RunningStage
    ) -> Union[DataLoader, "poptorch.DataLoader"]:
        """
        Get the poptorch dataloader for a given dataset

        Parameters:
            dataset: The dataset from which to load the data
            shuffle: set to ``True`` to have the data reshuffled at every epoch.
            stage: Whether in Training, Validating, Testing, Sanity-checking, Predicting, or Tuning phase.

        Returns:
            The poptorch dataloader to sample from
        """
        kwargs = self.get_dataloader_kwargs(stage=stage, shuffle=shuffle)
        sampler = None
        # use sampler only when sampler_task_dict is set in the config and during training
        if DatasetSubSampler.check_sampling_required(self.sampler_task_dict) and stage in [
            RunningStage.TRAINING
        ]:
            sampler = DatasetSubSampler(
                dataset, self.sampler_task_dict, self.processed_graph_data_path, self.data_hash
            )
            # turn shuffle off when sampler is used as sampler option is mutually exclusive with shuffle
            kwargs["shuffle"] = False
        is_ipu = ("ipu_options" in kwargs.keys()) and (kwargs.get("ipu_options") is not None)
        if is_ipu:
            loader = IPUDataModuleModifier._dataloader(self, dataset=dataset, sampler=sampler, **kwargs)
        else:
            loader = BaseDataModule._dataloader(self, dataset=dataset, sampler=sampler, **kwargs)

        return loader

    def get_collate_fn(self, collate_fn):
        if collate_fn is None:
            # Some values become `inf` when changing data type. `mask_nan` deals with that
            collate_fn = partial(
                graphium_collate_fn,
                mask_nan=0,
                do_not_collate_keys=["smiles", "mol_ids"],
                batch_size_per_pack=self.batch_size_per_pack,
            )
            collate_fn.__name__ = graphium_collate_fn.__name__
        return collate_fn

    # Cannot be used as is for the multitask version, because sample_idx does not apply.
    def _featurize_molecules(self, smiles: Iterable[str]) -> Tuple[List, List]:
        """
        Precompute the features (graphs, fingerprints, etc.) from the SMILES.
        Features are computed from `self.smiles_transformer`.
        A warning is issued to mention which molecules failed featurization.

        Note:
            (hadim): in case of very large dataset we could:
            - or cache the data and read from it during `next(iter(dataloader))`
            - or compute the features on-the-fly during `next(iter(dataloader))`
            For now we compute in advance and hold everything in memory.

        Parameters:
            smiles: A list of all the molecular SMILES to featurize
            sample_idx: The indexes corresponding to the sampled SMILES.
                If not provided, computed from `numpy.arange`.

        Returns:
            features: A list of all the featurized molecules
            idx_none: A list of the indexes that failed featurization
        """

        batch_size = BatchingSmilesTransform.parse_batch_size(
            numel=len(smiles),
            desired_batch_size=self.featurization_batch_size,
            n_jobs=self.featurization_n_jobs,
        )

        # Loop all the smiles and compute the features
        features = dm.parallelized_with_batches(
            BatchingSmilesTransform(self.smiles_transformer),
            smiles,
            batch_size=batch_size,
            progress=True,
            n_jobs=self.featurization_n_jobs,
            backend=self.featurization_backend,
            tqdm_kwargs={"desc": f"featurizing_smiles, batch={batch_size}"},
        )

        # Warn about None molecules
        idx_none = [ii for ii, feat in enumerate(features) if did_featurization_fail(feat)]
        if len(idx_none) > 0:
            mols_to_msg = [
                f"idx={idx} - smiles={smiles[idx]} - Error_msg[:-200]=\n{str(features[idx])[:-200]}"
                for idx in idx_none
            ]
            msg = "\n".join(mols_to_msg)
            logger.warning(
                (f"{len(idx_none)} molecules will be removed since they failed featurization:\n" + msg)
            )

        return features, idx_none

    @staticmethod
    def _filter_none_molecules(
        idx_none: Iterable,
        *args: Union[pd.DataFrame, pd.Series, np.ndarray, torch.Tensor, list, tuple, Dict[Any, Iterable]],
    ) -> List[Union[pd.DataFrame, pd.Series, np.ndarray, torch.Tensor, list, tuple, Dict[Any, Iterable]]]:
        """
        Filter the molecules, labels, etc. for the molecules that failed featurization.

        Parameters:
            idx_none: A list of the indexes that failed featurization
            args: Any argument from which to filter the failed SMILES.
                Can be a `list`, `tuple`, `Tensor`, `np.array`, `Dict`, `pd.DataFrame`, `pd.Series`.
                Otherwise, it is not filtered.
                WARNING: If a `pd.DataFrame` or `pd.Series` is passed, it filters by the row indexes,
                NOT by the `DataFrame.index` or `Series.index`! Be careful!

        Returns:
            out: All the `args` with the indexes from `idx_none` removed.
        """
        if len(idx_none) == 0:
            return args
        idx_none = np.asarray(idx_none)

        out = []
        for arg in args:
            if isinstance(arg, pd.DataFrame):
                new = arg.drop(arg.index[idx_none], axis=0)
            elif isinstance(arg, pd.Series):
                new = arg.drop(arg.index[idx_none], axis=0)
            elif isinstance(arg, np.ndarray):
                new = np.delete(arg, idx_none, axis=0)
            elif isinstance(arg, torch.Tensor):
                not_none = torch.ones(arg.shape[0], dtype=bool)
                not_none[idx_none] = False
                new = arg[not_none]
            elif isinstance(arg, (list, tuple)):
                arg = list(arg)
                new = [elem for ii, elem in enumerate(arg) if ii not in idx_none]
            elif isinstance(arg, dict):
                new = {}
                for key, val in arg.items():
                    new[key] = MultitaskFromSmilesDataModule._filter_none_molecules(idx_none, val)  # Careful
            else:
                new = arg
            out.append(new)

        out = tuple(out) if len(out) > 1 else out[0]

        return out

    def _parse_label_cols(
        self,
        df: pd.DataFrame,
        df_path: Optional[Union[str, os.PathLike]],
        label_cols: Union[Type[None], str, List[str]],
        smiles_col: str,
    ) -> List[str]:
        r"""
        Parse the choice of label columns depending on the type of input.
        The input parameters `label_cols` and `smiles_col` are described in
        the `__init__` method.
        Parameters:
            df: The dataframe containing the labels.
            df_path: The path to the dataframe containing the labels.
            label_cols: The columns to use as labels.
            smiles_col: The column to use as SMILES
        Returns:
            the parsed label columns
        """
        if df is None:
            files = self._glob(df_path)
            if len(files) == 0:
                raise FileNotFoundError(f"No such file or directory `{df_path}`")

            cols = BaseDataModule._get_table_columns(files[0])
            for file in files[1:]:
                _cols = BaseDataModule._get_table_columns(file)
                if set(cols) != set(_cols):
                    raise RuntimeError(
                        f"Multiple data files have different columns. \nColumn set 1: {cols}\nColumn set 2: {_cols}"
                    )
        else:
            cols = list(df.columns)

        # A star `*` at the beginning or end of the string specifies to look for all
        # columns that starts/end with a specific string
        if isinstance(label_cols, str):
            if label_cols[0] == "*":
                label_cols = [col for col in cols if str(col).endswith(label_cols[1:])]
            elif label_cols[-1] == "*":
                label_cols = [col for col in cols if str(col).startswith(label_cols[:-1])]
            else:
                label_cols = [label_cols]

        elif label_cols is None:
            label_cols = [col for col in cols if col != smiles_col]

        return check_arg_iterator(label_cols, enforce_type=list)

    @property
    def is_prepared(self):
        if not hasattr(self, "dataset"):
            return False
        return getattr(self, "dataset") is not None

    @property
    def is_setup(self):
        if not (hasattr(self, "train_ds") or hasattr(self, "test_ds")):
            return False
        return (getattr(self, "train_ds") is not None) or (getattr(self, "test_ds") is not None)

    @property
    def num_node_feats(self):
        """Return the number of node features in the first graph"""
        graph = self.get_fake_graph()
        num_feats = graph.feat.shape[1]
        return num_feats

    @property
    def in_dims(self):
        """
        Return all input dimensions for the set of graphs.
        Including node/edge features, and
        raw positional encoding dimensions such eigval, eigvec, rwse and more
        """

        graph = self.get_fake_graph()
        if isinstance(graph, (GraphDict)):
            graph = graph.data

        # get list of all keys corresponding to positional encoding
        pe_dim_dict = {}
        g_keys = get_keys(graph)
        # ignore the normal keys for node feat and edge feat etc.
        for key in g_keys:
            prop = graph.get(key, None)
            if hasattr(prop, "shape"):
                pe_dim_dict[key] = prop.shape[-1]
        return pe_dim_dict

    @property
    def num_edge_feats(self):
        """Return the number of edge features in the first graph"""

        graph = self.get_fake_graph()
        empty = torch.Tensor([])
        num_feats = graph.get("edge_feat", empty).shape[-1]

        return num_feats

    def get_fake_graph(self):
        """
        Low memory footprint method to get the featurization of a fake graph
        without reading the dataset. Useful for getting the number of node/edge features.

        Returns:
            graph: A fake graph with the right featurization
        """

        smiles = "C1=CC=CC=C1"
        trans = deepcopy(self.smiles_transformer)
        trans.keywords.setdefault("on_error", "raise")
        trans.keywords.setdefault("mask_nan", 0.0)
        graph = trans(smiles)
        return graph

    ########################## Private methods ######################################
    def _save_to_cache(self):
        raise NotImplementedError()

    def _load_from_cache(self):
        raise NotImplementedError()

    def _extract_smiles_labels(
        self,
        df: pd.DataFrame,
        task_level: str,
        smiles_col: Optional[str] = None,
        label_cols: List[str] = [],
        idx_col: Optional[str] = None,
        mol_ids_col: Optional[str] = None,
        weights_col: Optional[str] = None,
        weights_type: Optional[str] = None,
    ) -> Tuple[
        np.ndarray, np.ndarray, Union[Type[None], np.ndarray], Dict[str, Union[Type[None], np.ndarray]]
    ]:
        """
        For a given dataframe extract the SMILES and labels columns. Smiles is returned as a list
        of string while labels are returned as a 2D numpy array.

        Parameters:
            df: Pandas dataframe
            smiles_col: Name of the column containing the SMILES
            label_cols: List of column names containing the labels
            idx_col: Name of the column containing the index
            mol_ids_col: Name of the column containing the molecule ids
            weights_col: Name of the column containing the weights
            weights_type: Type of weights to use.
        Returns:
            smiles, labels, sample_idx, extras
        """

        if smiles_col is None:  # Should we specify which dataset has caused the potential issue?
            smiles_col_all = [col for col in df.columns if "smile" in str(col).lower()]
            if len(smiles_col_all) == 0:
                raise ValueError(f"No SMILES column found in dataframe. Columns are {df.columns}")
            elif len(smiles_col_all) > 1:
                raise ValueError(
                    f"Multiple SMILES column found in dataframe. SMILES Columns are {smiles_col_all}"
                )

            smiles_col = smiles_col_all[0]

        if label_cols is None:
            label_cols = df.columns.drop(smiles_col)

        label_cols = check_arg_iterator(label_cols, enforce_type=list)
        smiles = df[smiles_col].values
        if len(label_cols) > 0:
            if task_level == "graph":
                labels = extract_labels(df, "graph", label_cols)
            elif task_level == "node":
                labels = extract_labels(df, "node", label_cols)
            elif task_level == "edge":
                labels = extract_labels(df, "edge", label_cols)
            elif task_level == "nodepair":
                labels = extract_labels(df, "nodepair", label_cols)
            else:
                raise ValueError(f"Unknown task level: {task_level}")
        else:
            labels = float("nan") + np.zeros([len(smiles), 0])

        # Get the indices, used for sub-sampling and splitting the dataset
        if idx_col is not None:
            df = df.set_index(idx_col)
        sample_idx = df.index.values

        # Get the molecule ids
        mol_ids = None
        if mol_ids_col is not None:
            mol_ids = df[mol_ids_col].values

        # Extract the weights
        weights = None
        if weights_col is not None:
            weights = df[weights_col].values
        elif weights_type is not None:
            if not np.all((labels == 0) | (labels == 1)):
                raise ValueError("Labels must be binary for `weights_type`")

            if weights_type == "sample_label_balanced":
                ratio_pos_neg = np.sum(labels, axis=0, keepdims=1) / labels.shape[0]
                weights = np.zeros(labels.shape)
                weights[labels == 0] = ratio_pos_neg
                weights[labels == 1] = ratio_pos_neg**-1

            elif weights_type == "sample_balanced":
                ratio_pos_neg = np.sum(labels, axis=0, keepdims=1) / labels.shape[0]
                weights = np.zeros(labels.shape)
                weights[labels == 0] = ratio_pos_neg
                weights[labels == 1] = ratio_pos_neg**-1
                weights = np.prod(weights, axis=1)

            else:
                raise ValueError(f"Undefined `weights_type` {weights_type}")

            weights /= np.max(weights)  # Put the max weight to 1

        extras = {"weights": weights, "mol_ids": mol_ids}
        return smiles, labels, sample_idx, extras

    def _get_split_indices(
        self,
        dataset_size: int,
        split_val: float,
        split_test: float,
        sample_idx: Optional[Iterable[int]] = None,
        split_seed: int = None,
        splits_path: Union[str, os.PathLike] = None,
        split_names: Optional[List[str]] = ["train", "val", "test"],
    ):
        r"""
        Compute indices of random splits.
        Parameters:
            dataset_size: Size of the dataset
            split_val: Fraction of the dataset to use for validation
            split_test: Fraction of the dataset to use for testing
            sample_idx: Indices of the samples to use for splitting
            split_seed: Seed for the random splitting
            splits_path: Path to a file containing the splits
        Returns:
            train_indices, val_indices, test_indices
        """

        if sample_idx is None:
            sample_idx = np.arange(dataset_size)

        if splits_path is None:
            # Random splitting
            if split_test + split_val > 0:
                train_indices, val_test_indices = train_test_split(
                    sample_idx,
                    test_size=split_val + split_test,
                    random_state=split_seed,
                )
                sub_split_test = split_test / (split_test + split_val)
            else:
                train_indices = sample_idx
                val_test_indices = np.array([])
                sub_split_test = 0

            if split_test > 0:
                val_indices, test_indices = train_test_split(
                    val_test_indices,
                    test_size=sub_split_test,
                    random_state=split_seed,
                )
            else:
                val_indices = val_test_indices
                test_indices = np.array([])

        else:
            # Split from an indices file
            file_type = self._get_data_file_type(splits_path)

            train, val, test = split_names

            if file_type == "pt":
                splits = torch.load(splits_path)
            elif file_type in ["csv", "tsv"]:
                with fsspec.open(str(splits_path)) as f:
                    splits = self._read_csv(splits_path)
            else:
                raise ValueError(
                    f"file type `{file_type}` for `{splits_path}` not recognised, please use .pt, .csv or .tsv"
                )
            train, val, test = split_names
            train_indices = np.asarray(splits[train]).astype("int")
            train_indices = train_indices[~np.isnan(train_indices)].tolist()
            val_indices = np.asarray(splits[val]).astype("int")
            val_indices = val_indices[~np.isnan(val_indices)].tolist()
            test_indices = np.asarray(splits[test]).astype("int")
            test_indices = test_indices[~np.isnan(test_indices)].tolist()

        # Filter train, val and test indices
        _, train_idx, _ = np.intersect1d(sample_idx, train_indices, return_indices=True)
        train_indices = train_idx.tolist()
        _, valid_idx, _ = np.intersect1d(sample_idx, val_indices, return_indices=True)
        val_indices = valid_idx.tolist()
        _, test_idx, _ = np.intersect1d(sample_idx, test_indices, return_indices=True)
        test_indices = test_idx.tolist()

        return train_indices, val_indices, test_indices

    def _sub_sample_df(
        self, df: pd.DataFrame, sample_size: Union[int, float, None], seed: Optional[int] = None
    ) -> pd.DataFrame:
        r"""
        subsample from a pandas dataframe
        Parameters:
            df: pandas dataframe to subsample
            sample_size: number of samples to subsample
        Returns:
            subsampled pandas dataframe
        """
        # Sub-sample the dataframe
        if isinstance(sample_size, int):
            n = min(sample_size, df.shape[0])
            df = df.sample(n=n, random_state=seed)
        elif isinstance(sample_size, float):
            df = df.sample(f=sample_size, random_state=seed)
        elif sample_size is None:
            pass
        else:
            raise ValueError(
                f"Wrong value for `sample_size`: {sample_size}"
            )  # Maybe specify which task it was for?

        return df

    def get_data_hash(self):
        """
        Get a hash specific to a dataset and smiles_transformer.
        Useful to cache the pre-processed data.
        """
        args = deepcopy(self.task_specific_args)
        # pop epoch_sampling_fraction out when creating hash
        # so that the data cache does not need to be regenerated
        # when epoch_sampling_fraction has changed.
        for task in self.task_specific_args.keys():
            try:
                if "epoch_sampling_fraction" in args[task].keys():
                    args[task].pop("epoch_sampling_fraction")
            except:
                pass
        hash_dict = {
            "smiles_transformer": self.smiles_transformer,
            "task_specific_args": args,
        }
        data_hash = get_md5_hash(hash_dict)
        return data_hash

    def get_data_cache_fullname(self, compress: bool = False) -> str:
        """
        Create a hash for the dataset, and use it to generate a file name

        Parameters:
            compress: Whether to compress the data
        Returns:
            full path to the data cache file
        """
        if self.cache_data_path is None:
            return
        ext = ".datacache"
        if compress:
            ext += ".gz"
        data_cache_fullname = fs.join(self.cache_data_path, self.data_hash + ext)
        return data_cache_fullname

    def save_data_to_cache(self, verbose: bool = True, compress: bool = False) -> None:
        """
        Save the datasets from cache. First create a hash for the dataset, use it to
        generate a file name. Then save to the path given by `self.cache_data_path`.

        Parameters:
            verbose: Whether to print the progress
            compress: Whether to compress the data

        """
        full_cache_data_path = self.get_data_cache_fullname(compress=compress)
        if full_cache_data_path is None:
            logger.info("No cache data path specified. Skipping saving the data to cache.")
            return

        save_params = {
            "single_task_datasets": self.single_task_datasets,
            "task_train_indices": self.task_train_indices,
            "task_val_indices": self.task_val_indices,
            "task_test_indices": self.task_test_indices,
        }

        fs.mkdir(self.cache_data_path)
        with fsspec.open(full_cache_data_path, mode="wb", compression="infer") as file:
            if verbose:
                logger.info(f"Saving the data to cache at path:\n`{full_cache_data_path}`")
            now = time.time()
            torch.save(save_params, file)
            elapsed = round(time.time() - now)
            if verbose:
                logger.info(
                    f"Successfully saved the data to cache in {elapsed}s at path: `{full_cache_data_path}`"
                )

        # At the moment, we need to merge the `SingleTaskDataset`'s into `MultitaskDataset`s in order to save label stats
        #     This is because the combined labels need to be stored together. We can investigate not doing this if this is a problem
        temp_train_dataset = self._make_multitask_dataset(
            stage="train", save_smiles_and_ids=False, load_from_file=False
        )

        self.get_label_statistics(self.cache_data_path, self.data_hash, temp_train_dataset, train=True)

    def load_data_from_cache(self, verbose: bool = True, compress: bool = False) -> bool:
        """
        Load the datasets from cache. First create a hash for the dataset, and verify if that
        hash is available at the path given by `self.cache_data_path`.

        Parameters:
            verbose: Whether to print the progress
            compress: Whether to compress the data

        Returns:
            cache_data_exists: Whether the cache exists (if the hash matches) and the loading succeeded
        """
        full_cache_data_path = self.get_data_cache_fullname(compress=compress)

        if full_cache_data_path is None:
            logger.info("No cache data path specified. Skipping loading the data from cache.")
            return False

        cache_data_exists = fs.exists(full_cache_data_path)

        if cache_data_exists:
            try:
                logger.info(f"Loading the data from cache at path `{full_cache_data_path}`")
                now = time.time()
                with fsspec.open(full_cache_data_path, mode="rb", compression="infer") as file:
                    load_params = torch.load(file)
                    self.__dict__.update(load_params)
                    (
                        self.train_singletask_datasets,
                        self.val_singletask_datasets,
                        self.test_singletask_datasets,
                    ) = self.get_subsets_of_datasets(
                        self.single_task_datasets,
                        self.task_train_indices,
                        self.task_val_indices,
                        self.task_test_indices,
                    )
                elapsed = round(time.time() - now)
                logger.info(
                    f"Successfully loaded the data from cache in {elapsed}s at path: `{full_cache_data_path}`"
                )
                return True
            except Exception as e:
                if verbose:
                    logger.warning(
                        f"Data cache failed to load path: `{full_cache_data_path}`.\nThe data will be prepared and cache will be created for future runs."
                    )
                    logger.warning(e.__str__())
                return False
        else:
            if verbose:
                logger.info(
                    f"Data cache not found at path: `{full_cache_data_path}`.\nThe data will be prepared and cache will be created for future runs."
                )
            return False

    def get_subsets_of_datasets(
        self,
        single_task_datasets: Dict[str, Datasets.SingleTaskDataset],
        task_train_indices: Dict[str, Iterable],
        task_val_indices: Dict[str, Iterable],
        task_test_indices: Dict[str, Iterable],
    ) -> Tuple[Subset, Subset, Subset]:
        """
        From a dictionary of datasets and their associated indices, subset the train/val/test sets

        Parameters:
            single_task_datasets: Dictionary of datasets
            task_train_indices: Dictionary of train indices
            task_val_indices: Dictionary of val indices
            task_test_indices: Dictionary of test indices
        Returns:
            train_singletask_datasets: Dictionary of train subsets
            val_singletask_datasets: Dictionary of val subsets
            test_singletask_datasets: Dictionary of test subsets
        """
        train_singletask_datasets = {}
        val_singletask_datasets = {}
        test_singletask_datasets = {}
        for task in task_train_indices.keys():
            train_singletask_datasets[task] = Subset(single_task_datasets[task], task_train_indices[task])
            val_singletask_datasets[task] = Subset(single_task_datasets[task], task_val_indices[task])
            test_singletask_datasets[task] = Subset(single_task_datasets[task], task_test_indices[task])
        return train_singletask_datasets, val_singletask_datasets, test_singletask_datasets

    def __len__(self) -> int:
        r"""
        Returns the number of elements of the current DataModule, which is the combined size of all single-task datasets given.
        Returns:
            num_elements: Number of elements in the current DataModule
        """
        num_elements = 0
        for task, args in self.task_dataset_processing_params.items():
            if args.df is None:
                df = self._read_table(args.df_path, usecols=[args.smiles_col])
                num_elements += len(df)
            else:
                num_elements += len(args.df)
        return num_elements

    def to_dict(self) -> Dict[str, Any]:
        """
        Returns a dictionary representation of the current DataModule
        Returns:
            obj_repr: Dictionary representation of the current DataModule
        """
        # TODO: Change to make more multi-task friendly
        obj_repr = {}
        obj_repr["name"] = self.__class__.__name__
        obj_repr["len"] = len(self)
        obj_repr["train_size"] = len(self.train_indices) if self.train_indices is not None else None
        obj_repr["val_size"] = len(self.val_indices) if self.val_indices is not None else None
        obj_repr["test_size"] = len(self.test_indices) if self.test_indices is not None else None
        obj_repr["batch_size_training"] = self.batch_size_training
        obj_repr["batch_size_inference"] = self.batch_size_inference
        obj_repr["batch_size_per_pack"] = self.batch_size_per_pack
        obj_repr["num_node_feats"] = self.num_node_feats
        obj_repr["num_node_feats_with_positional_encoding"] = self.num_node_feats_with_positional_encoding
        obj_repr["num_edge_feats"] = self.num_edge_feats
        obj_repr["num_tasks"] = len(self.task_dataset_processing_params)
        obj_repr["num_labels"] = len(self.label_cols)
        obj_repr["collate_fn"] = self.collate_fn.__name__
        obj_repr["featurization"] = self.featurization
        return obj_repr

    def __repr__(self) -> str:
        r"""
        Controls how the class is printed

        Returns:

        """
        return omegaconf.OmegaConf.to_yaml(self.to_dict())


class GraphOGBDataModule(MultitaskFromSmilesDataModule):
    def __init__(
        self,
        task_specific_args: Dict[str, Union[DatasetProcessingParams, Dict[str, Any]]],
        cache_data_path: Optional[Union[str, os.PathLike]] = None,
        featurization: Optional[Union[Dict[str, Any], omegaconf.DictConfig]] = None,
        batch_size_training: int = 16,
        batch_size_inference: int = 16,
        batch_size_per_pack: Optional[int] = None,
        num_workers: int = 0,
        pin_memory: bool = True,
        persistent_workers: bool = False,
        featurization_n_jobs: int = -1,
        featurization_progress: bool = False,
        featurization_backend: str = "loky",
        collate_fn: Optional[Callable] = None,
        prepare_dict_or_graph: str = "pyg:graph",
        **kwargs,
    ):
        r"""
        Load an OGB (Open-graph-benchmark) GraphProp dataset.

        Parameters:
            task_specific_args: Arguments related to each task, with the task-name being the key,
              and the specific arguments being the values. The arguments must be a
              Dict containing the following keys:

              - "dataset_name": Name of the OGB dataset to load. Examples of possible datasets are
                "ogbg-molhiv", "ogbg-molpcba", "ogbg-moltox21", "ogbg-molfreesolv".
              - "sample_size": The number of molecules to sample from the dataset. Default=None,
                meaning that all molecules will be considered.
            cache_data_path: path where to save or reload the cached data. The path can be
                remote (S3, GS, etc).
            featurization: args to apply to the SMILES to Graph featurizer.
            batch_size_training: batch size for training and val dataset.
            batch_size_inference: batch size for test dataset.
            num_workers: Number of workers for the dataloader. Use -1 to use all available
                cores.
            pin_memory: Whether to pin on paginated CPU memory for the dataloader.
            featurization_n_jobs: Number of cores to use for the featurization.
            featurization_progress: whether to show a progress bar during featurization.
            featurization_backend: The backend to use for the molecular featurization.

                - "multiprocessing": Found to cause less memory issues.
                - "loky": joblib's Default. Found to cause memory leaks.
                - "threading": Found to be slow.

            collate_fn: A custom torch collate function. Default is to `graphium.data.graphium_collate_fn`
            sample_size:

                - `int`: The maximum number of elements to take from the dataset.
                - `float`: Value between 0 and 1 representing the fraction of the dataset to consider
                - `None`: all elements are considered.
        """

        new_task_specific_args = {}
        self.metadata = {}
        for task_name, task_args in task_specific_args.items():
            # Get OGB metadata
            this_metadata = self._get_dataset_metadata(task_args["dataset_name"])
            # Get dataset
            df, mol_ids_col, smiles_col, label_cols, splits_path = self._load_dataset(
                this_metadata, sample_size=task_args.get("sample_size", None)
            )
            new_task_specific_args[task_name] = {
                "df": df,
                "mol_ids_col": mol_ids_col,
                "smiles_col": smiles_col,
                "label_cols": label_cols,
                "splits_path": splits_path,
                "task_level": task_args["task_level"],
            }
            self.metadata[task_name] = this_metadata

        # Config for datamodule
        dm_args = {}
        dm_args["task_specific_args"] = new_task_specific_args
        dm_args["cache_data_path"] = cache_data_path
        dm_args["featurization"] = featurization
        dm_args["batch_size_training"] = batch_size_training
        dm_args["batch_size_inference"] = batch_size_inference
        dm_args["batch_size_per_pack"] = batch_size_per_pack
        dm_args["num_workers"] = num_workers
        dm_args["pin_memory"] = pin_memory
        dm_args["featurization_n_jobs"] = featurization_n_jobs
        dm_args["featurization_progress"] = featurization_progress
        dm_args["featurization_backend"] = featurization_backend
        dm_args["persistent_workers"] = persistent_workers
        dm_args["collate_fn"] = collate_fn
        dm_args["prepare_dict_or_graph"] = prepare_dict_or_graph

        super().__init__(**dm_args, **kwargs)

    def to_dict(self) -> Dict[str, Any]:
        r"""
        geenrate a dictionary representation of the class
        Returns:
            dict: dictionary representation of the class
        """
        # TODO: Change to make more multi-task friendly
        obj_repr = {}
        obj_repr["dataset_name"] = self.dataset_name
        obj_repr.update(super().to_dict())
        return obj_repr

    # Private methods

    def _load_dataset(
        self,
        metadata: dict,
        sample_size: Optional[int] = None,
    ) -> Tuple[pd.DataFrame, str, str, List[str], str]:
        """
        Download, extract and load an OGB dataset.
        Parameters:
            metadata: Metadata for the dataset to load.
            sample_size: The number of molecules to sample from the dataset. Default=None,
                meaning that all molecules will be considered.
        Returns:
            df: Pandas dataframe containing the dataset.
            mol_ids_col: Name of the column containing the molecule ids.
            smiles_col: Name of the column containing the SMILES.
            label_cols: List of column names containing the labels.
            splits_path: Path to the file containing the train/val/test splits.
        """

        base_dir = fs.get_cache_dir("ogb")
        dataset_dir = base_dir / metadata["download_name"]

        if not dataset_dir.exists():
            # Create cache filepath for zip file and associated folder
            dataset_path = base_dir / f"{metadata['download_name']}.zip"

            # Download it
            if not dataset_path.exists():
                logger.info(f"Downloading {metadata['url']} to {dataset_path}")
                fs.copy(metadata["url"], dataset_path, progress=True)

            # Extract
            zf = zipfile.ZipFile(dataset_path)
            zf.extractall(base_dir)

        # Load CSV file
        if metadata["download_name"].startswith("pcqm4m"):
            df_path = dataset_dir / "raw" / "data.csv.gz"
        else:
            df_path = dataset_dir / "mapping" / "mol.csv.gz"
        logger.info(f"Loading {df_path} in memory.")
        df = pd.read_csv(df_path)

        # Subsample the dataset
        df = self._sub_sample_df(df, sample_size)

        # Load split from the OGB dataset and save them in a single CSV file
        if metadata["download_name"].startswith("pcqm4m"):
            split_name = metadata["split"]
            split_dict = torch.load(dataset_dir / "split_dict.pt")
            train_split = pd.DataFrame(split_dict["train"])
            val_split = pd.DataFrame(split_dict["valid"])
            if "test" in split_dict.keys():
                test_split = pd.DataFrame(split_dict["test"])
            else:
                test_split = pd.DataFrame(split_dict["test-dev"])

            splits = pd.concat([train_split, val_split, test_split], axis=1)  # type: ignore
            splits.columns = ["train", "val", "test"]

            splits_path = dataset_dir / "split"
            if not splits_path.exists():
                os.makedirs(splits_path)
                splits_path = dataset_dir / f"{split_name}.csv.gz"
            else:
                splits_path = splits_path / f"{split_name}.csv.gz"
        else:
            split_name = metadata["split"]
            train_split = pd.read_csv(dataset_dir / "split" / split_name / "train.csv.gz", header=None)  # type: ignore
            val_split = pd.read_csv(dataset_dir / "split" / split_name / "valid.csv.gz", header=None)  # type: ignore
            test_split = pd.read_csv(dataset_dir / "split" / split_name / "test.csv.gz", header=None)  # type: ignore

            splits = pd.concat([train_split, val_split, test_split], axis=1)  # type: ignore
            splits.columns = ["train", "val", "test"]

            splits_path = dataset_dir / "split" / f"{split_name}.csv.gz"

        logger.info(f"Saving splits to {splits_path}")
        splits.to_csv(splits_path, index=None)

        # Get column names: OGB columns are predictable
        if metadata["download_name"].startswith("pcqm4m"):
            mol_ids_col = df.columns[0]
            smiles_col = df.columns[-2]
            label_cols = df.columns[-1:].to_list()
        else:
            mol_ids_col = df.columns[-1]
            smiles_col = df.columns[-2]
            label_cols = df.columns[:-2].to_list()

        return df, mol_ids_col, smiles_col, label_cols, splits_path

    def _get_dataset_metadata(self, dataset_name: str) -> Dict[str, Any]:
        ogb_metadata = self._get_ogb_metadata()
        if dataset_name not in ogb_metadata.index:
            raise ValueError(f"'{dataset_name}' is not a valid dataset name.")

        return ogb_metadata.loc[dataset_name].to_dict()

    def _get_ogb_metadata(self):
        """
        Get the metadata of OGB GraphProp datasets.
        """

        with importlib.resources.open_text("ogb.graphproppred", "master.csv") as f:
            ogb_metadata = pd.read_csv(f)
        ogb_metadata = ogb_metadata.set_index(ogb_metadata.columns[0])
        ogb_metadata = ogb_metadata.T

        # Add metadata related to PCQM4M
        ogb_metadata = pd.concat([ogb_metadata, pd.DataFrame(PCQM4M_meta, index=["ogbg-lsc-pcqm4m"])])
        ogb_metadata = pd.concat([ogb_metadata, pd.DataFrame(PCQM4Mv2_meta, index=["ogbg-lsc-pcqm4mv2"])])

        # Only keep datasets of type 'mol'
        ogb_metadata = ogb_metadata[ogb_metadata["data type"] == "mol"]

        return ogb_metadata


class ADMETBenchmarkDataModule(MultitaskFromSmilesDataModule):
    """
    Wrapper to use the ADMET benchmark group from the TDC (Therapeutics Data Commons).

    !!! warning "Dependency"

        This class requires [PyTDC](https://pypi.org/project/PyTDC/) to be installed.

    !!! note "Citation"

        Huang, K., Fu, T., Gao, W., Zhao, Y., Roohani, Y., Leskovec, J., Coley, C., Xiao, C., Sun, J., & Zitnik, M. (2021).
        Therapeutics Data Commons: Machine Learning Datasets and Tasks for Drug Discovery and Development.
        Proceedings of Neural Information Processing Systems, NeurIPS Datasets and Benchmarks.


    Parameters:
        tdc_benchmark_names: This can be any subset of the benchmark names that make up the ADMET benchmarking group.
            If `None`, uses the complete benchmarking group. For all full list of options, see
            [the TDC website](https://tdcommons.ai/benchmark/admet_group/overview/) or use:

           ```python
           import tdc.utils.retrieve_benchmark_names
           retrieve_benchmark_names("admet_group")
           ```
        tdc_train_val_seed: TDC recommends a default splitting method for the train-val split. This parameter
          is used to seed that splitting method.
    """

    def __init__(
        self,
        # TDC-specific
        tdc_benchmark_names: Optional[Union[str, List[str]]] = None,
        tdc_train_val_seed: int = 0,
        # Inherited arguments from superclass
        cache_data_path: Optional[Union[str, os.PathLike]] = None,
        featurization: Optional[Union[Dict[str, Any], omegaconf.DictConfig]] = None,
        batch_size_training: int = 16,
        batch_size_inference: int = 16,
        batch_size_per_pack: Optional[int] = None,
        num_workers: int = 0,
        pin_memory: bool = True,
        persistent_workers: bool = False,
        featurization_n_jobs: int = -1,
        featurization_progress: bool = False,
        featurization_backend: str = "loky",
        collate_fn: Optional[Callable] = None,
        prepare_dict_or_graph: str = "pyg:graph",
        **kwargs,
    ):
        try:
            from tdc.benchmark_group import admet_group
            from tdc.utils import retrieve_benchmark_names
        except ImportError as error:
            # To make sure we use the exact same train-test set and preprocessing as other benchmark entries,
            # we rely on the PyTDC library.
            raise RuntimeError(
                f"To use {self.__class__.__name__}, `PyTDC` needs to be installed. "
                f"Please install it with `pip install PyTDC`"
            ) from error

        # Pick a path to save the TDC data to
        tdc_cache_dir = fs.get_cache_dir("tdc")
        tdc_cache_dir = fs.join(tdc_cache_dir, "ADMET_Benchmark")
        fs.mkdir(tdc_cache_dir, exist_ok=True)

        # Create the benchmark group object
        # NOTE (cwognum): We redirect stderr and stdout to a file since TDC uses print statements,
        #  which quickly pollute the logs. Ideally, we would use `redirect_stderr(None)`, but that breaks TQDM.
        with tempfile.TemporaryFile("w") as f:
            with redirect_stderr(f):
                with redirect_stdout(f):
                    self.group = admet_group(path=tdc_cache_dir)

        # By default, use all available benchmarks in a benchmark group
        if tdc_benchmark_names is None:
            tdc_benchmark_names = retrieve_benchmark_names("admet_group")
        if isinstance(tdc_benchmark_names, str):
            tdc_benchmark_names = [tdc_benchmark_names]

        # Create the task-specific arguments
        logger.info(
            f"Preparing the TDC ADMET Benchmark Group splits for each of the {len(tdc_benchmark_names)} benchmarks."
        )

        task_specific_args = {
            t: self._get_task_specific_arguments(t, tdc_train_val_seed, tdc_cache_dir)
            for t in tdc_benchmark_names
        }

        super().__init__(
            task_specific_args=task_specific_args,
            cache_data_path=cache_data_path,
            featurization=featurization,
            batch_size_training=batch_size_training,
            batch_size_inference=batch_size_inference,
            batch_size_per_pack=batch_size_per_pack,
            num_workers=num_workers,
            pin_memory=pin_memory,
            persistent_workers=persistent_workers,
            featurization_n_jobs=featurization_n_jobs,
            featurization_progress=featurization_progress,
            featurization_backend=featurization_backend,
            collate_fn=collate_fn,
            prepare_dict_or_graph=prepare_dict_or_graph,
            **kwargs,
        )

    def _get_task_specific_arguments(self, name: str, seed: int, cache_dir: str) -> DatasetProcessingParams:
        """
        Loads the data and split from the TDC benchmark group object.

        For the train-test split, this is fixed.

        For the train-val split, this does not have to be fixed. Here we use the default splitting method
        from TDC to do the train-val split and allow the seed to be changed. This is likely to best match
        other entries in the benchmarking group.
        """

        benchmark = self.group.get(name)

        # Get the default train-val-test split

        # NOTE (cwognum): TDC prints by default to stderr, which pollutes the logs quite a bit.
        #  This context manager mutes these by temporarily writing to a file.
        #  Ideally, we would use `redirect_stderr(None)`, but that breaks TQDM.

        with tempfile.TemporaryFile("w") as f:
            with redirect_stderr(f):
                train, val = self.group.get_train_valid_split(seed, name)
        test = benchmark["test"]

        # Convert to the Graphium format
        n_val = len(val)
        n_test = len(test)
        n_train = len(train)
        max_len = max(n_train, n_val, n_test)
        total_len = n_train + n_val + n_test

        data = pd.concat([train, val, test], ignore_index=True)

        # NOTE (cwognum): We need to convert the labels to float, since we use NaNs down the line.
        #  If you uncomment this line, collating the labels will raise an overflow by converting a NaN to the int dtype.
        data["Y"] = data["Y"].astype(float)

        split = pd.DataFrame(
            {
                "train": list(range(n_train)),
                "val": list(range(n_train, n_train + n_val)) + [float("nan")] * (max_len - n_val),
                "test": list(range(n_train + n_val, total_len)) + [float("nan")] * (max_len - n_test),
            }
        )
        split_path = fs.join(cache_dir, f"{name}_split.csv")
        split.to_csv(split_path, index=False)

        return DatasetProcessingParams(
            df=data,
            idx_col=None,
            smiles_col="Drug",
            label_cols=["Y"],
            splits_path=split_path,
            split_names=["train", "val", "test"],
            task_level="graph",
        )


class FakeDataModule(MultitaskFromSmilesDataModule):
    """
    A fake datamodule that generates artificial data by mimicking the true data coming
    from the provided dataset.
    It is useful to test the speed and performance of the model on a dataset without
    having to featurize it and wait for the workers to load it.
    """

    def __init__(
        self,
        task_specific_args: Dict[str, Dict[str, Any]],  # TODO: Replace this with DatasetParams
        cache_data_path: Optional[Union[str, os.PathLike]] = None,
        featurization: Optional[Union[Dict[str, Any], omegaconf.DictConfig]] = None,
        batch_size_training: int = 16,
        batch_size_inference: int = 16,
        num_workers: int = 0,
        pin_memory: bool = True,
        persistent_workers: bool = False,
        collate_fn: Optional[Callable] = None,
        prepare_dict_or_graph: str = "pyg:graph",
        num_mols_to_generate: int = 1000000,
        indexing_single_elem: bool = True,
        **kwargs,
    ):
        super().__init__(
            task_specific_args=task_specific_args,
            cache_data_path=cache_data_path,
            featurization=featurization,
            batch_size_training=batch_size_training,
            batch_size_inference=batch_size_inference,
            num_workers=num_workers,
            pin_memory=pin_memory,
            persistent_workers=persistent_workers,
            collate_fn=collate_fn,
            prepare_dict_or_graph=prepare_dict_or_graph,
            **kwargs,
        )
        self.num_mols_to_generate = num_mols_to_generate
        self.indexing_single_elem = indexing_single_elem

    def generate_data(self, label_cols: List[str], smiles_col: str):
        """
        Parameters:
            labels_cols
            smiles_col
        Returns:
            pd.DataFrame
        """
        num_generated_mols = int(1)
        # Create a dummy generated dataset - singel smiles string, duplicated N times
        example_molecules = dict(
            smiles="C1N2C3C4C5OC13C2C45",
            cxsmiles="[H]C1C2=C(NC(=O)[C@@]1([H])C1=C([H])C([H])=C(C([H])([H])[H])C([H])=C1[H])C([H])=C([H])N=C2[H] |(6.4528,-1.5789,-1.2859;5.789,-0.835,-0.8455;4.8499,-0.2104,-1.5946;3.9134,0.7241,-0.934;3.9796,1.1019,0.3172;5.0405,0.6404,1.1008;5.2985,1.1457,2.1772;5.9121,-0.5519,0.613;6.9467,-0.2303,0.8014;5.677,-1.7955,1.4745;4.7751,-2.7953,1.0929;4.2336,-2.7113,0.154;4.5521,-3.9001,1.914;3.8445,-4.6636,1.5979;5.215,-4.0391,3.1392;4.9919,-5.2514,4.0126;5.1819,-5.0262,5.0671;5.6619,-6.0746,3.7296;3.966,-5.6247,3.925;6.1051,-3.0257,3.52;6.6247,-3.101,4.4725;6.3372,-1.9217,2.7029;7.0168,-1.1395,3.0281;2.8586,1.2252,-1.7853;2.1303,1.9004,-1.3493;2.8118,0.8707,-3.0956;2.0282,1.2549,-3.7434;3.716,0.0207,-3.7371;4.6658,-0.476,-3.0127;5.3755,-1.1468,-3.5021)|",
        )
        example_df_entry = {smiles_col: example_molecules[smiles_col]}
        for label in label_cols:
            example_df_entry[label] = np.random.random()
        df = pd.DataFrame([example_df_entry])
        logger.info(f"Generating fake dataset on host... \n Generating {num_generated_mols} rows in the df.")
        df = pd.concat([df] * num_generated_mols, ignore_index=True)
        return df

    def prepare_data(self):
        """Called only from a single process in distributed settings. Steps:

        - If each cache is set and exists, reload from cache and return. Otherwise,
        - For each single-task dataset:
            - Load its dataframe from a path (if provided)
            - Subsample the dataframe
            - Extract the smiles, labels from the dataframe
        - In the previous step, we were also able to get the unique smiles, which we use to compute the features
        - For each single-task dataframe and associated data (smiles, labels, etc.):
            - Filter out the data corresponding to molecules which failed featurization.
            - Create a corresponding SingletaskDataset
            - Split the SingletaskDataset according to the task-specific splits for train, val and test
        """

        """Load all single-task dataframes."""
        if self.num_mols_to_generate is None:
            num_mols = 0

        task_df = {}
        for task, args in self.task_dataset_processing_params.items():
            logger.info(f"Reading data for task '{task}'")
            if args.df is None:
                # Only load the useful columns, as some datasets can be very large when loading all columns.
                label_cols = self._parse_label_cols(
                    df=None, df_path=args.df_path, label_cols=args.label_cols, smiles_col=args.smiles_col
                )
                task_df[task] = self.generate_data(label_cols=args.label_cols, smiles_col=args.smiles_col)
                if self.num_mols_to_generate is None:
                    num_mols = max(num_mols, len(task_df[task]))
            task_df[task] = task_df[task].iloc[0:1]

            args.label_cols = label_cols
        if self.num_mols_to_generate is None:
            self.num_mols_to_generate = num_mols
        logger.info("Done reading datasets")

        """Subsample the data frames and extract the necessary data to create SingleTaskDatasets for each task (smiles, labels, extras)."""
        task_dataset_args = {}
        for task in task_df.keys():
            task_dataset_args[task] = {}

        for task, df in task_df.items():
            logger.info(f"Prepare single-task dataset for task '{task}' with {len(df)} data points.")
            # Extract smiles, labels, extras
            args = self.task_dataset_processing_params[task]
            smiles, labels, sample_idx, extras = self._extract_smiles_labels(
                df,
                task_level=args.task_level,
                smiles_col=args.smiles_col,
                label_cols=args.label_cols,
                idx_col=args.idx_col,
                weights_col=args.weights_col,
                weights_type=args.weights_type,
            )

            # Store the relevant information for each task's dataset
            task_dataset_args[task]["smiles"] = smiles
            task_dataset_args[task]["labels"] = labels
            task_dataset_args[task]["sample_idx"] = sample_idx
            task_dataset_args[task]["extras"] = extras

        """Convert SMILES to features (graphs, fingerprints, etc.) for the unique molecules found."""
        all_smiles = []
        idx_per_task = {}
        total_len = 0
        for task, dataset_args in task_dataset_args.items():
            all_smiles.extend(dataset_args["smiles"])
            num_smiles = len(dataset_args["smiles"])
            idx_per_task[task] = (total_len, total_len + num_smiles)
            total_len += num_smiles
        # Get all unique mol ids
        all_unique_mol_ids = smiles_to_unique_mol_ids(
            all_smiles,
            n_jobs=self.featurization_n_jobs,
            featurization_batch_size=self.featurization_batch_size,
            backend=self.featurization_backend,
        )
        # Convert SMILES to features
        features, _ = self._featurize_molecules(all_smiles)
        task_dataset_args[task]["features"] = features
        """Filter data based on molecules which failed featurization. Create single task datasets as well."""
        self.single_task_datasets = {}
        for task, args in task_dataset_args.items():
            self.single_task_datasets[task] = Datasets.SingleTaskDataset(
                features=task_dataset_args[task]["features"],
                labels=task_dataset_args[task]["labels"],
                smiles=task_dataset_args[task]["smiles"],
                indices=task_dataset_args[task]["sample_idx"],
                unique_ids=all_unique_mol_ids[idx_per_task[task][0] : idx_per_task[task][1]],
                **task_dataset_args[task]["extras"],
            )

        """We split the data up to create train, val and test datasets"""
        self.train_singletask_datasets = {}
        self.val_singletask_datasets = {}
        self.test_singletask_datasets = {}
        for task, df in task_df.items():
            self.train_singletask_datasets[task] = Subset(self.single_task_datasets[task], [0])
            self.val_singletask_datasets[task] = Subset(self.single_task_datasets[task], [0])
            self.test_singletask_datasets[task] = Subset(self.single_task_datasets[task], [0])

    def setup(self, stage=None):
        # TODO
        """
        Prepare the torch dataset. Called on every GPUs. Setting state here is ok.
        Parameters:
            stage (str): Either 'fit', 'test', or None.
        """
        labels_size = {}

        if stage == "fit" or stage is None:
            self.train_ds = Datasets.FakeDataset(self.train_singletask_datasets, num_mols=self.num_mols_to_generate, indexing_same_elem=self.indexing_single_elem)  # type: ignore
            self.val_ds = Datasets.FakeDataset(self.val_singletask_datasets, num_mols=self.num_mols_to_generate, indexing_same_elem=self.indexing_single_elem)  # type: ignore
            print(self.train_ds)
            print(self.val_ds)

            labels_size.update(
                self.train_ds.labels_size
            )  # Make sure that all task label sizes are contained in here. Maybe do the update outside these if statements.
            labels_size.update(self.val_ds.labels_size)

        if stage == "test" or stage is None:
            self.test_ds = Datasets.FakeDataset(self.test_singletask_datasets, num_mols=self.num_mols_to_generate, indexing_same_elem=self.indexing_single_elem)  # type: ignore
            print(self.test_ds)
            labels_size.update(self.test_ds.labels_size)

        default_labels_size_dict = self.collate_fn.keywords.get("labels_size_dict", None)

        if default_labels_size_dict is None:
            self.collate_fn.keywords["labels_size_dict"] = labels_size

    def get_fake_graph(self):
        """
        Low memory footprint method to get the first datapoint DGL graph.
        The first 10 rows of the data are read in case the first one has a featurization
        error. If all 20 first element, then `None` is returned, otherwise the first
        graph to not fail is returned.
        """
        keys = list(self.task_dataset_processing_params.keys())
        task = keys[0]
        args = self.task_dataset_processing_params[task]
        if args.df is None:
            df = self._read_csv(args.df_path, nrows=20)
        else:
            df = args.df.iloc[0:20, :]

        df = df.iloc[0:20, :]
        label_cols = self._parse_label_cols(
            df, df_path=None, label_cols=args.label_cols, smiles_col=args.smiles_col
        )

        smiles, labels, sample_idx, extras = self._extract_smiles_labels(
            df,
            task_level=args.task_level,
            smiles_col=args.smiles_col,
            label_cols=label_cols,
            idx_col=args.idx_col,
            weights_col=args.weights_col,
            weights_type=args.weights_type,
        )

        graph = None
        for s in smiles:
            graph = self.smiles_transformer(s, mask_nan=0.0)
            num_nodes = graph.num_nodes
            num_edges = graph.num_edges
            if (graph is not None) and (num_edges > 0) and (num_nodes > 0):
                break
        return graph


# graphium/data/dataset.py

from typing import Type, List, Dict, Union, Any, Callable, Optional, Tuple, Iterable

from multiprocessing import Manager
import numpy as np
from functools import lru_cache
from loguru import logger
from copy import deepcopy
import os
import numpy as np

import torch
from torch.utils.data.dataloader import Dataset
from torch_geometric.data import Data, Batch

from graphium.data.smiles_transform import smiles_to_unique_mol_ids
from graphium.features import GraphDict


class SingleTaskDataset(Dataset):
    def __init__(
        self,
        labels: List[Union[torch.Tensor, np.ndarray]],
        features: Optional[List[Union[Data, "GraphDict"]]] = None,
        smiles: Optional[List[str]] = None,
        indices: Optional[List[int]] = None,
        weights: Optional[Union[torch.Tensor, np.ndarray]] = None,
        unique_ids: Optional[List[str]] = None,
        mol_ids: Optional[List[str]] = None,
    ):
        r"""
        dataset for a single task
        Parameters:
            labels: A list of labels for the given task (one per graph)
            features: A list of graphs
            smiles: A list of smiles
            indices: A list of indices
            weights: A list of weights
            unique_ids: A list of unique ids for each molecule generated from `datamol.unique_id`
            mol_ids: A list of ids coming from the original dataset. Useful to identify the molecule in the original dataset.
        """

        # Verify that all lists are the same length
        numel = len(labels)

        def _check_if_same_length(to_check, label):
            """Simple utility method to throw an error if the length is not as expected."""
            if to_check is not None and len(to_check) != numel:
                raise ValueError(
                    f"{label} must be the same length as `labels`, got {len(to_check)} and {numel}"
                )

        _check_if_same_length(features, "features")
        _check_if_same_length(indices, "indices")
        _check_if_same_length(weights, "weights")
        _check_if_same_length(unique_ids, "unique_ids")
        _check_if_same_length(mol_ids, "mol_ids")

        self.labels = labels
        if smiles is not None:
            manager = Manager()  # Avoid memory leaks with `num_workers > 0` by using the Manager
            self.smiles = manager.list(smiles)
        else:
            self.smiles = None
        self.features = features
        self.indices = indices
        if self.indices is not None:
            self.indices = np.array(
                self.indices
            )  # Avoid memory leaks with `num_workers > 0` by using numpy array
        self.weights = weights
        self.unique_ids = unique_ids
        self.mol_ids = mol_ids

    def __len__(self):
        r"""
        return the size of the dataset
        Returns:
            size: the size of the dataset
        """
        return len(self.labels)

    def __getitem__(self, idx):
        """
        get the data at the given index
        Parameters:
            idx: the index to get the data at
        Returns:
            datum: a dictionary containing the data at the given index, with keys "features", "labels", "smiles", "indices", "weights", "unique_ids"
        """
        datum = {}

        if self.features is not None:
            datum["features"] = self.features[idx]

        if self.labels is not None:
            datum["labels"] = self.labels[idx]

        if self.smiles is not None:
            datum["smiles"] = self.smiles[idx]

        if self.indices is not None:
            datum["indices"] = self.indices[idx]

        if self.weights is not None:
            datum["weights"] = self.weights[idx]

        if self.unique_ids is not None:
            datum["unique_ids"] = self.unique_ids[idx]

        if self.mol_ids is not None:
            datum["mol_ids"] = self.mol_ids[idx]

        return datum

    def __getstate__(self):
        """Serialize the class for pickling."""
        state = {}
        state["labels"] = self.labels
        state["smiles"] = list(self.smiles) if self.smiles is not None else None
        state["features"] = self.features
        state["indices"] = self.indices
        state["weights"] = self.weights
        state["unique_ids"] = self.unique_ids
        state["mol_ids"] = self.mol_ids
        return state

    def __setstate__(self, state: dict):
        """Reload the class from pickling."""
        if state["smiles"] is not None:
            manager = Manager()
            state["smiles"] = manager.list(state["smiles"])

        self.__dict__.update(state)


class MultitaskDataset(Dataset):
    pass

    def __init__(
        self,
        datasets: Dict[str, SingleTaskDataset],
        n_jobs=-1,
        backend: str = "loky",
        featurization_batch_size=1000,
        progress: bool = True,
        save_smiles_and_ids: bool = False,
        about: str = "",
        data_path: Optional[Union[str, os.PathLike]] = None,
        load_from_file: bool = False,
        files_ready: bool = False,
    ):
        r"""
        This class holds the information for the multitask dataset.
        Several single-task datasets can be merged to create a multi-task dataset. After merging the dictionary of single-task datasets.
        we will have a multitask dataset of the following form:
        - self.mol_ids will be a list to contain the unique molecular IDs to identify the molecules
        - self.smiles will be a list to contain the corresponding smiles for that molecular ID across all single-task datasets
        - self.labels will be a list of dictionaries where the key is the task name and the value is the label(s) for that task.
            At this point, any particular molecule will only have entries for tasks for which it has a label. Later, in the collate
            function, we fill up the missing task labels with NaNs.
        - self.features will be a list of featurized graphs corresponding to that particular unique molecule.
            However, for testing purposes we may not require features so that we can make sure that this merge function works.

        Parameters:
            datasets: A dictionary of single-task datasets
            n_jobs: Number of jobs to run in parallel
            backend: Parallelization backend
            featurization_batch_size: The batch size to use for the parallelization of the featurization
            progress: Whether to display the progress bar
            save_smiles_and_ids: Whether to save the smiles and ids for the dataset. If `False`, `mol_ids` and `smiles` are set to `None`
            about: A description of the dataset
            progress: Whether to display the progress bar
            about: A description of the dataset
            data_path: The location of the data if saved on disk
            load_from_file: Whether to load the data from disk
            files_ready: Whether the files to load from were prepared ahead of time
        """
        super().__init__()
        # self.datasets = datasets
        self.n_jobs = n_jobs
        self.backend = backend
        self.featurization_batch_size = featurization_batch_size
        self.progress = progress
        self.about = about
        self.data_path = data_path
        self.load_from_file = load_from_file

        if files_ready:
            assert load_from_file
            self._load_metadata()
            self.features = None
            self.labels = None

        else:
            task = next(iter(datasets))
            self.features = None
            if (len(datasets[task]) > 0) and ("features" in datasets[task][0]):
                self.mol_ids, self.smiles, self.labels, self.features = self.merge(datasets)
            else:
                self.mol_ids, self.smiles, self.labels = self.merge(datasets)
            # Set mol_ids and smiles to None to save memory as they are not needed.
            if not save_smiles_and_ids:
                self.mol_ids = None
                self.smiles = None
            self.labels_size = self.set_label_size_dict(datasets)
            self.labels_dtype = self.set_label_dtype_dict(datasets)
            self.dataset_length = len(self.labels)
            self._num_nodes_list = None
            self._num_edges_list = None
            if self.features is not None:
                self._num_nodes_list = get_num_nodes_per_graph(self.features)
                self._num_edges_list = get_num_edges_per_graph(self.features)
            if self.load_from_file:
                self.features = None
                self.labels = None

    def save_metadata(self, directory: str):
        """
        Save everything other than features/labels
        """
        attrs_to_save = [
            "mol_ids",
            "smiles",
            "labels_size",
            "labels_dtype",
            "dataset_length",
            "_num_nodes_list",
            "_num_edges_list",
        ]
        attrs = {attr: getattr(self, attr) for attr in attrs_to_save}

        path = os.path.join(directory, "multitask_metadata.pkl")

        torch.save(attrs, path, pickle_protocol=4)

    def _load_metadata(self):
        """
        Load everything other than features/labels
        """
        attrs_to_load = [
            "mol_ids",
            "smiles",
            "labels_size",
            "labels_dtype",
            "dataset_length",
            "_num_nodes_list",
            "_num_edges_list",
        ]
        path = os.path.join(self.data_path, "multitask_metadata.pkl")
        attrs = torch.load(path)

        if not set(attrs_to_load).issubset(set(attrs.keys())):
            raise ValueError(
                f"The metadata in the cache at {self.data_path} does not contain the right information. "
                f"This may be because the cache was prepared using an earlier version of Graphium. "
                f"You can try deleting the cache and running the data preparation again. "
                f"\nMetadata keys found: {attrs.keys()}"
                f"\nMetadata keys required: {attrs_to_load}"
            )

        for attr, value in attrs.items():
            setattr(self, attr, value)

    def __len__(self):
        r"""
        Returns the number of molecules
        """
        return self.dataset_length

    @property
    def num_nodes_list(self):
        """
        The number of nodes per graph
        """
        if self._num_nodes_list is None:
            if len(self) == 0:
                self._num_nodes_list = []
            else:
                self._num_nodes_list = get_num_nodes_per_graph(self.features)
        return self._num_nodes_list

    @property
    def num_edges_list(self):
        """
        The number of edges per graph
        """
        if self._num_edges_list is None:
            if len(self) == 0:
                self._num_edges_list = []
            else:
                self._num_edges_list = get_num_edges_per_graph(self.features)
        return self._num_edges_list

    @property
    def num_graphs_total(self):
        r"""
        number of graphs (molecules) in the dataset
        """
        return len(self)

    @property
    def num_nodes_total(self):
        """Total number of nodes for all graphs"""
        if len(self) == 0:
            return
        return sum(self.num_nodes_list)

    @property
    def max_num_nodes_per_graph(self):
        """Maximum number of nodes per graph"""
        if len(self) == 0:
            return
        return max(self.num_nodes_list)

    @property
    def std_num_nodes_per_graph(self):
        """Standard deviation of number of nodes per graph"""
        if len(self) == 0:
            return
        return np.std(self.num_nodes_list)

    @property
    def min_num_nodes_per_graph(self):
        """Minimum number of nodes per graph"""
        if len(self) == 0:
            return
        return min(self.num_nodes_list)

    @property
    def mean_num_nodes_per_graph(self):
        """Average number of nodes per graph"""
        if len(self) == 0:
            return
        return self.num_nodes_total / self.num_graphs_total

    @property
    def num_edges_total(self):
        """Total number of edges for all graphs"""
        if len(self) == 0:
            return
        return sum(self.num_edges_list)

    @property
    def max_num_edges_per_graph(self):
        """Maximum number of edges per graph"""
        if len(self) == 0:
            return
        return max(self.num_edges_list)

    @property
    def min_num_edges_per_graph(self):
        """Minimum number of edges per graph"""
        if len(self) == 0:
            return
        return min(self.num_edges_list)

    @property
    def std_num_edges_per_graph(self):
        """Standard deviation of number of nodes per graph"""
        if len(self) == 0:
            return
        return np.std(self.num_edges_list)

    @property
    def mean_num_edges_per_graph(self):
        """Average number of edges per graph"""
        if len(self) == 0:
            return
        return self.num_edges_total / self.num_graphs_total

    def __getitem__(self, idx):
        r"""
        get the data for at the specified index
        Parameters:
            idx: The index of the data to retrieve
        Returns:
            A dictionary containing the data for the specified index with keys "mol_ids", "smiles", "labels", and "features"
        """
        datum = {}
        if self.load_from_file:
            data_dict = self.load_graph_from_index(idx)
            datum["features"] = data_dict["graph_with_features"]
            datum["labels"] = data_dict["labels"]
            if "smiles" in data_dict.keys():
                datum["smiles"] = data_dict["smiles"]
        else:
            if self.mol_ids is not None:
                datum["mol_ids"] = self.mol_ids[idx]

            if self.smiles is not None:
                datum["smiles"] = self.smiles[idx]

            if self.labels is not None:
                datum["labels"] = self.labels[idx]

            if self.features is not None:
                datum["features"] = self.features[idx]

        return datum

    def load_graph_from_index(self, data_idx):
        r"""
        load the graph (in pickle file) from the disk
        Parameters:
            data_idx: The index of the data to retrieve
        Returns:
            A dictionary containing the data for the specified index with keys "graph_with_features", "labels" and "smiles" (optional).
        """
        filename = os.path.join(
            self.data_path, format(data_idx // 1000, "04d"), format(data_idx, "07d") + ".pkl"
        )
        data_dict = torch.load(filename)
        return data_dict

    def merge(
        self, datasets: Dict[str, SingleTaskDataset]
    ) -> Tuple[List[str], List[str], List[Dict[str, Any]], List[Any]]:
        r"""This function merges several single task datasets into a multitask dataset.

        The idea: for each of the smiles, labels, features and tasks, we create a corresponding list that concatenates these items across all tasks.
        In particular, for any index, the elements in the smiles, labels, features and task lists at that index will correspond to each other (i.e. match up).
        Over this list of all smiles (which we created by concatenating the smiles across all tasks), we compute their molecular ID using functions from Datamol.
        Once again, we will have a list of molecular IDs which is the same size as the list of smiles, labels, features and tasks.
        We then use numpy's `unique` function to find the exact list of unique molecular IDs as these will identify the molecules in our dataset. We also get the
        inverse from numpy's `unique`, which will allow us to index in addition to the list of all molecular IDs, the list of all smiles, labels, features and tasks.
        Finally, we use this inverse to construct the list of list of smiles, list of label dictionaries (indexed by task) and the list of features such that
        the indices match up. This is what is needed for the `get_item` function to work.

        Parameters:
            datasets: A dictionary of single-task datasets
        Returns:
            A tuple of (list of molecular IDs, list of smiles, list of label dictionaries, list of features)
        """

        # Get all the smiles, labels, features and tasks.
        all_lists = self._get_all_lists_ids(datasets=datasets)
        mol_ids, inv = self._get_inv_of_mol_ids(all_mol_ids=all_lists["mol_ids"])

        # Store the smiles.
        smiles = [[] for _ in range(len(mol_ids))]
        for all_idx, unique_idx in enumerate(inv):
            smiles[unique_idx].append(all_lists["smiles"][all_idx])

        # Store the labels.
        labels = [Data() for _ in range(len(mol_ids))]
        for all_idx, unique_idx in enumerate(inv):
            task: str = all_lists["tasks"][all_idx]
            label = all_lists["labels"][all_idx]
            labels[unique_idx][task] = label

            if all_idx < len(all_lists["features"]):
                features = all_lists["features"][all_idx]
                labels[unique_idx]["x"] = torch.empty(
                    (features.num_nodes, 1)
                )  # IPU is not happy with zero-sized tensors, so use shape (features.num_nodes, 1) here
                labels[unique_idx]["edge_index"] = torch.empty((2, features.num_edges))

        # Store the features
        if len(all_lists["features"]) > 0:
            features = [-1 for i in range(len(mol_ids))]
            for all_idx, unique_idx in enumerate(inv):
                features[unique_idx] = all_lists["features"][all_idx]
            return mol_ids, smiles, labels, features
        else:
            return mol_ids, smiles, labels

    def _get_all_lists_ids(self, datasets: Dict[str, SingleTaskDataset]) -> Dict[str, Any]:
        all_smiles = []
        all_features = []
        all_labels = []
        all_mol_ids = []
        all_tasks = []

        for task, ds in datasets.items():
            if len(ds) == 0:
                continue
            # Get data from single task dataset
            ds_smiles = [ds[i]["smiles"] for i in range(len(ds))]
            ds_labels = [ds[i]["labels"] for i in range(len(ds))]
            if "unique_ids" in ds[0].keys():
                ds_mol_ids = [ds[i]["unique_ids"] for i in range(len(ds))]
            else:
                ds_mol_ids = smiles_to_unique_mol_ids(
                    ds_smiles,
                    n_jobs=self.n_jobs,
                    featurization_batch_size=self.featurization_batch_size,
                    backend=self.backend,
                    progress=self.progress,
                    progress_desc=f"{task}: mol to ids",
                )
            if "features" in ds[0]:
                ds_features = [ds[i]["features"] for i in range(len(ds))]
            else:
                ds_features = None
            all_smiles.extend(ds_smiles)
            all_labels.extend(ds_labels)
            all_mol_ids.extend(ds_mol_ids)
            if ds_features is not None:
                all_features.extend(ds_features)

            task_list = [task] * ds.__len__()
            all_tasks.extend(task_list)

        all_lists = {
            "smiles": all_smiles,
            "features": all_features,
            "labels": all_labels,
            "mol_ids": all_mol_ids,
            "tasks": all_tasks,
        }

        return all_lists

    def _get_inv_of_mol_ids(self, all_mol_ids):
        mol_ids, inv = np.unique(all_mol_ids, return_inverse=True)
        return mol_ids, inv

    def _find_valid_label(self, task, ds):
        r"""
        For a given dataset, find a genuine label for that dataset
        """
        valid_label = None
        for i in range(len(ds)):
            if ds[i] is not None:
                valid_label = ds[i]["labels"]
                break

        if valid_label is None:
            raise ValueError(f"Dataset for task {task} has no valid labels.")

        return valid_label

    def set_label_size_dict(self, datasets: Dict[str, SingleTaskDataset]):
        r"""
        This gives the number of labels to predict for a given task.
        """
        task_labels_size = {}
        for task, ds in datasets.items():
            if len(ds) == 0:
                continue

            valid_label = self._find_valid_label(task, ds)

            # Assume for a fixed task, the label dimension is the same across data points
            torch_label = torch.as_tensor(valid_label)

            # First dimension is graph-specific
            task_labels_size[task] = torch_label.size()
        return task_labels_size

    def set_label_dtype_dict(self, datasets: Dict[str, SingleTaskDataset]):
        r"""
        Gets correct dtype for a given label
        """
        task_labels_dtype = {}
        for task, ds in datasets.items():
            if len(ds) == 0:
                continue

            valid_label = self._find_valid_label(task, ds)

            torch_label = torch.as_tensor(valid_label)
            task_labels_dtype[task] = torch_label.dtype
        return task_labels_dtype

    def __repr__(self) -> str:
        """
        summarizes the dataset in a string
        Returns:
            A string representation of the dataset.
        """
        if len(self) == 0:
            out_str = (
                f"-------------------\n{self.__class__.__name__}\n"
                + f"\tabout = {self.about}\n"
                + f"\tnum_graphs_total = {self.num_graphs_total}\n"
                + f"-------------------\n"
            )
            return out_str

        # Faster to compute the statistics if we unbatch first.
        features = self.features
        if isinstance(self.features, Batch):
            self.features = self.features.to_data_list()

        out_str = (
            f"-------------------\n{self.__class__.__name__}\n"
            + f"\tabout = {self.about}\n"
            + f"\tnum_graphs_total = {self.num_graphs_total}\n"
            + f"\tnum_nodes_total = {self.num_nodes_total}\n"
            + f"\tmax_num_nodes_per_graph = {self.max_num_nodes_per_graph}\n"
            + f"\tmin_num_nodes_per_graph = {self.min_num_nodes_per_graph}\n"
            + f"\tstd_num_nodes_per_graph = {self.std_num_nodes_per_graph}\n"
            + f"\tmean_num_nodes_per_graph = {self.mean_num_nodes_per_graph}\n"
            + f"\tnum_edges_total = {self.num_edges_total}\n"
            + f"\tmax_num_edges_per_graph = {self.max_num_edges_per_graph}\n"
            + f"\tmin_num_edges_per_graph = {self.min_num_edges_per_graph}\n"
            + f"\tstd_num_edges_per_graph = {self.std_num_edges_per_graph}\n"
            + f"\tmean_num_edges_per_graph = {self.mean_num_edges_per_graph}\n"
            + f"-------------------\n"
        )

        # Restore the original features.
        self.features = features

        return out_str


class FakeDataset(MultitaskDataset):
    """
    A dataset to hold the fake data.
    """

    def __init__(
        self, datasets: Dict[str, SingleTaskDataset], num_mols: int = 1234, indexing_same_elem: bool = False
    ):
        """
        Parameters:
            datasets:
                A dictionary of datasets. The keys are the task names and the values are the datasets.
            num_mols:
                The number of molecules to generate. In reality, it is the same molecule,
                but `num_mols` will change the length of the dataset.
            indexing_same_elem:
                If True, the same molecule is used for all samples.
                Otherwise, a deepcopied molecule is used for each sample.
        """
        self.indexing_same_elem = indexing_same_elem
        self.num_mols = num_mols
        self.num_datasets = len(datasets)

        self.about = "FakeDatasets"
        task = next(iter(datasets))
        if "features" in datasets[task][0]:
            self.mol_ids, self.smiles, self.labels, self.features = self.merge(datasets)
            if self.indexing_same_elem is False:
                self.mol_ids, self.smiles, self.labels, self.features = self.deepcopy_mol(
                    self.mol_ids, self.smiles, self.labels, self.features
                )
        else:
            self.mol_ids, self.smiles, self.labels = self.merge(datasets)
            if self.indexing_same_elem is False:
                self.mol_ids, self.smiles, self.labels, _ = self.deepcopy_mol(
                    self.mol_ids, self.smiles, self.labels
                )

        self.labels_size = self.set_label_size_dict(datasets)
        self.labels_dtype = self.set_label_dtype_dict(datasets)
        self.features = self.features

    def _get_inv_of_mol_ids(self, all_mol_ids):
        # The generated data is a single molecule duplicated
        mol_ids = np.array(all_mol_ids)
        inv = [_ for _ in range(len(mol_ids) // self.num_datasets)] * self.num_datasets
        mol_ids = np.unique(inv)
        return mol_ids, inv

    def deepcopy_mol(self, mol_ids, labels, smiles, features=None):
        """
        Create a deepcopy of the single molecule num_mols times

        Args:
            mol_ids (array): The single value for the mol ID
            labels (List[Dict]): List containing one dict with the label name-value pairs
            smiles (List[List[str]]): List of list containing SMILE sting
            features (List[Data], optional): list containing Data object. Defaults to None.

        Returns:
            The deep copy of the inputs
        """
        logger.info("Duplicating the single dataset element...")
        mol_ids = [deepcopy(mol_ids[0]) for _ in range(self.num_mols)]
        logger.info("Finished `mol_ids`")
        labels = [deepcopy(labels[0]) for _ in range(self.num_mols)]
        logger.info("Finished `labels`")
        smiles = [deepcopy(smiles[0]) for _ in range(self.num_mols)]
        logger.info("Finished `smiles`")
        if features is not None:
            features = [deepcopy(features[0]) for _ in range(self.num_mols)]
            logger.info("Finished `features`")
        return mol_ids, labels, smiles, features

    def __len__(self):
        r"""
        Returns the number of molecules
        """
        return self.num_mols

    def __getitem__(self, idx):
        r"""
        get the data for at the specified index
        Parameters:
            idx: The index of the data to retrieve
        Returns:
            A dictionary containing the data for the specified index with keys "mol_ids", "smiles", "labels", and "features"
        """
        datum = {}
        if self.indexing_same_elem is True:
            # If using a single memory location override the idx value passed
            idx = 0
        if self.labels is not None:
            datum["labels"] = self.labels[idx]

        if self.features is not None:
            datum["features"] = self.features[idx]

        return datum


def get_num_nodes_per_graph(graphs):
    r"""
    number of nodes per graph
    """
    if isinstance(graphs, Batch):
        graphs = graphs.to_data_list()
    counts = [graph.num_nodes for graph in graphs]
    return counts


def get_num_edges_per_graph(graphs):
    r"""
    number of edges per graph
    """
    if isinstance(graphs, Batch):
        graphs = graphs.to_data_list()
    counts = [graph.num_edges for graph in graphs]
    return counts


# graphium/data/multilevel_utils.py

import pandas as pd
import ast
import numpy as np
from typing import List
import itertools
import math


def extract_labels(df: pd.DataFrame, task_level: str, label_cols: List[str]):
    """Extracts labels in label_cols from dataframe df for a given task_level.
    Returns a list of numpy arrays converted to the correct shape. Multiple
    targets are concatenated for each graph.
    """

    def unpack(graph_data):
        graph_data = pd.to_numeric(graph_data, errors="coerce")
        if isinstance(graph_data, str):
            graph_data_list = ast.literal_eval(graph_data)
            return np.array(graph_data_list)
        elif isinstance(graph_data, (int, float)):
            return np.array([graph_data])
        elif isinstance(graph_data, list):
            return np.array(graph_data)
        elif isinstance(graph_data, np.ndarray):
            if len(graph_data.shape) == 0:
                graph_data = np.expand_dims(graph_data, 0)
            if graph_data.shape[0] == 0:
                graph_data = np.array([np.nan])
                # TODO: Warning
            return graph_data
        else:
            raise ValueError(
                f"Graph data should be one of str, float, int, list, np.ndarray, got {type(graph_data)}"
            )

    def unpack_column(data: pd.Series):
        return data.apply(unpack)

    def merge_columns(data: pd.Series):
        data = data.to_list()
        data = [np.array([np.nan]) if not isinstance(d, np.ndarray) and math.isnan(d) else d for d in data]
        padded_data = itertools.zip_longest(*data, fillvalue=np.nan)
        data = np.stack(list(padded_data), 1).T
        return data

    unpacked_df: pd.DataFrame = df[label_cols].apply(unpack_column)
    output = unpacked_df.apply(merge_columns, axis="columns").to_list()

    if task_level == "graph":
        return np.concatenate(output)
    return output


# graphium/data/normalization.py

from typing import Optional
from loguru import logger
import numpy as np
import torch
from torch import Tensor


class LabelNormalization:
    def __init__(
        self,
        method: Optional[str] = None,
        min_clipping: Optional[int] = None,
        max_clipping: Optional[int] = None,
        normalize_val_test: Optional[bool] = False,
        verbose: Optional[bool] = True,
    ):
        """
        Parameters:
        method: str
            Normalization method. Supports the following values:
            - `None` (default): No normalization applied
            - `normal`: Normalize to have 0-mean and 1-variance
            - `unit`: Normalize to have all values in the range 0-1

        min_clipping: int
            Minimum value to clip to. If `None` (default), no clipping is applied.
            For example, if `min_clipping` is -2, all values below -2 will be clipped to -2.
            This is applied before the normalization.

        max_clipping: int
            Maximum value to clip to. If `None` (default), no clipping is applied.
            For example, if `max_clipping` is 2, all values above 2 will be clipped to 2.
            This is applied before the normalization.
        """

        self.method = method
        self.min_clipping = min_clipping
        self.max_clipping = max_clipping
        self.normalize_val_test = normalize_val_test
        self.verbose = verbose
        self.data_max = None
        self.data_min = None
        self.data_mean = None
        self.data_std = None

    def calculate_statistics(self, array):
        """
        Saves the normalization parameters (e.g. mean and variance) to the object.
        """
        # add axis = 0 to make sure that the statistics for multiple column labels is a vector or list instead of a scalar
        self.data_max = np.nanmax(array, axis=0).tolist()
        self.data_min = np.nanmin(array, axis=0).tolist()
        self.data_mean = np.nanmean(array, axis=0).tolist()  # 5.380503871833475 for pcqm4mv2
        self.data_std = np.nanstd(array, axis=0).tolist()  # 1.17850688410978995 for pcqm4mv2
        if self.verbose:
            logger.info(f"Max value for normalization '{self.data_max}'")
            logger.info(f"Min value for normalization '{self.data_min}'")
            logger.info(f"Mean value for normalization '{self.data_mean}'")
            logger.info(f"STD value for normalization '{self.data_std}'")

    def normalize(self, input):
        """
        Apply the normalization method to the data.
        Saves the normalization parameters (e.g. mean and variance) to the object.
        """
        assert self.data_max is not None, "calculate_statistic must be called before applying normalization"
        if self.min_clipping is not None:
            self.data_min = max(self.min_clipping, self.data_min)
        if self.max_clipping is not None:
            self.data_max = min(self.max_clipping, self.data_max)
        clipping = self.min_clipping is not None and self.max_clipping is not None
        # Need to check since np.clip fails if both a_min and a_max are None
        if clipping:
            if isinstance(input, np.ndarray):
                input = np.clip(input, a_min=self.data_min, a_max=self.data_max)
            elif isinstance(input, Tensor):
                input = torch.clip(input, min=self.data_min, max=self.data_max)
        if self.method is None:
            return input
        elif self.method == "normal":
            return (input - self.data_mean) / self.data_std
        elif self.method == "unit":
            return (input - self.data_min) / (self.data_max - self.data_min)
        else:
            raise ValueError(f"normalization method {self.method} not recognised.")

    def denormalize(self, input):
        """
        Apply the inverse of the normalization method to the data.
        """
        if self.method is None:
            return input
        elif self.method == "normal":
            mean, std = torch.tensor(self.data_mean), torch.tensor(self.data_std)
            if input.device.type != "ipu":  # Cast to device if not on IPU
                mean, std = mean.to(input.device), std.to(input.device)
            return (input * std) + mean
        elif self.method == "unit":
            dmax, dmin = torch.tensor(self.data_max), torch.tensor(self.data_min)
            if input.device.type != "ipu":  # Cast to device if not on IPU
                dmax, dmin = dmax.to(input.device), dmin.to(input.device)
            return input * (dmax - dmin) + dmin
        else:
            raise ValueError(f"normalization method {self.method} not recognised.")


# graphium/data/sampler.py

from typing import Dict, Optional
from torch.utils.data.dataloader import Dataset

import torch.utils.data as data_utils
import numpy as np
import time
import os
import torch


class DatasetSubSampler(data_utils.Sampler):
    def __init__(
        self, dataset: Dataset, sampler_task_dict: Dict[str, Optional[float]], data_path: str, data_hash: str
    ):
        """
        Random sample a subset of the dataset for each task each epoch and combine them for training.

        Parameters:
            dataset: the whole training dataset
            sampler_task_dict: a dict which indicates the sampled amount of data for each task.
            data_path: path to save the data indices.
            data_hash: hash folder name for the data indices.
        """
        self.dataset = dataset
        self.sampler_task_dict = sampler_task_dict
        self.task_indices = {}
        now = time.time()
        path_with_hash = os.path.join(data_path, data_hash)
        os.makedirs(path_with_hash, exist_ok=True)
        filename = os.path.join(path_with_hash, "dataset_indices.pkl")
        # if file dataset_indices.pkl exists, load indices from file.
        if os.path.isfile(filename):
            print(f"Sampler--loading dataset indices from disk.")
            self.task_indices = torch.load(filename)
        # if not, iterate through all data items in dataset and save indices to disk
        else:
            for i in range(len(dataset)):
                # the dataset[i]["labels"].keys() are not deterministic, need to sort by key length
                task_name = sorted(dataset[i]["labels"].keys(), key=len)[-1]
                if task_name not in self.task_indices:
                    self.task_indices[task_name] = []
                self.task_indices[task_name].append(i)
            elapsed = round(time.time() - now)
            print(f"Sampler-->time spent on getting indices: {elapsed}.")
            torch.save(self.task_indices, filename, pickle_protocol=4)

    def __iter__(self):
        indices = []
        for task_name in self.task_indices.keys():
            task_size = int(len(self.task_indices[task_name]) * self.sampler_task_dict[task_name])
            indices += np.random.choice(self.task_indices[task_name], task_size, replace=False).tolist()
            indices_set = set(indices)
            self.total_size = len(indices_set)
        return iter(indices_set)

    def __len__(self):
        return self.total_size

    @classmethod
    def check_sampling_required(cls, sampler_task_dict):
        """
        Check if we need subsampling: if all items in the sampler_task_dict are 1.0,
        skip subsampling.
        """
        return not all(value == 1.0 for value in sampler_task_dict.values())


# graphium/data/sdf2csv.py

import datamol as dm
import zipfile
import pandas as pd
from rdkit import Chem
from ogb.lsc import PCQM4Mv2Dataset
import graphium
from graphium.data.datamodule import BaseDataModule
import csv


def extract_zip(fname):
    """
    #extract sdf from zip
    """
    zf = zipfile.ZipFile(fname)
    zf.extractall(".")


def extract_mols_from_sdf(fname):
    """
    load sdf into mols
    """
    mol_df = BaseDataModule._read_sdf(fname)
    mols = mol_df["_rdkit_molecule_obj"]
    return mols


def mols2cxs(mols):
    """
    convert into a smiles that contains the 3D structure
    """
    cxs = []
    for mol in mols:
        cxs.append(dm.to_smiles(mol, cxsmiles=True))
    return cxs


def write_csv(cxs: any, homos: any, fname: str):
    """
    write cxsmiles and homo lomo to file
    write the training molecules with cxsmiles first
    """
    outname = fname + ".csv"
    fieldnames = ["cxsmiles", "homo_lumo_gap"]

    with open(outname, "w") as outfile:
        writer = csv.DictWriter(outfile, fieldnames=fieldnames)
        writer.writeheader()
        for i in range(len(cxs)):
            writer.writerow({"cxsmiles": cxs[i], "homo_lumo_gap": homos[i]})


def sdf2csv(sdf_name: str = "pcqm4m-v2-train", outname: str = "pcqm4m-v2-train"):
    """
    function converting sdf file molecules into csv format using CXSmiles and combine with mols without 3d positions from ogb
    Parameters:
        sdf_name: name to the extracted sdf file
        outname: output name of the cxsmile file
    """
    mols = extract_mols_from_sdf(sdf_name + ".sdf")

    # download ogb smiles
    dataset = PCQM4Mv2Dataset(root=".", only_smiles=True)  # (smiles, homo_lomo)
    homos = []
    for i in range(len(mols)):
        homo = dataset[i][1]
        homos.append(homo)

    # write the trainning set molecules first with cxsmiles
    cxs = mols2cxs(mols)
    for j in range(len(mols), len(dataset)):
        cxs.append(dataset[j][0])
        homos.append(dataset[j][1])

    write_csv(cxs, homos, outname)


if __name__ == "__main__":
    """
    #* main function
    #! this script need to be located at the specific data folder as it uses relative dependencies
    for example   #* graphium/data/PCQM4Mv2

    instruction on how to generate the csv file:
    1. download the extract the sdf file from ogb: https://ogb.stanford.edu/docs/lsc/pcqm4mv2/
    $ wget http://ogb-data.stanford.edu/data/lsc/pcqm4m-v2-train.sdf.tar.gz
    $ md5sum pcqm4m-v2-train.sdf.tar.gz
    $ tar -xf pcqm4m-v2-train.sdf.tar.gz
    2. run this function (the smiles csv file will be directed downloaded in code)
    """
    sdf_name = "pcqm4m-v2-train"
    outname = "pcqm4m-v2-train"
    sdf2csv(sdf_name=sdf_name, outname=outname)

    #! check how many warning you get from loading cxsmiles
    # path = "pcqm4m-v2-train.csv"
    # df = pd.read_csv(path)
    # smiles = df["cxsmiles"]
    # print (smiles[0])
    # graphium.data.datamodule.smiles_to_unique_mol_ids(smiles)


# graphium/data/smiles_transform.py

from typing import Type, List, Dict, Union, Any, Callable, Optional, Tuple, Iterable

import os
import datamol as dm


def smiles_to_unique_mol_id(smiles: str) -> Optional[str]:
    """
    Convert a smiles to a unique MD5 Hash ID. Returns None if featurization fails.
    Parameters:
        smiles: A smiles string to be converted to a unique ID
    Returns:
        mol_id: a string unique ID
    """
    try:
        mol = dm.to_mol(mol=smiles)
        mol_id = dm.unique_id(mol)
    except:
        mol_id = ""
    if mol_id is None:
        mol_id = ""
    return mol_id


def did_featurization_fail(features: Any) -> bool:
    """
    Check if a featurization failed.
    """
    return (features is None) or isinstance(features, str)


class BatchingSmilesTransform:
    """
    Class to transform a list of smiles using a transform function
    """

    def __init__(self, transform: Callable):
        """
        Parameters:
            transform: Callable function to transform a single smiles
        """
        self.transform = transform

    def __call__(self, smiles_list: Iterable[str]) -> Any:
        """
        Function to transform a list of smiles
        """
        mol_id_list = []
        for smiles in smiles_list:
            mol_id_list.append(self.transform(smiles))
        return mol_id_list

    @staticmethod
    def parse_batch_size(numel: int, desired_batch_size: int, n_jobs: int) -> int:
        """
        Function to parse the batch size.
        The batch size is limited by the number of elements divided by the number of jobs.
        """
        assert ((n_jobs >= 0) or (n_jobs == -1)) and isinstance(
            n_jobs, int
        ), f"n_jobs must be a positive integer or -1, got {n_jobs}"
        assert (
            isinstance(desired_batch_size, int) and desired_batch_size >= 0
        ), f"desired_batch_size must be a positive integer, got {desired_batch_size}"

        if n_jobs == -1:
            n_jobs = os.cpu_count()
        if (n_jobs == 0) or (n_jobs == 1):
            batch_size = 1
        else:
            batch_size = min(desired_batch_size, numel // n_jobs)
        batch_size = max(1, batch_size)
        return batch_size


def smiles_to_unique_mol_ids(
    smiles: Iterable[str],
    n_jobs=-1,
    featurization_batch_size=1000,
    backend="loky",
    progress=True,
    progress_desc="mols to ids",
) -> List[Optional[str]]:
    """
    This function takes a list of smiles and finds the corresponding datamol unique_id
    in an element-wise fashion, returning the corresponding unique_ids.

    The ID is an MD5 hash of the non-standard InChiKey provided
    by `dm.to_inchikey_non_standard()`. It guarantees uniqueness for
    different tautomeric forms of the same molecule.

    Parameters:
        smiles: a list of smiles to be converted to mol ids
        n_jobs: number of jobs to run in parallel
        backend: Parallelization backend
        progress: Whether to display the progress bar

    Returns:
        ids: A list of MD5 hash ids
    """

    batch_size = BatchingSmilesTransform.parse_batch_size(
        numel=len(smiles), desired_batch_size=featurization_batch_size, n_jobs=n_jobs
    )

    unique_mol_ids = dm.parallelized_with_batches(
        BatchingSmilesTransform(smiles_to_unique_mol_id),
        smiles,
        batch_size=batch_size,
        progress=progress,
        n_jobs=n_jobs,
        backend=backend,
        tqdm_kwargs={"desc": f"{progress_desc}, batch={batch_size}"},
    )

    return unique_mol_ids


# graphium/data/utils.py

from typing import Union, List, Callable, Dict, Tuple, Any, Optional

import importlib.resources
import zipfile

from loguru import logger

import pandas as pd
import numpy as np

import graphium

from torch_geometric.data import Data
from graphium.features.featurizer import GraphDict

GRAPHIUM_DATASETS_BASE_URL = "gs://graphium-public/datasets"
GRAPHIUM_DATASETS = {
    "graphium-zinc-micro": "zinc-micro.zip",
    "graphium-zinc-bench-gnn": "zinc-bench-gnn.zip",
}


def load_micro_zinc() -> pd.DataFrame:
    """
    Return a dataframe of micro ZINC (1000 data points).
    Returns:
        pd.DataFrame: A dataframe of micro ZINC.
    """

    with importlib.resources.open_text("graphium.data.micro_ZINC", "micro_ZINC.csv") as f:
        df = pd.read_csv(f)

    return df  # type: ignore


def load_tiny_zinc() -> pd.DataFrame:
    """
    Return a dataframe of tiny ZINC (100 data points).
    Returns:
        pd.DataFrame: A dataframe of tiny ZINC.
    """

    with importlib.resources.open_text("graphium.data.micro_ZINC", "micro_ZINC.csv") as f:
        df = pd.read_csv(f, nrows=100)

    return df  # type: ignore


def graphium_package_path(graphium_path: str) -> str:
    """Return the path of a graphium file in the package."""

    assert graphium_path.startswith(
        "graphium://"
    ), f"Invalid graphium path, must start with 'graphium://': {graphium_path}"

    graphium_path = graphium_path.replace("graphium://", "")
    package, ressource = graphium_path.split("/")
    with importlib.resources.path(package, ressource) as data_path:
        pass
    return str(data_path)


def list_graphium_datasets() -> set:
    """
    List Graphium datasets available to download.
    Returns:
        set: A set of Graphium dataset names.
    """
    return set(GRAPHIUM_DATASETS.keys())


def download_graphium_dataset(
    name: str, output_path: str, extract_zip: bool = True, progress: bool = False
) -> str:
    r"""Download a Graphium dataset to a specified location.

    Args:
        name: Name of the Graphium dataset from `graphium.data.utils.get_graphium_datasets()`.
        output_path: Directory path where to download the dataset to.
        extract_zip: Whether to extract the dataset if it's a zip file.
        progress: Whether to show a progress bar during download.

    Returns:
        str: Path to the downloaded dataset.
    """

    if name not in GRAPHIUM_DATASETS:
        raise ValueError(f"'{name}' is not a valid Graphium dataset name. Choose from {GRAPHIUM_DATASETS}")

    fname = GRAPHIUM_DATASETS[name]

    dataset_path_source = graphium.utils.fs.join(GRAPHIUM_DATASETS_BASE_URL, fname)
    dataset_path_destination = graphium.utils.fs.join(output_path, fname)

    if not graphium.utils.fs.exists(dataset_path_destination):
        graphium.utils.fs.copy(dataset_path_source, dataset_path_destination, progress=progress)

        if extract_zip and str(dataset_path_destination).endswith(".zip"):
            # Unzip the dataset
            with zipfile.ZipFile(dataset_path_destination, "r") as zf:
                zf.extractall(output_path)

    if extract_zip:
        # Set the destination path to the folder
        # NOTE(hadim): this is a bit fragile.
        dataset_path_destination = dataset_path_destination.split(".")[0]

    return dataset_path_destination


def get_keys(pyg_data):
    if isinstance(type(pyg_data).keys, property):
        return pyg_data.keys
    else:
        return pyg_data.keys()


def found_size_mismatch(task: str, features: Union[Data, GraphDict], labels: np.ndarray, smiles: str) -> bool:
    """Check if a size mismatch exists between features and labels with respect to node/edge/nodepair.

    Args:
        task: The task name is needed to determine the task level (graph, node, edge or nodepair)
        features: Features/information of molecule/graph (e.g., edge_index, feat, edge_feat, num_nodes, etc.)
        labels: Target label of molecule for the task
        smiles: Smiles string of molecule

    Returns:
        mismatch: Boolean variable indicating if a size mismatch was found between featurs and labels.
    """

    mismatch = False

    if np.isnan(labels).any():
        pass

    elif task.startswith("graph_"):
        pass

    elif task.startswith("node_"):
        if labels.shape[0] != features.num_nodes:
            mismatch = True
            logger.warning(
                (
                    f"Inconsistent number of nodes between labels and features in {task} task for {smiles}: {labels.shape[0]} vs {features.num_nodes}"
                )
            )

    elif task.startswith("edge_"):
        if labels.shape[0] != features.num_edges:
            mismatch = True
            logger.warning(
                (
                    f"Inconsistent number of edges between labels and features in {task} task for {smiles}: {labels.shape[0]} vs {features.num_edges}"
                )
            )

    elif task.startswith("nodepair_"):
        if list(labels.shape[:2]) != 2 * [features.num_nodes]:
            mismatch = True
            logger.warning(
                (
                    f"Inconsistent shape of nodepairs between labels and features in {task} task for {smiles}: {list(labels.shape[:2])} vs {2 * [features.num_nodes]}"
                )
            )

    else:
        raise ValueError("Unkown task level")

    return mismatch


# graphium/features/__init__.py

from .featurizer import get_mol_atomic_features_onehot
from .featurizer import get_mol_atomic_features_float
from .featurizer import get_mol_edge_features
from .featurizer import mol_to_adj_and_features
from .featurizer import mol_to_graph_dict
from .featurizer import mol_to_graph_signature
from .featurizer import GraphDict
from .featurizer import mol_to_pyggraph
from .featurizer import to_dense_array


# graphium/features/commute.py

from typing import Tuple, Union, Dict, Any

import numpy as np

from scipy.sparse import spmatrix, issparse
from scipy.linalg import pinv


def compute_commute_distances(
    adj: Union[np.ndarray, spmatrix], num_nodes: int, cache: Dict[str, Any]
) -> Tuple[np.ndarray, str, Dict[str, Any]]:
    """
    Compute avg. commute time/distance between nodepairs. This is the avg. number of steps a random walker, starting
    at node i, will take before reaching a given node j for the first time, and then return to node i.

    Reference: Saerens et al. "The principal components analysis of a graph, and its relationships to spectral clustering." ECML. 2004.

    Parameters:
        adj [num_nodes, num_nodes]: Adjacency matrix
        num_nodes: Number of nodes in the graph
        cache: Dictionary of cached objects
    Returns:
        dist [num_nodes, num_nodes]: 2D array with avg. commute distances between nodepairs
        base_level: Indicator of the output pos_level (node, edge, nodepair, graph) -> here nodepair
        cache: Updated dictionary of cached objects
    """

    base_level = "nodepair"

    if "commute" in cache:
        dist = cache["commute"]

    else:
        if issparse(adj):
            adj = adj.toarray()

        volG = adj.sum()

        if "pinvL" in cache:
            pinvL = cache["pinvL"]

        else:
            L = np.diagflat(np.sum(adj, axis=1)) - adj
            pinvL = pinv(L)
            cache["pinvL"] = pinvL

        dist = volG * np.asarray(
            [
                [pinvL[i, i] + pinvL[j, j] - 2 * pinvL[i, j] for j in range(num_nodes)]
                for i in range(num_nodes)
            ]
        )
        cache["commute"] = dist

    return dist, base_level, cache


# graphium/features/electrostatic.py

from typing import Tuple, Union, Dict, Any

import numpy as np

from scipy.linalg import pinv
from scipy.sparse import spmatrix, issparse


def compute_electrostatic_interactions(
    adj: Union[np.ndarray, spmatrix], cache: Dict[str, Any]
) -> Tuple[np.ndarray, str, Dict[str, Any]]:
    """
    Compute electrostatic interaction of nodepairs.

    Parameters:
        adj [num_nodes, num_nodes]: Adjacency matrix
        cache: Dictionary of cached objects
    Returns:
        electrostatic [num_nodes, num_nodes]: 2D array with electrostatic interactions of node nodepairs
        base_level: Indicator of the output pos_level (node, edge, nodepair, graph) -> here nodepair
        cache: Updated dictionary of cached objects
    """

    base_level = "nodepair"

    if "electrostatic" in cache:
        electrostatic = cache["electrostatic"]

    else:
        if "pinvL" in cache:
            pinvL = cache["pinvL"]

        else:
            if issparse(adj):
                adj = adj.toarray()

            L = np.diagflat(np.sum(adj, axis=1)) - adj
            pinvL = pinv(L)
            cache["pinvL"] = pinvL

        electrostatic = pinvL - np.diag(pinvL)  # This means that the "ground" is set to any given atom
        cache["electrostatic"] = electrostatic

    return electrostatic, base_level, cache


# graphium/features/featurizer.py

from typing import Union, List, Callable, Dict, Tuple, Any, Optional

import inspect
from loguru import logger
import numpy as np
from scipy.sparse import issparse, coo_matrix
import torch
from torch import Tensor

from torch_geometric.data import Data

from rdkit import Chem
import datamol as dm

from graphium.features import nmp
from graphium.utils.tensor import one_of_k_encoding
from graphium.features.positional_encoding import get_all_positional_encodings


def to_dense_array(array: np.ndarray, dtype: str = None) -> np.ndarray:
    r"""
    Assign the node data
    Parameters:
        array: The array to convert to dense
        dtype: The dtype of the array
    Returns:
        The dense array
    """
    if array is not None:
        if issparse(array):
            if array.dtype == np.float16:  # float16 doesn't support `todense`
                array = array.astype(np.float32)
            array = array.todense()

        if dtype is not None:
            array = array.astype(dtype)
    return array


def to_dense_tensor(tensor: Tensor, dtype: str = None) -> Tensor:
    r"""
    Assign the node data
    Parameters:
        array: The array to convert to dense
        dtype: The dtype of the array
    Returns:
        The dense array
    """
    if tensor is not None:
        if tensor.is_sparse:
            tensor = tensor.todense()
        if dtype is not None:
            tensor = tensor.to(dtype)
    return tensor


def _mask_nans_inf(mask_nan: Optional[str], array: np.ndarray, array_name: str) -> np.ndarray:
    r"""
    mask the NaNs in the array
    Parameters:
        mask_nan: How to mask the NaNs
        array: The array to mask
        array_name: The name of the array
    Returns:
        The masked array
    """
    if (mask_nan is None) or (array is None):
        return array

    new_array = array
    if issparse(new_array):
        new_array = new_array.data
    nans = ~np.isfinite(new_array)

    # Mask the NaNs
    if nans.any():
        msg = f"There are {np.sum(nans)} NaNs in `{array_name}`"
        if mask_nan == "raise":
            raise ValueError(msg)
        elif mask_nan == "warn":
            logger.warning(msg)
        else:
            new_array[nans] = mask_nan
            if issparse(array):
                array.data = new_array
                new_array = array
    return new_array


def get_mol_atomic_features_onehot(mol: dm.Mol, property_list: List[str]) -> Dict[str, Tensor]:
    r"""
    Get the following set of features for any given atom

    * One-hot representation of the atom
    * One-hot representation of the atom degree
    * One-hot representation of the atom implicit valence
    * One-hot representation of the the atom hybridization
    * Whether the atom is aromatic
    * The atom's formal charge
    * The atom's number of radical electrons

    Additionally, the following features can be set, depending on the value of input Parameters

    * One-hot representation of the number of hydrogen atom in the the current atom neighborhood if `explicit_H` is false
    * One-hot encoding of the atom chirality, and whether such configuration is even possible

    Parameters:

        mol:
            molecule from which to extract the properties

        property_list:
            A list of integer atomic properties to get from the molecule.
            The integer values are converted to a one-hot vector.
            Callables are not supported by this function.

            Accepted properties are:

            - "atomic-number"
            - "degree"
            - "valence", "total-valence"
            - "implicit-valence"
            - "hybridization"
            - "chirality"
            - "phase"
            - "type"
            - "group"
            - "period"

    Returns:
        prop_dict:
            A dictionnary where the element of ``property_list`` are the keys
            and the values are np.ndarray of shape (N, OH). N is the number of atoms
            in ``mol`` and OH the lenght of the one-hot encoding.

    """

    prop_dict = {}

    for prop in property_list:
        prop = prop.lower()
        prop_name = prop

        property_array = []
        for ii, atom in enumerate(mol.GetAtoms()):
            if prop in ["atomic-number"]:
                one_hot = one_of_k_encoding(atom.GetSymbol(), nmp.ATOM_LIST)
            elif prop in ["degree"]:
                one_hot = one_of_k_encoding(atom.GetDegree(), nmp.ATOM_DEGREE_LIST)
            elif prop in ["valence", "total-valence"]:
                prop_name = "valence"
                one_hot = one_of_k_encoding(atom.GetTotalValence(), nmp.VALENCE)
            elif prop in ["implicit-valence"]:
                one_hot = one_of_k_encoding(atom.GetImplicitValence(), nmp.VALENCE)
            elif prop in ["hybridization"]:
                one_hot = one_of_k_encoding(atom.GetHybridization(), nmp.HYBRIDIZATION_LIST)
            elif prop in ["chirality"]:
                try:
                    one_hot = one_of_k_encoding(atom.GetProp("_CIPCode"), nmp.CHIRALITY_LIST)
                    one_hot.append(int(atom.HasProp("_ChiralityPossible")))
                except:
                    one_hot = [0, 0, int(atom.HasProp("_ChiralityPossible"))]
            elif prop in "phase":
                one_hot = one_of_k_encoding(nmp.PHASE[atom.GetAtomicNum() - 1], nmp.PHASE_SET)
            elif prop in "type":
                one_hot = one_of_k_encoding(nmp.TYPE[atom.GetAtomicNum() - 1], nmp.TYPE_SET)
            elif prop in "group":
                one_hot = one_of_k_encoding(nmp.GROUP[atom.GetAtomicNum() - 1], nmp.GROUP_SET)
            elif prop in "period":
                one_hot = one_of_k_encoding(nmp.PERIOD[atom.GetAtomicNum() - 1], nmp.PERIOD_SET)
            else:
                raise ValueError(f"Unsupported property `{prop}`")

            property_array.append(np.asarray(one_hot, dtype=np.float16))

        prop_dict[prop_name] = np.stack(property_array, axis=0)

    return prop_dict


def get_mol_conformer_features(
    mol: dm.Mol,
    property_list: Union[List[str], List[Callable]],
    mask_nan: Optional[Union[float, str]] = None,
) -> Dict[str, np.ndarray]:
    r"""obtain the conformer features of a molecule
    Parameters:

        mol:
            molecule from which to extract the properties

        property_list:
            A list of conformer property to get from the molecule
            Accepted properties are:
            - "positions_3d"

    Returns:
        prop_dict: a dictionary where the element of ``property_list`` are the keys
    """
    prop_dict = {}
    has_conf = True

    try:
        mol.GetConformer()
    except:
        has_conf = False
    # * currently only accepts "positions_3d", raise errors otherwise
    for prop in property_list:
        if isinstance(prop, str):
            if prop in ["positions_3d"]:  # locating 3d conformer coordinates
                if not has_conf:
                    positions = np.full((mol.GetNumAtoms(), 3), float("nan"), dtype=np.float16)
                else:
                    positions = [[], [], []]
                    for i in range(mol.GetNumAtoms()):
                        pos = mol.GetConformer().GetAtomPosition(i)
                        positions[0].append(pos.x)
                        positions[1].append(pos.y)
                        positions[2].append(pos.z)
                    positions = np.asarray(positions, dtype=np.float16).T
                prop_dict[prop] = positions
            else:
                raise ValueError(
                    str(prop) + " is not currently supported as a conformer property in `property_list`"
                )
        else:
            raise ValueError(f"Elements in `property_list` must be str or callable, provided `{type(prop)}`")

        prop_dict[prop] = _mask_nans_inf(mask_nan, prop_dict[prop], prop)

    return prop_dict


def get_mol_atomic_features_float(
    mol: dm.Mol,
    property_list: Union[List[str], List[Callable]],
    offset_carbon: bool = True,
    mask_nan: Union[str, float, type(None)] = "raise",
) -> Dict[str, np.ndarray]:
    r"""
    Get a dictionary of floating-point arrays of atomic properties.
    To ensure all properties are at a similar scale, some of the properties
    are divided by a constant.

    There is also the possibility of offseting by the carbon value using
    the `offset_carbon` parameter.

    Parameters:

        mol:
            molecule from which to extract the properties

        property_list:
            A list of atomic properties to get from the molecule, such as 'atomic-number',
            'mass', 'valence', 'degree', 'electronegativity'.
            Some elements are divided by a factor to avoid feature explosion.

            Accepted properties are:

            - "atomic-number"
            - "mass", "weight"
            - "valence", "total-valence"
            - "implicit-valence"
            - "hybridization"
            - "chirality"
            - "hybridization"
            - "aromatic"
            - "ring", "in-ring"
            - "min-ring"
            - "max-ring"
            - "num-ring"
            - "degree"
            - "radical-electron"
            - "formal-charge"
            - "vdw-radius"
            - "covalent-radius"
            - "electronegativity"
            - "ionization", "first-ionization"
            - "melting-point"
            - "metal"
            - "single-bond"
            - "aromatic-bond"
            - "double-bond"
            - "triple-bond"
            - "is-carbon"
            - "group"
            - "period"

        offset_carbon:
            Whether to subract the Carbon property from the desired atomic property.
            For example, if we want the mass of the Lithium (6.941), the mass of the
            Carbon (12.0107) will be subracted, resulting in a value of -5.0697

        mask_nan:
            Deal with molecules that fail a part of the featurization.
            NaNs can happen when taking the of a noble gas,
            or other properties that are not measured for specific atoms.

            - "raise": Raise an error when there is a nan or inf in the featurization
            - "warn": Raise a warning when there is a nan or inf in the featurization
            - "None": DEFAULT. Don't do anything
            - "Floating value": Replace nans or inf by the specified value

    Returns:

        prop_dict:
            A dictionnary where the element of ``property_list`` are the keys
            and the values are np.ndarray of shape (N,). N is the number of atoms
            in ``mol``.

    """

    periodic_table = Chem.GetPeriodicTable()
    prop_dict = {}
    C = Chem.Atom("C")
    C_num = C.GetAtomicNum()
    offC = bool(offset_carbon)
    atom_list = list(mol.GetAtoms())

    for prop in property_list:
        prop_name = None

        property_array = np.zeros(mol.GetNumAtoms(), dtype=np.float16)
        for ii, atom in enumerate(atom_list):
            val = None
            atomic_num = atom.GetAtomicNum()

            if isinstance(prop, str):
                prop = prop.lower()
                prop_name = prop

                if prop in ["atomic-number"]:
                    val = (atomic_num - (offC * C_num)) / 5
                elif prop in ["mass", "weight"]:
                    prop_name = "mass"
                    val = (atom.GetMass() - (offC * C.GetMass())) / 10
                elif prop in ["valence", "total-valence"]:
                    prop_name = "valence"
                    val = atom.GetTotalValence() - (offC * 4)
                elif prop in ["implicit-valence"]:
                    val = atom.GetImplicitValence()
                elif prop in ["hybridization"]:
                    val = atom.GetHybridization()
                elif prop in ["chirality"]:
                    val = (atom.GetProp("_CIPCode") == "R") if atom.HasProp("_CIPCode") else 2
                elif prop in ["hybridization"]:
                    val = atom.GetHybridization()
                elif prop in ["aromatic"]:
                    val = atom.GetIsAromatic()
                elif prop in ["ring", "in-ring"]:
                    prop_name = "in-ring"
                    val = atom.IsInRing()
                elif prop in ["min-ring"]:
                    ring_info = mol.GetRingInfo()
                    val = ring_info.MinAtomRingSize(atom.GetIdx())
                elif prop in ["max-ring"]:
                    rings = mol.GetRingInfo().AtomRings()
                    val = 0
                    for ring in rings:
                        if atom.GetIdx() in ring:
                            if len(ring) > val:
                                val = len(ring)
                elif prop in ["num-ring"]:
                    ring_info = mol.GetRingInfo()
                    val = ring_info.NumAtomRings(atom.GetIdx())
                elif prop in ["degree"]:
                    val = atom.GetTotalDegree() - (offC * 2)
                elif prop in ["radical-electron"]:
                    val = atom.GetNumRadicalElectrons()
                elif prop in ["formal-charge"]:
                    val = atom.GetFormalCharge()
                elif prop in ["vdw-radius"]:
                    val = periodic_table.GetRvdw(atom.GetAtomicNum()) - offC * periodic_table.GetRvdw(C_num)
                elif prop in ["covalent-radius"]:
                    val = periodic_table.GetRcovalent(atomic_num) - offC * periodic_table.GetRcovalent(C_num)
                elif prop in ["electronegativity"]:
                    val = (
                        nmp.ELECTRONEGATIVITY[atom.GetAtomicNum() - 1]
                        - offC * nmp.ELECTRONEGATIVITY[C_num - 1]
                    )
                elif prop in ["ionization", "first-ionization"]:
                    prop_name = "ionization"
                    val = (nmp.FIRST_IONIZATION[atomic_num - 1] - offC * nmp.FIRST_IONIZATION[C_num - 1]) / 5
                elif prop in ["melting-point"]:
                    val = (nmp.MELTING_POINT[atomic_num - 1] - offC * nmp.MELTING_POINT[C_num - 1]) / 200
                elif prop in ["metal"]:
                    val = nmp.METAL[atomic_num - 1]
                elif prop in "group":
                    val = float(nmp.GROUP[atomic_num - 1]) - offC * float(nmp.GROUP[C_num - 1])
                elif prop in "period":
                    val = float(nmp.PERIOD[atomic_num - 1]) - offC * float(nmp.PERIOD[C_num - 1])
                elif "-bond" in prop:
                    bonds = [bond.GetBondTypeAsDouble() for bond in atom.GetBonds()]
                    if prop in ["single-bond"]:
                        val = len([bond == 1 for bond in bonds])
                    elif prop in ["aromatic-bond"]:
                        val = len([bond == 1.5 for bond in bonds])
                    elif prop in ["double-bond"]:
                        val = len([bond == 2 for bond in bonds])
                    elif prop in ["triple-bond"]:
                        val = len([bond == 3 for bond in bonds])
                    else:
                        raise ValueError(f"{prop} is not a correct bond.")
                    val -= offC * 1
                elif prop in ["is-carbon"]:
                    val = atom.GetAtomicNum() == 6
                    val -= offC * 1
                else:
                    raise ValueError(f"Unsupported property `{prop}`")

            elif callable(prop):
                prop_name = str(prop)
                val = prop(atom)
            else:
                ValueError(f"Elements in `property_list` must be str or callable, provided `{type(prop)}`")

            if val is None:
                raise ValueError("val is undefined.")

            property_array[ii] = val

        if prop_name is None:
            raise ValueError("prop_name is undefined.")

        # Mask the NaNs
        prop_dict[prop_name] = _mask_nans_inf(mask_nan, property_array, "atom featurization")

    return prop_dict


def get_simple_mol_conformer(mol: dm.Mol) -> Union[Chem.rdchem.Conformer, None]:
    r"""
    If the molecule has a conformer, then it will return the conformer at idx `0`.
    Otherwise, it generates a simple molecule conformer using `rdkit.Chem.rdDistGeom.EmbedMolecule`
    and returns it. This is meant to be used in simple functions like `GetBondLength`,
    not in functions requiring complex 3D structure.

    Parameters:

        mol: Rdkit Molecule

    Returns:
        conf: A conformer of the molecule, or `None` if it fails
    """

    val = 0
    if mol.GetNumConformers() == 0:
        val = Chem.rdDistGeom.EmbedMolecule(mol)
    if val == -1:
        val = Chem.rdDistGeom.EmbedMolecule(
            mol,
            enforceChirality=False,
            ignoreSmoothingFailures=True,
            useBasicKnowledge=True,
            useExpTorsionAnglePrefs=True,
            forceTol=0.1,
        )

    if val == -1:
        conf = None
        logger.warn("Couldn't compute conformer for molecule `{}`".format(Chem.MolToSmiles(mol)))
    else:
        conf = mol.GetConformer(0)

    return conf


def get_estimated_bond_length(bond: Chem.rdchem.Bond, mol: dm.Mol) -> float:
    r"""
    Estimate the bond length between atoms by looking at the estimated atomic radius
    that depends both on the atom type and the bond type. The resulting bond-length is
    then the sum of the radius.

    Keep in mind that this function only provides an estimate of the bond length and not
    the true one based on a conformer. The vast majority od estimated bond lengths will
    have an error below 5% while some bonds can have an error up to 20%. This function
    is mostly useful when conformer generation fails for some molecules, or for
    increased computation speed.

    Parameters:
        bond: The bond to measure its lenght
        mol: The molecule containing the bond (used to get neighbouring atoms)

    Returns:
        bond_length: The bond length in Angstrom, typically a value around 1-2.

    """

    # Get the atoms connected by the bond
    idx1 = bond.GetBeginAtomIdx()
    idx2 = bond.GetEndAtomIdx()
    atom1 = mol.GetAtomWithIdx(idx1).GetAtomicNum()
    atom2 = mol.GetAtomWithIdx(idx2).GetAtomicNum()
    bond_type = bond.GetBondType()

    # Get single bond atomic radius
    if bond_type == Chem.rdchem.BondType.SINGLE:
        rad1 = [nmp.BOND_RADIUS_SINGLE[atom1 - 1]]
        rad2 = [nmp.BOND_RADIUS_SINGLE[atom2 - 1]]
    # Get double bond atomic radius
    elif bond_type == Chem.rdchem.BondType.DOUBLE:
        rad1 = [nmp.BOND_RADIUS_DOUBLE[atom1 - 1]]
        rad2 = [nmp.BOND_RADIUS_DOUBLE[atom2 - 1]]
    # Get triple bond atomic radius
    elif bond_type == Chem.rdchem.BondType.TRIPLE:
        rad1 = [nmp.BOND_RADIUS_TRIPLE[atom1 - 1]]
        rad2 = [nmp.BOND_RADIUS_TRIPLE[atom2 - 1]]
    # Get average of single bond and double bond atomic radius
    elif bond_type == Chem.rdchem.BondType.AROMATIC:
        rad1 = [nmp.BOND_RADIUS_SINGLE[atom1 - 1], nmp.BOND_RADIUS_DOUBLE[atom1 - 1]]
        rad2 = [nmp.BOND_RADIUS_SINGLE[atom2 - 1], nmp.BOND_RADIUS_DOUBLE[atom2 - 1]]

    # Average the bond lengths, while ignoring nans in case some missing value
    rad1_float = [elem for elem in rad1 if elem is not None]
    rad2_float = [elem for elem in rad2 if elem is not None]

    if len(rad1_float) > 0:
        rad1_float = sum(rad1_float) / len(rad1_float)
    else:
        rad1_float = float(nmp.BOND_RADIUS_SINGLE[atom1 - 1])

    if len(rad2_float) > 0:
        rad2_float = sum(rad2_float) / len(rad2_float)
    else:
        rad2_float = float(nmp.BOND_RADIUS_SINGLE[atom2 - 1])

    bond_length = rad1_float + rad2_float
    return bond_length


def get_mol_edge_features(
    mol: dm.Mol, property_list: List[str], mask_nan: Union[str, float, type(None)] = "raise"
) -> Dict[str, np.ndarray]:
    r"""
    Get the following set of features for any given bond
    See `graphium.features.nmp` for allowed values in one hot encoding

    * One-hot representation of the bond type. Note that you should not kekulize your
        molecules, if you expect this to take aromatic bond into account.
    * Bond stereo type, following CIP classification
    * Whether the bond is conjugated
    * Whether the bond is in a ring

    Parameters:
        mol: rdkit.Chem.Molecule
            the molecule of interest

        property_list:
            A list of edge properties to return for the given molecule.
            Accepted properties are:

            - "bond-type-onehot"
            - "bond-type-float"
            - "stereo"
            - "in-ring"
            - "conjugated"
            - "conformer-bond-length" (might cause problems with complex molecules)
            - "estimated-bond-length"

    Returns:
        prop_dict:
            A dictionnary where the element of ``property_list`` are the keys
            and the values are np.ndarray of shape (N,). N is the number of atoms
            in ``mol``.

    """

    prop_dict = {}

    # Compute features for each bond
    num_bonds = mol.GetNumBonds()
    for prop in property_list:
        property_array = []
        for ii in range(num_bonds):
            prop = prop.lower()
            bond = mol.GetBondWithIdx(ii)

            if prop in ["bond-type-onehot"]:
                encoding = one_of_k_encoding(bond.GetBondType(), nmp.BOND_TYPES)
            elif prop in ["bond-type-float"]:
                encoding = [bond.GetBondTypeAsDouble()]
            elif prop in ["stereo"]:
                encoding = one_of_k_encoding(bond.GetStereo(), nmp.BOND_STEREO)
            elif prop in ["in-ring"]:
                encoding = [bond.IsInRing()]
            elif prop in ["conjugated"]:
                encoding = [bond.GetIsConjugated()]
            elif prop in ["conformer-bond-length"]:
                conf = get_simple_mol_conformer(mol)
                if conf is not None:
                    idx1 = bond.GetBeginAtomIdx()
                    idx2 = bond.GetEndAtomIdx()
                    encoding = [Chem.rdMolTransforms.GetBondLength(conf, idx1, idx2)]
                else:
                    encoding = [0]
            elif prop in ["estimated-bond-length"]:
                encoding = [get_estimated_bond_length(bond, mol)]

            else:
                raise ValueError(f"Unsupported property `{prop}`")

            property_array.append(np.asarray(encoding, dtype=np.float16))

        if num_bonds > 0:
            property_array = np.stack(property_array, axis=0)
            # Mask the NaNs
            prop_dict[prop] = _mask_nans_inf(mask_nan, property_array, "edge property")
        else:
            # Add an empty vector with the right shape
            arr_len = 1
            if prop in ["bond-type-onehot"]:
                arr_len = len(nmp.BOND_TYPES) + 1
            elif prop in ["stereo"]:
                arr_len = len(nmp.BOND_STEREO) + 1

            prop_dict[prop] = np.zeros((0, arr_len))

    return prop_dict


def mol_to_adj_and_features(
    mol: Union[str, dm.Mol],
    atom_property_list_onehot: List[str] = [],
    atom_property_list_float: List[Union[str, Callable]] = [],
    conformer_property_list: List[str] = [],
    edge_property_list: List[str] = [],
    add_self_loop: bool = False,
    explicit_H: bool = False,
    use_bonds_weights: bool = False,
    pos_encoding_as_features: Dict[str, Any] = None,
    dtype: np.dtype = np.float16,
    mask_nan: Union[str, float, type(None)] = "raise",
) -> Union[
    coo_matrix,
    Union[Tensor, None],
    Union[Tensor, None],
    Dict[str, Tensor],
    Union[Tensor, None],
    Dict[str, Tensor],
]:
    r"""
    Transforms a molecule into an adjacency matrix representing the molecular graph
    and a set of atom and bond features.

    It also returns the positional encodings associated to the graph.

    Parameters:

        mol:
            The molecule to be converted

        atom_property_list_onehot:
            List of the properties used to get one-hot encoding of the atom type,
            such as the atom index represented as a one-hot vector.
            See function `get_mol_atomic_features_onehot`

        atom_property_list_float:
            List of the properties used to get floating-point encoding of the atom type,
            such as the atomic mass or electronegativity.
            See function `get_mol_atomic_features_float`

        conformer_property_list:
            list of properties used to encode the conformer information, outside of atom properties, currently support "positions_3d"

        edge_property_list:
            List of the properties used to encode the edges, such as the edge type
            and the stereo type.

        add_self_loop:
            Whether to add a value of `1` on the diagonal of the adjacency matrix.

        explicit_H:
            Whether to consider the Hydrogens explicitely. If `False`, the hydrogens
            are implicit.

        use_bonds_weights:
            Whether to use the floating-point value of the bonds in the adjacency matrix,
            such that single bonds are represented by 1, double bonds 2, triple 3, aromatic 1.5

        pos_encoding_as_features: keyword arguments for function `graph_positional_encoder`
            to generate positional encoding for node features.

        dtype:
            The torch data type used to build the graph

        mask_nan:
            Deal with molecules that fail a part of the featurization.
            NaNs can happen when taking the of a noble gas,
            or other properties that are not measured for specific atoms.

            - "raise": Raise an error when there is a nan or inf in the featurization
            - "warn": Raise a warning when there is a nan or inf in the featurization
            - "None": DEFAULT. Don't do anything
            - "Floating value": Replace nans or inf by the specified value
    Returns:

        adj:
            torch coo sparse adjacency matrix of the molecule

        ndata:
            Concatenated node data of the atoms, based on the properties from
            `atom_property_list_onehot` and `atom_property_list_float`.
            If no properties are given, it returns `None`

        edata:
            Concatenated node edge of the molecule, based on the properties from
            `edge_property_list`.
            If no properties are given, it returns `None`

        pe_dict:
            Dictionary of all positional encodings. Current supported keys:

            - "pos_enc_feats_sign_flip":
                Node positional encoding that requires augmentation via sign-flip.
                For example, eigenvectors of the Laplacian are ambiguous to the
                sign and are returned here.

            - "pos_enc_feats_no_flip":
                Node positional encoding that requires does not use sign-flip.
                For example, distance from centroid are returned here.

            - "rwse":
                Node structural encoding corresponding to the diagonal of the random
                walk matrix

        conf_dict:
            contains the 3d positions of a conformer of the molecule or 0s if none is found

    """

    if isinstance(mol, str):
        mol = dm.to_mol(mol)

    # Add or remove explicit hydrogens
    if explicit_H:
        mol = Chem.AddHs(mol)
    else:
        mol = Chem.RemoveHs(mol)

    num_nodes = mol.GetNumAtoms()

    adj = mol_to_adjacency_matrix(
        mol, use_bonds_weights=use_bonds_weights, add_self_loop=add_self_loop, dtype=dtype
    )

    # Get the node features
    atom_features_onehot = get_mol_atomic_features_onehot(mol, atom_property_list_onehot)
    atom_features_float = get_mol_atomic_features_float(mol, atom_property_list_float, mask_nan=mask_nan)
    conf_dict = get_mol_conformer_features(mol, conformer_property_list, mask_nan=mask_nan)
    ndata = list(atom_features_float.values()) + list(atom_features_onehot.values())
    ndata = [d[:, np.newaxis] if d.ndim == 1 else d for d in ndata]

    if len(ndata) > 0:
        ndata = np.concatenate(ndata, axis=1).astype(dtype=dtype)
    else:
        ndata = None

    # Get the edge features
    edge_features = get_mol_edge_features(mol, edge_property_list, mask_nan=mask_nan)
    edata = list(edge_features.values())
    edata = [np.expand_dims(d, axis=1) if d.ndim == 1 else d for d in edata]
    if len(edata) > 0:
        edata = np.concatenate(edata, axis=1).astype(dtype=dtype)
    else:
        edata = None

    # Get all positional encodings
    pe_dict = get_all_positional_encodings(adj, num_nodes, pos_encoding_as_features)

    # Mask the NaNs
    for pe_key, pe_val in pe_dict.items():
        pe_val = np.asarray(pe_val, dtype=dtype)
        pe_dict[pe_key] = _mask_nans_inf(mask_nan, pe_val, pe_key)

    return adj, ndata, edata, pe_dict, conf_dict


def mol_to_adjacency_matrix(
    mol: dm.Mol,
    use_bonds_weights: bool = False,
    add_self_loop: bool = False,
    dtype: np.dtype = np.float32,
) -> coo_matrix:
    r"""
    Convert a molecule to a sparse adjacency matrix, as a torch Tensor.
    Instead of using the Rdkit `GetAdjacencyMatrix()` method, this method
    uses the bond ordering from the molecule object, which is the same as
    the bond ordering in the bond features.

    Warning:
        Do not use `Tensor.coalesce()` on the returned adjacency matrix, as it
        will change the ordering of the bonds.

    Args:
        mol: A molecule in the form of a SMILES string or an RDKit molecule object.

        use_bonds_weights:
            If `True`, the adjacency matrix will contain the bond type as the
            value of the edge. If `False`, the adjacency matrix will contain
            `1` as the value of the edge.

        add_self_loop:
            If `True`, the adjacency matrix will contain a self-loop for each
            node.

        dtype:
            The data type used to build the graph

    Returns:
        adj:
            coo sparse adjacency matrix of the molecule
    """

    # Get the indices for the adjacency matrix, and the bond value
    adj_idx, adj_val = [], []
    for bond in mol.GetBonds():
        adj_idx.append([bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()])
        adj_idx.append([bond.GetEndAtomIdx(), bond.GetBeginAtomIdx()])
        if use_bonds_weights:
            val = nmp.BOND_TYPES[bond.GetBondType()]
        else:
            val = 1
        adj_val.extend([val, val])

    # Convert to torch coo sparse tensor
    if len(adj_val) > 0:  # ensure tensor is not empty
        adj = coo_matrix(
            (torch.as_tensor(adj_val), torch.as_tensor(adj_idx).T.reshape(2, -1)),
            shape=(mol.GetNumAtoms(), mol.GetNumAtoms()),
            dtype=dtype,
        )
    else:
        # Special case for molecules with one atom
        adj = coo_matrix(([], np.array([[], []])), shape=(mol.GetNumAtoms(), mol.GetNumAtoms()), dtype=dtype)

    # Add self loops
    if add_self_loop:
        arange = np.arange(adj.shape[0], dtype=int)
        adj[arange, arange] = 1
    return adj


class GraphDict(dict):
    def __init__(
        self,
        dic: Dict,
    ):
        r"""
        Store the parameters required to initialize a `pyg.data.Data`, but
        as a dictionary to reduce memory consumption.

        Possible keys for the dictionary:

        - adj: A sparse Tensor containing the adjacency matrix

        - ndata: A dictionnary containing different keys and Tensors
            associated to the node features.

        - edata: A dictionnary containing different keys and Tensors
            associated to the edge features.

        - dtype: The dtype for the floating data.

        - mask_nan:
            Deal with molecules that fail a part of the featurization.
            NaNs can happen when taking the of a noble gas,
            or other properties that are not measured for specific atoms.

            - "raise": Raise an error when there is a nan or inf in the featurization
            - "warn": Raise a warning when there is a nan or inf in the featurization
            - "None": DEFAULT. Don't do anything
            - "Floating value": Replace nans or inf by the specified value
        """
        default_dic = {
            "dtype": np.float16,
            "mask_nan": "raise",
        }
        data = dic.pop("data", {})
        # ndata = dic.pop("ndata", {})
        # edata = dic.pop("edata", {})
        # for key in edata.keys():
        #     assert key.startswith("edge_"), f"Edge features must start with 'edge_' but got {key}"
        default_dic.update(dic)
        default_dic.update(data)
        # default_dic.update(ndata)
        # default_dic.update(edata)
        super().__init__(default_dic)

    @property
    def keys(self):
        return list(super().keys())

    @property
    def values(self):
        return list(super().self.values())

    def make_pyg_graph(self, **kwargs) -> Data:
        """
        Convert the current dictionary of parameters, containing an adjacency matrix with node/edge data
        into a `pyg.data.Data` of torch Tensors.

        `**kwargs` can be used to overwrite any parameter from the current dictionary. See `GraphDict.__init__`
        for a list of parameters
        """

        num_nodes = self.adj.shape[0]
        data_dict = {}

        # Convert the numpy and numpy sparse data to torch
        for key, val in self.items():
            if key in ["adj", "dtype", "mask_nan"]:  # Skip the parameters
                continue
            elif isinstance(val, np.ndarray):
                # Convert the data to the specified dtype in torch format
                val = val.astype(self.dtype)
                data_dict[key] = torch.as_tensor(val)
            elif issparse(val):
                data_dict[key] = torch.as_tensor(val.astype(np.float32).todense())
                # `torch.sparse_coo_tensor` is too slow. Slows down the multiprocessing of features by >3x on 32 cores.
                # indices = torch.from_numpy(np.vstack((val.row, val.col)).astype(np.int64))
                # data_dict[key] = torch.sparse_coo_tensor(indices=indices, values=val.data, size=val.shape)
            elif isinstance(val, torch.Tensor):
                data_dict[key] = val
            else:
                pass  # Skip the other parameters

        # Create the PyG graph object `Data`
        edge_index = torch.as_tensor(np.vstack((self.adj.row, self.adj.col)))
        edge_weight = torch.as_tensor(self.adj.data)
        data = Data(edge_index=edge_index, edge_weight=edge_weight, num_nodes=num_nodes, **data_dict)
        return data

    @property
    def adj(self):
        return self["adj"]

    @property
    def dtype(self):
        return self["dtype"]

    @property
    def mask_nan(self):
        return self["mask_nan"]

    @property
    def num_nodes(self) -> int:
        return self.adj.shape[0]

    @property
    def num_edges(self) -> int:
        if issparse(self.adj):
            return self.adj.nnz
        else:
            return np.count_nonzero(self.adj)  # No division by 2 because edges are counted twice


def mol_to_graph_dict(
    mol: dm.Mol,
    atom_property_list_onehot: List[str] = [],
    atom_property_list_float: List[Union[str, Callable]] = [],
    conformer_property_list: List[str] = [],
    edge_property_list: List[str] = [],
    add_self_loop: bool = False,
    explicit_H: bool = False,
    use_bonds_weights: bool = False,
    pos_encoding_as_features: Dict[str, Any] = None,
    dtype: np.dtype = np.float16,
    on_error: str = "ignore",
    mask_nan: Union[str, float, type(None)] = "raise",
    max_num_atoms: Optional[int] = None,
) -> Union[GraphDict, str]:
    r"""
    Transforms a molecule into an adjacency matrix representing the molecular graph
    and a set of atom and bond features, and re-organizes them into a dictionary
    that allows to build a `pyg.data.Data` object.

    Compared to `mol_to_pyggraph`, this function does not build the graph directly,
    and is thus faster, less memory heavy, and compatible with other frameworks.

    Parameters:

        mol:
            The molecule to be converted

        atom_property_list_onehot:
            List of the properties used to get one-hot encoding of the atom type,
            such as the atom index represented as a one-hot vector.
            See function `get_mol_atomic_features_onehot`

        atom_property_list_float:
            List of the properties used to get floating-point encoding of the atom type,
            such as the atomic mass or electronegativity.
            See function `get_mol_atomic_features_float`

        conformer_property_list:
            list of properties used to encode the conformer information, outside of atom properties, currently support "positions_3d"

        edge_property_list:
            List of the properties used to encode the edges, such as the edge type
            and the stereo type.

        add_self_loop:
            Whether to add a value of `1` on the diagonal of the adjacency matrix.

        explicit_H:
            Whether to consider the Hydrogens explicitely. If `False`, the hydrogens
            are implicit.

        use_bonds_weights:
            Whether to use the floating-point value of the bonds in the adjacency matrix,
            such that single bonds are represented by 1, double bonds 2, triple 3, aromatic 1.5

        pos_encoding_as_features: keyword arguments for function `graph_positional_encoder`
            to generate positional encoding for node features.

        dtype:
            The numpy data type used to build the graph

        on_error:
            What to do when the featurization fails. This can change the
            behavior of `mask_nan`.

            - "raise": Raise an error
            - "warn": Raise a warning and return a string of the error
            - "ignore": Ignore the error and return a string of the error

        mask_nan:
            Deal with molecules that fail a part of the featurization.
            NaNs can happen when taking the of a noble gas,
            or other properties that are not measured for specific atoms.

            - "raise": Raise an error when there is a nan or inf in the featurization
            - "warn": Raise a warning when there is a nan or inf in the featurization
            - "None": DEFAULT. Don't do anything
            - "Floating value": Replace nans or inf by the specified value

        max_num_atoms:
            Maximum number of atoms for a given molecule. If a molecule with more atoms
            is give, an error is raised, but catpured according to the rules of
            `on_error`.

    Returns:

        graph_dict:
            A dictionary `GraphDict` containing the keys required to build a graph,
            and which can be used to build a PyG graph. If it fails
            to featurize the molecule, it returns a string with the error.

            - "adj": A sparse int-array containing the adjacency matrix

            - "data": A dictionnary containing different keys and numpy
              arrays associated to the (node, edge & graph) features.

            - "dtype": The numpy dtype for the floating data.
    """

    input_mol = mol
    try:
        if isinstance(mol, str):
            mol = dm.to_mol(mol)
        if explicit_H:
            mol = Chem.AddHs(mol)
        else:
            mol = Chem.RemoveHs(mol)

        num_atoms = mol.GetNumAtoms()
        if (max_num_atoms is not None) and (num_atoms > max_num_atoms):
            raise ValueError(f"Maximum number of atoms greater than permitted {num_atoms}>{max_num_atoms}")

        (
            adj,
            ndata,
            edata,
            pe_dict,
            conf_dict,
        ) = mol_to_adj_and_features(
            mol=mol,
            atom_property_list_onehot=atom_property_list_onehot,
            atom_property_list_float=atom_property_list_float,
            conformer_property_list=conformer_property_list,
            edge_property_list=edge_property_list,
            add_self_loop=add_self_loop,
            explicit_H=explicit_H,
            use_bonds_weights=use_bonds_weights,
            pos_encoding_as_features=pos_encoding_as_features,
            mask_nan=mask_nan,
        )
    except Exception as e:
        if on_error.lower() == "raise":
            raise e
        elif on_error.lower() == "warn":
            smiles = input_mol
            if isinstance(smiles, dm.Mol):
                smiles = Chem.MolToSmiles(input_mol)
            msg = str(e) + "\nIgnoring following molecule:" + smiles
            logger.warning(msg)
            return str(e)
        elif on_error.lower() == "ignore":
            return str(e)

    graph_dict = {"adj": adj, "data": {}, "dtype": dtype}

    # Assign the node data
    if ndata is not None:
        graph_dict["data"]["feat"] = ndata

    # Assign the edge data
    if edata is not None:
        if issparse(edata):
            edata = to_dense_array(edata, dtype=dtype)
        hetero_edata = edata.repeat(2, axis=0)
        graph_dict["data"]["edge_feat"] = hetero_edata

    # Put the positional encodings as node features
    # TODO: add support for PE on edges
    for key, pe in pe_dict.items():
        graph_dict["data"][key] = pe

    # put the conformer positions here
    for key, val in conf_dict.items():
        graph_dict["data"][key] = val

    graph_dict = GraphDict(graph_dict)
    return graph_dict


def mol_to_pyggraph(
    mol: dm.Mol,
    atom_property_list_onehot: List[str] = [],
    atom_property_list_float: List[Union[str, Callable]] = [],
    conformer_property_list: List[str] = [],
    edge_property_list: List[str] = [],
    add_self_loop: bool = False,
    explicit_H: bool = False,
    use_bonds_weights: bool = False,
    pos_encoding_as_features: Dict[str, Any] = None,
    dtype: np.dtype = np.float16,
    on_error: str = "ignore",
    mask_nan: Union[str, float, type(None)] = "raise",
    max_num_atoms: Optional[int] = None,
) -> Union[Data, str]:
    r"""
    Transforms a molecule into an adjacency matrix representing the molecular graph
    and a set of atom and bond features.

    Then, the adjacency matrix and node/edge features are used to build a
    `pyg.data.Data` with pytorch Tensors.

    Parameters:

        mol:
            The molecule to be converted

        atom_property_list_onehot:
            List of the properties used to get one-hot encoding of the atom type,
            such as the atom index represented as a one-hot vector.
            See function `get_mol_atomic_features_onehot`

        atom_property_list_float:
            List of the properties used to get floating-point encoding of the atom type,
            such as the atomic mass or electronegativity.
            See function `get_mol_atomic_features_float`

        conformer_property_list:
            list of properties used to encode the conformer information, outside of atom properties, currently support "positions_3d"

        edge_property_list:
            List of the properties used to encode the edges, such as the edge type
            and the stereo type.

        add_self_loop:
            Whether to add a value of `1` on the diagonal of the adjacency matrix.

        explicit_H:
            Whether to consider the Hydrogens explicitely. If `False`, the hydrogens
            are implicit.

        use_bonds_weights:
            Whether to use the floating-point value of the bonds in the adjacency matrix,
            such that single bonds are represented by 1, double bonds 2, triple 3, aromatic 1.5

        pos_encoding_as_features: keyword arguments for function `graph_positional_encoder`
            to generate positional encoding for node features.

        dtype:
            The numpy data type used to build the graph

        on_error:
            What to do when the featurization fails. This can change the
            behavior of `mask_nan`.

            - "raise": Raise an error
            - "warn": Raise a warning and return a string of the error
            - "ignore": Ignore the error and return a string of the error

        mask_nan:
            Deal with molecules that fail a part of the featurization.
            NaNs can happen when taking the of a noble gas,
            or other properties that are not measured for specific atoms.

            - "raise": Raise an error when there is a nan in the featurization
            - "warn": Raise a warning when there is a nan in the featurization
            - "None": DEFAULT. Don't do anything
            - "Floating value": Replace nans by the specified value

        max_num_atoms:
            Maximum number of atoms for a given molecule. If a molecule with more atoms
            is give, an error is raised, but catpured according to the rules of
            `on_error`.
    Returns:

        graph:
            Pyg graph, with `graph['feat']` corresponding to the concatenated
            node data from `atom_property_list_onehot` and `atom_property_list_float`,
            `graph['edge_feat']` corresponding to the concatenated edge data from `edge_property_list`.
            There are also additional entries for the positional encodings.

    """
    graph_dict = mol_to_graph_dict(
        mol=mol,
        atom_property_list_onehot=atom_property_list_onehot,
        atom_property_list_float=atom_property_list_float,
        conformer_property_list=conformer_property_list,
        edge_property_list=edge_property_list,
        add_self_loop=add_self_loop,
        explicit_H=explicit_H,
        use_bonds_weights=use_bonds_weights,
        pos_encoding_as_features=pos_encoding_as_features,
        dtype=dtype,
        on_error=on_error,
        mask_nan=mask_nan,
        max_num_atoms=max_num_atoms,
    )

    if (graph_dict is not None) and not isinstance(graph_dict, str):
        return graph_dict.make_pyg_graph()
    else:
        return graph_dict


def mol_to_graph_signature(featurizer_args: Dict[str, Any] = None) -> Dict[str, Any]:
    """
    Get the default arguments of `mol_to_graph_dict` and update it
    with a provided dict of arguments in order to get a fulle signature
    of the featurizer args actually used for the features computation.

    Parameters:
        featurizer_args: A dictionary of featurizer arguments to update
    Returns:
        A dictionary of featurizer arguments
    """

    # Get the signature of `mol_to_graph_dict`
    signature = inspect.signature(mol_to_graph_dict)

    # Filter out empty arguments (without default value)
    parameters = list(filter(lambda param: param.default is not param.empty, signature.parameters.values()))

    # Convert to dict
    parameters = {param.name: param.default for param in parameters}

    # Update the parameters with the supplied ones
    if featurizer_args is not None:
        parameters.update(featurizer_args)

    return parameters


# graphium/features/graphormer.py

from typing import Tuple, Union, Dict, Any

import numpy as np
import networkx as nx

from scipy.sparse import spmatrix, issparse


def compute_graphormer_distances(
    adj: Union[np.ndarray, spmatrix], num_nodes: int, cache: Dict[str, Any]
) -> Tuple[np.ndarray, str, Dict[str, Any]]:
    """
    Compute Graphormer distance between nodepairs.

    Parameters:
        adj [num_nodes, num_nodes]: Adjacency matrix
        num_nodes: Number of nodes in the graph
        cache: Dictionary of cached objects
    Returns:
        dist [num_nodes, num_nodes]: 2D array with Graphormer distances between nodepairs
        base_level: Indicator of the output pos_level (node, edge, nodepair, graph) -> here nodepair
        cache: Updated dictionary of cached objects
    """

    base_level = "nodepair"

    if "graphormer" in cache:
        dist = cache["graphormer"]

    else:
        if issparse(adj):
            adj = adj.toarray()

        G = nx.from_numpy_array(adj)
        paths = nx.all_pairs_shortest_path(G)

        dist_dict = {i: {j: len(path) - 1 for j, path in paths_from_i.items()} for i, paths_from_i in paths}
        dist = np.asarray([[dist_dict[i][j] for j in range(num_nodes)] for i in range(num_nodes)])
        cache["graphormer"] = dist

    return dist, base_level, cache


# graphium/features/nmp.py

from typing import Tuple, Optional, Dict, Union
import importlib.resources
from copy import deepcopy
import pandas as pd
import numpy as np
import math
from rdkit import Chem

# NOTE(hadim): usually it's best to embed this in a function.
with importlib.resources.open_text("graphium.features", "periodic_table.csv") as f:
    PERIODIC_TABLE = pd.read_csv(f)
PERIODIC_TABLE = PERIODIC_TABLE.set_index("AtomicNumber")


# Small function to convert strings to floats
def float_or_none(string: str) -> Union[float, None]:
    """
    check if a string can be converted to float, return none if it can't
    Parameters:
        string: str
    Returns:
        val: float or None
    """
    try:
        val = float(string)
    except:
        val = None
    return val


# It's much faster to index from a list than a DataFrame, which can have big impact
# when featurizing millions of molecules
BOND_RADIUS_SINGLE = [float_or_none(elem) for elem in PERIODIC_TABLE["SingleBondRadius"]]
BOND_RADIUS_DOUBLE = [float_or_none(elem) for elem in PERIODIC_TABLE["DoubleBondRadius"]]
BOND_RADIUS_TRIPLE = [float_or_none(elem) for elem in PERIODIC_TABLE["TripleBondRadius"]]
ELECTRONEGATIVITY = [float_or_none(elem) for elem in PERIODIC_TABLE["Electronegativity"]]
FIRST_IONIZATION = [float_or_none(elem) for elem in PERIODIC_TABLE["FirstIonization"]]
MELTING_POINT = [float_or_none(elem) for elem in PERIODIC_TABLE["MeltingPoint"]]
METAL = (2 * (PERIODIC_TABLE["Metal"] == "yes") + (PERIODIC_TABLE["Metalloid"] == "yes")).tolist()

PHASE = list(PERIODIC_TABLE["Phase"].values)
PHASE_SET = list(set(PHASE))
TYPE = list(deepcopy(PERIODIC_TABLE["Type"]))
TYPE = ["none" if (isinstance(t, float) and math.isnan(t)) else t for t in TYPE]
TYPE_SET = list(set(TYPE))
GROUP = deepcopy(PERIODIC_TABLE["Group"].values)
GROUP[np.isnan(GROUP)] = 19
GROUP_SET = list(set(GROUP))
PERIOD = list(PERIODIC_TABLE["Period"].values)
PERIOD_SET = list(set(PERIOD))

ATOM_LIST = [
    "C",
    "N",
    "O",
    "S",
    "F",
    "Si",
    "P",
    "Cl",
    "Br",
    "Mg",
    "Na",
    "Ca",
    "Fe",
    "As",
    "Al",
    "I",
    "B",
    "V",
    "K",
    "Tl",
    "Yb",
    "Sb",
    "Sn",
    "Ag",
    "Pd",
    "Co",
    "Se",
    "Ti",
    "Zn",
    "H",
    "Li",
    "Ge",
    "Cu",
    "Au",
    "Ni",
    "Cd",
    "In",
    "Mn",
    "Zr",
    "Cr",
    "Pt",
    "Hg",
    "Pb",
]

ATOM_NUM_H = [0, 1, 2, 3, 4]
VALENCE = [0, 1, 2, 3, 4, 5, 6]
CHARGE_LIST = [-3, -2, -1, 0, 1, 2, 3]
RADICAL_E_LIST = [0, 1, 2]
ATOM_DEGREE_LIST = [0, 1, 2, 3, 4]

HYBRIDIZATION_LIST = [
    Chem.rdchem.HybridizationType.names[k]
    for k in sorted(Chem.rdchem.HybridizationType.names.keys(), reverse=True)
    if k != "OTHER"
]


CHIRALITY_LIST = ["R"]  # alternative is just S


BOND_TYPES = [
    Chem.rdchem.BondType.SINGLE,
    Chem.rdchem.BondType.DOUBLE,
    Chem.rdchem.BondType.TRIPLE,
    Chem.rdchem.BondType.AROMATIC,
]

BOND_STEREO = [
    Chem.rdchem.BondStereo.STEREONONE,
    Chem.rdchem.BondStereo.STEREOANY,
    Chem.rdchem.BondStereo.STEREOZ,
    Chem.rdchem.BondStereo.STEREOE,
    Chem.rdchem.BondStereo.STEREOCIS,
    Chem.rdchem.BondStereo.STEREOTRANS,
]


# graphium/features/positional_encoding.py

from typing import Tuple, Union, Optional, Dict, Any, OrderedDict
from copy import deepcopy
import numpy as np
import torch
from scipy.sparse import spmatrix
from collections import OrderedDict as OderedDictClass

from graphium.features.spectral import compute_laplacian_pe
from graphium.features.rw import compute_rwse
from graphium.features.electrostatic import compute_electrostatic_interactions
from graphium.features.commute import compute_commute_distances
from graphium.features.graphormer import compute_graphormer_distances
from graphium.features.transfer_pos_level import transfer_pos_level


def get_all_positional_encodings(
    adj: Union[np.ndarray, spmatrix],
    num_nodes: int,
    pos_kwargs: Optional[Dict] = None,
) -> Tuple["OrderedDict[str, np.ndarray]"]:
    r"""
    Get features positional encoding.

    Parameters:
        adj [num_nodes, num_nodes]: Adjacency matrix of the graph
        num_nodes: Number of nodes in the graph
        pos_encoding_as_features: keyword arguments for function `graph_positional_encoder`
            to generate positional encoding for node features.

    Returns:
        pe_dict: Dictionary of positional and structural encodings
    """

    pos_kwargs = {} if pos_kwargs is None else pos_kwargs

    pe_dict = OderedDictClass()

    # Initialize cache
    cache = {}

    # Get the positional encoding for the features
    if len(pos_kwargs) > 0:
        for pos_name, this_pos_kwargs in pos_kwargs["pos_types"].items():
            this_pos_kwargs = deepcopy(this_pos_kwargs)
            pos_type = this_pos_kwargs.pop("pos_type", None)
            pos_level = this_pos_kwargs.pop("pos_level", None)
            this_pe, cache = graph_positional_encoder(
                deepcopy(adj),
                num_nodes,
                pos_type=pos_type,
                pos_level=pos_level,
                pos_kwargs=this_pos_kwargs,
                cache=cache,
            )
            if pos_level == "node":
                pe_dict.update({f"{pos_type}": this_pe})
            else:
                pe_dict.update({f"{pos_level}_{pos_type}": this_pe})

    return pe_dict


def graph_positional_encoder(
    adj: Union[np.ndarray, spmatrix],
    num_nodes: int,
    pos_type: Optional[str] = None,
    pos_level: Optional[str] = None,
    pos_kwargs: Optional[Dict[str, Any]] = None,
    cache: Optional[Dict[str, Any]] = None,
) -> Tuple[Dict[str, np.ndarray], Dict[str, Any]]:
    r"""
    Get a positional encoding that depends on the parameters.

    Parameters:
        adj [num_nodes, num_nodes]: Adjacency matrix of the graph
        num_nodes: Number of nodes in the graph
        pos_type: The type of positional encoding to use. If None, it must be provided by `pos_kwargs["pos_type"]`. Supported types are:
            - laplacian_eigvec              \
            - laplacian_eigval               \  -> cache connected comps. & eigendecomp.
            - rwse
            - electrostatic                 \
            - commute                        \  -> cache pinvL
            - graphormer
        pos_level: Positional level to output. If None, it must be provided by `pos_kwargs["pos_level"]`.
            - node
            - edge
            - nodepair
            - graph
        pos_kwargs: Extra keyword arguments for the positional encoding. Can include the keys pos_type and pos_level.
        cache: Dictionary of cached objects

    Returns:
        pe: Positional or structural encoding
        cache: Updated dictionary of cached objects
    """

    pos_kwargs = deepcopy(pos_kwargs)
    if pos_kwargs is None:
        pos_kwargs = {}
    if cache is None:
        cache = {}

    # Get the positional type
    pos_type2 = pos_kwargs.pop("pos_type", None)
    if pos_type is None:
        pos_type = pos_type2
    if pos_type2 is not None:
        assert (
            pos_type == pos_type2
        ), f"The positional type must be the same in `pos_type` and `pos_kwargs['pos_type']`. Provided: {pos_type} and {pos_type2}"
    assert pos_type is not None, "Either `pos_type` or `pos_kwargs['pos_type']` must be provided."

    # Get the positional level
    pos_level2 = pos_kwargs.pop("pos_level", None)
    if pos_level is None:
        pos_level = pos_level2
    if pos_level2 is not None:
        assert (
            pos_level == pos_level2
        ), f"The positional level must be the same in `pos_level` and `pos_kwargs['pos_level']`. Provided: {pos_level} and {pos_level2}"
    assert pos_level is not None, "Either `pos_level` or `pos_kwargs['pos_level']` must be provided."

    # Convert to numpy array
    if isinstance(adj, torch.sparse.Tensor):
        adj = adj.to_dense().numpy()
    elif isinstance(adj, torch.Tensor):
        adj = adj.numpy()
    adj = adj.astype(np.float64)

    # Calculate positional encoding
    if pos_type == "laplacian_eigvec":
        _, pe, base_level, cache = compute_laplacian_pe(adj, cache=cache, **pos_kwargs)

    elif pos_type == "laplacian_eigval":
        pe, _, base_level, cache = compute_laplacian_pe(adj, cache=cache, **pos_kwargs)

    elif pos_type == "rw_return_probs":
        pe, base_level, cache = compute_rwse(
            adj.astype(np.float32), num_nodes=num_nodes, cache=cache, pos_type=pos_type, **pos_kwargs
        )

    elif pos_type == "rw_transition_probs":
        pe, base_level, cache = compute_rwse(
            adj.astype(np.float32), num_nodes=num_nodes, cache=cache, pos_type=pos_type, **pos_kwargs
        )

    elif pos_type == "electrostatic":
        pe, base_level, cache = compute_electrostatic_interactions(adj, cache, **pos_kwargs)

    elif pos_type == "commute":
        pe, base_level, cache = compute_commute_distances(adj, num_nodes, cache, **pos_kwargs)

    elif pos_type == "graphormer":
        pe, base_level, cache = compute_graphormer_distances(adj, num_nodes, cache, **pos_kwargs)

    else:
        raise ValueError(f"Unknown `pos_type`: {pos_type}")

    # Convert to float32 and Convert between different pos levels
    if isinstance(pe, (list, tuple)):
        pe = [this_pe.astype(np.float32) for this_pe in pe]
        pe = [transfer_pos_level(this_pe, base_level, pos_level, adj, num_nodes, cache) for this_pe in pe]
    else:
        pe = np.real(pe).astype(np.float32)
        pe = transfer_pos_level(pe, base_level, pos_level, adj, num_nodes, cache)

    return pe, cache


# graphium/features/properties.py

from typing import Union, List, Callable

import numpy as np
import datamol as dm

from rdkit.Chem import rdMolDescriptors as rdMD


def get_prop_or_none(
    prop: Callable, n: int, *args: Union[dm.Mol, str], **kwargs: Union[dm.Mol, str]
) -> Union[List[float], List[None]]:
    r"""
    return properties. If error, return list of `None` with lenght `n`.
    Parameters:
        prop: The property to compute.
        n: The number of elements in the property.
        *args: The arguments to pass to the property.
        **kwargs: The keyword arguments to pass to the property.
    Returns:
        The property or a list of `None` with lenght `n`.
    """
    try:
        return prop(*args, **kwargs)
    except RuntimeError:
        return [None] * n


def get_props_from_mol(
    mol: Union[dm.Mol, str],
    properties: Union[List[str], str] = "autocorr3d",
) -> np.ndarray:
    r"""
    Function to get a given set of desired properties from a molecule,
    and output a property list.

    Parameters:
        mol: The molecule from which to compute the properties.
        properties:
            The list of properties to compute for each molecule. It can be the following:

            - 'descriptors'
            - 'autocorr3d'
            - 'rdf'
            - 'morse'
            - 'whim'
            - 'all'

    Returns:
        props: np.array(float)
            The array of properties for the desired molecule
        classes_start_idx: list(int)
            The list of index specifying the start of each new class of
            descriptor or property. For example, if props has 20 elements,
            the first 5 are rotatable bonds, the next 8 are morse, and
            the rest are whim, then ``classes_start_idx = [0, 5, 13]``.
            This will mainly be useful to normalize the features of
            each class.
        classes_names: list(str)
            The name of the classes associated to each starting index.
            Will be usefull to understand what property is the network learning.

    """

    if isinstance(mol, str):
        mol = dm.to_mol(mol)

    if isinstance(properties, str):
        properties = [properties]

    properties = [p.lower() for p in properties]

    # Initialize arrays
    props = []  # Property vector for the features
    classes_start_idx = []  # The starting index for each property class
    classes_names = []

    # Generate a 3D structure for the molecule
    mol = dm.add_hs(mol)

    if ("autocorr3d" in properties) or ("all" in properties):
        # Some kind of 3D description of the molecule
        classes_names.append("autocorr3d")
        classes_start_idx.append(len(props))
        props.extend(get_prop_or_none(rdMD.CalcAUTOCORR3D, 80, mol))

    if ("rdf" in properties) or ("all" in properties):
        # The radial distribution function (better than the inertia)
        # https://en.wikipedia.org/wiki/Radial_distribution_function
        classes_names.append("rdf")
        classes_start_idx.append(len(props))
        props.extend(get_prop_or_none(rdMD.CalcRDF, 210, mol))

    if ("morse" in properties) or ("all" in properties):
        # Molecule Representation of Structures based on Electron diffraction descriptors
        classes_names.append("morse")
        classes_start_idx.append(len(props))
        props.extend(get_prop_or_none(rdMD.CalcMORSE, 224, mol))

    if ("whim" in properties) or ("all" in properties):
        # WHIM descriptors are 3D structural descriptors obtained from the
        # (x,y,z)‐atomic coordinates of a molecular conformation of a chemical,
        # and are used successfully in QSAR modelling.
        classes_names.append("whim")
        classes_start_idx.append(len(props))
        props.extend(get_prop_or_none(rdMD.CalcWHIM, 114, mol))

    return np.array(props), classes_start_idx, classes_names


# graphium/features/rw.py

from typing import Tuple, Union, Optional, List, Dict, Any, Iterable

from scipy.sparse import issparse, spmatrix, coo_matrix
import numpy as np
import torch

from torch_geometric.utils import to_dense_adj, from_scipy_sparse_matrix
from torch_scatter import scatter_add
from torch_geometric.utils.num_nodes import maybe_num_nodes


def compute_rwse(
    adj: Union[np.ndarray, spmatrix],
    ksteps: Union[int, List[int]],
    num_nodes: int,
    cache: Dict[str, Any],
    pos_type: str = "rw_return_probs" or "rw_transition_probs",
    space_dim: int = 0,
) -> Tuple[np.ndarray, str, Dict[str, Any]]:
    """
    Compute Random Walk Spectral Embedding (RWSE) for given list of K steps.

    Parameters:
        adj [num_nodes, num_nodes]: Adjacency matrix
        ksteps: List of numbers of steps for the random walks. If int, a list is generated from 1 to ksteps.
        num_nodes: Number of nodes in the graph
        cache: Dictionary of cached objects
        pos_type: Desired output
        space_dim: Estimated dimensionality of the space. Used to
            correct the random-walk diagonal by a factor `k^(space_dim/2)`.
            In euclidean space, this correction means that the height of
            the gaussian distribution stays almost constant across the number of
            steps, if `space_dim` is the dimension of the euclidean space.
    Returns:
        Two possible outputs:
            rw_return_probs [num_nodes, len(ksteps)]: Random-Walk k-step landing probabilities
            rw_transition_probs [num_nodes, num_nodes, len(ksteps)]:  Random-Walk k-step transition probabilities
        base_level: Indicator of the output pos_level (node, edge, nodepair, graph) -> here either node or nodepair
        cache: Updated dictionary of cached objects
    """

    base_level = "node" if pos_type == "rw_return_probs" else "nodepair"

    # Manually handles edge case of 1 atom molecules here
    if not isinstance(ksteps, Iterable):
        ksteps = list(range(1, ksteps + 1))
    if num_nodes == 1:
        if pos_type == "rw_return_probs":
            return np.ones((1, len(ksteps))), base_level, cache
        else:
            return np.ones((1, 1, len(ksteps))), base_level, cache

    # Get the edge indices from the adjacency matrix
    if not issparse(adj):
        if "coo_adj" in cache:
            adj = cache["coo_adj"]
        elif "csr_adj" in cache:
            adj = cache["csr_adj"]
        else:
            adj = coo_matrix(adj, dtype=np.float64)
            cache["coo_adj"] = adj

    edge_index, edge_weight = from_scipy_sparse_matrix(adj)

    # Compute the random-walk transition probabilities
    if "ksteps" in cache:
        cached_k = cache["ksteps"]
        missing_k = [k for k in ksteps if k not in cached_k]
        if missing_k == []:
            pass
        elif min(missing_k) < min(cached_k):
            Pk_dict = get_Pks(missing_k, edge_index=edge_index, edge_weight=edge_weight, num_nodes=num_nodes)
            cache["ksteps"] = sorted(missing_k + cache["ksteps"])
            for k in missing_k:
                cache["Pk"][k] = Pk_dict[k]
        else:
            start_k = min([max(cached_k), min(missing_k)])
            start_Pk = cache["Pk"][start_k]
            Pk_dict = get_Pks(
                missing_k,
                edge_index=edge_index,
                edge_weight=edge_weight,
                num_nodes=num_nodes,
                start_Pk=start_Pk,
                start_k=start_k,
            )
            cache["ksteps"] = sorted(cache["ksteps"] + missing_k)
            for k in missing_k:
                cache["Pk"][k] = Pk_dict[k]
    else:
        Pk_dict = get_Pks(ksteps, edge_index=edge_index, edge_weight=edge_weight, num_nodes=num_nodes)

        cache["ksteps"] = list(Pk_dict.keys())
        cache["Pk"] = Pk_dict

    pe_list = []
    if pos_type == "rw_return_probs":
        for k in ksteps:
            pe_list.append(torch.diagonal(cache["Pk"][k], dim1=-2, dim2=-1) * (k ** (space_dim / 2)))
    else:
        for k in ksteps:
            pe_list.append(cache["Pk"][k])

    pe = torch.stack(pe_list, dim=-1).numpy()

    return pe, base_level, cache


def get_Pks(
    ksteps: List[int],
    edge_index: Tuple[torch.Tensor, torch.Tensor],
    edge_weight: Optional[torch.Tensor] = None,
    num_nodes: Optional[int] = None,
    start_Pk: Optional[torch.Tensor] = None,
    start_k: Optional[int] = None,
) -> Dict[int, np.ndarray]:
    """
    Compute Random Walk landing probabilities for given list of K steps.

    Parameters:
        ksteps: List of numbers of k-steps for which to compute the RW landings
        edge_index: PyG sparse representation of the graph
        edge_weight: Edge weights
        num_nodes: Number of nodes in the graph

    Returns:
        2D Tensor with shape (num_nodes, len(ksteps)) with RW landing probs
    """
    if edge_weight is None:
        edge_weight = torch.ones(edge_index.size(1), device=edge_index.device)
    num_nodes = maybe_num_nodes(edge_index, num_nodes)
    src = edge_index[0]
    deg = scatter_add(edge_weight, src, dim=0, dim_size=num_nodes)  # Out degrees.
    deg_inv = deg.pow(-1.0)
    deg_inv.masked_fill_(deg_inv == float("inf"), 0)

    if edge_index.numel() == 0:
        P = edge_index.new_zeros((1, num_nodes, num_nodes))
    else:
        # P = D^-1 * A
        P = torch.diag(deg_inv).float() @ to_dense_adj(
            edge_index, max_num_nodes=num_nodes
        )  # 1 x (Num nodes) x (Num nodes)

    if start_Pk is not None:
        Pk = start_Pk @ P.clone().detach().matrix_power(min(ksteps) - start_k)
    else:
        Pk = P.clone().detach().matrix_power(min(ksteps))

    Pk_dict = {}
    for k in range(min(ksteps), max(ksteps) + 1):
        Pk_dict[k] = Pk.squeeze(0)
        Pk = Pk @ P

    return Pk_dict


# graphium/features/spectral.py

from typing import Tuple, Union, Dict, Any
from scipy.linalg import eig
from scipy.sparse import csr_matrix, diags, issparse, spmatrix
import numpy as np
import torch
import networkx as nx

from graphium.utils.tensor import is_dtype_torch_tensor, is_dtype_numpy_array


def compute_laplacian_pe(
    adj: Union[np.ndarray, spmatrix],
    num_pos: int,
    cache: Dict[str, Any],
    disconnected_comp: bool = True,
    normalization: str = "none",
) -> Tuple[np.ndarray, str, Dict[str, Any]]:
    r"""
    Compute the Laplacian eigenvalues and eigenvectors of the Laplacian of the graph.

    Parameters:
        adj [num_nodes, num_nodes]: Adjacency matrix of the graph
        num_pos: Number of Laplacian eigenvectors to compute
        cache: Dictionary of cached objects
        disconnected_comp: Whether to compute the eigenvectors for each connected component
        normalization: Normalization to apply to the Laplacian

    Returns:
        Two possible outputs:
            eigvals [num_nodes, num_pos]: Eigenvalues of the Laplacian repeated for each node.
                This repetition is necessary in case of disconnected components, where
                the eigenvalues of the Laplacian are not the same for each node.
            eigvecs [num_nodes, num_pos]: Eigenvectors of the Laplacian
        base_level: Indicator of the output pos_level (node, edge, nodepair, graph) -> here node
        cache: Updated dictionary of cached objects
    """

    base_level = "node"

    # Sparsify the adjacency patrix
    if not issparse(adj):
        if "csr_adj" not in cache:
            adj = csr_matrix(adj, dtype=np.float64)
            cache["csr_adj"] = adj
        else:
            adj = cache["csr_adj"]

    # Compute the Laplacian, and normalize it
    if f"L_{normalization}_sp" not in cache:
        D = np.array(np.sum(adj, axis=1)).flatten()
        D_mat = diags(D)
        L = -adj + D_mat
        L_norm = normalize_matrix(L, degree_vector=D, normalization=normalization)
        cache[f"L_{normalization}_sp"] = L_norm
    else:
        L_norm = cache[f"L_{normalization}_sp"]

    components = []

    if disconnected_comp:
        if "components" not in cache:
            # Get the list of connected components
            components = list(nx.connected_components(nx.from_scipy_sparse_array(adj)))
            cache["components"] = components

        else:
            components = cache["components"]

    # Compute the eigenvectors for each connected component, and stack them together
    if len(components) > 1:
        if "lap_eig_comp" not in cache:
            eigvals = np.zeros((adj.shape[0], num_pos), dtype=np.complex64)
            eigvecs = np.zeros((adj.shape[0], num_pos), dtype=np.complex64)
            for component in components:
                comp = list(component)
                this_L = L_norm[comp][:, comp]
                this_eigvals, this_eigvecs = _get_positional_eigvecs(this_L, num_pos=num_pos)

                # Eigenvalues previously set to infinity are now set to 0
                # Any NaN in the eigvals or eigvecs will be set to 0
                this_eigvecs[~np.isfinite(this_eigvecs)] = 0.0
                this_eigvals[~np.isfinite(this_eigvals)] = 0.0

                eigvals[comp, :] = np.expand_dims(this_eigvals, axis=0)
                eigvecs[comp, :] = this_eigvecs
            cache["lap_eig_comp"] = (eigvals, eigvecs)

        else:
            eigvals, eigvecs = cache["lap_eig_comp"]

    else:
        if "lap_eig" not in cache:
            eigvals, eigvecs = _get_positional_eigvecs(L, num_pos=num_pos)

            # Eigenvalues previously set to infinity are now set to 0
            # Any NaN in the eigvals or eigvecs will be set to 0
            eigvecs[~np.isfinite(eigvecs)] = 0.0
            eigvals[~np.isfinite(eigvals)] = 0.0
            eigvals = np.repeat(np.expand_dims(eigvals, axis=0), adj.shape[0], axis=0)

            cache["lap_eig"] = (eigvals, eigvecs)

        else:
            eigvals, eigvecs = cache["lap_eig"]

    return eigvals, eigvecs, base_level, cache


def _get_positional_eigvecs(
    matrix: Union[np.ndarray, spmatrix],
    num_pos: int,
) -> Tuple[np.ndarray, np.ndarray]:
    r"""
    compute the eigenvalues and eigenvectors of a matrix
    Parameters:
        matrix: Matrix to compute the eigenvalues and eigenvectors of
        num_pos: Number of eigenvalues and eigenvectors to compute
    Returns:
        eigvals: Eigenvalues of the matrix
        eigvecs: Eigenvectors of the matrix
    """
    mat_len = matrix.shape[0]
    eigvals, eigvecs = eig(matrix.todense())

    # Pad with non-sense eigenvectors if required
    if num_pos > mat_len:
        temp_EigVal = np.ones(num_pos - mat_len, dtype=np.float64) + float("inf")
        temp_EigVec = np.zeros((mat_len, num_pos - mat_len), dtype=np.float64)
        eigvals = np.concatenate([eigvals, temp_EigVal], axis=0)
        eigvecs = np.concatenate([eigvecs, temp_EigVec], axis=1)

    # Sort and keep only the first `num_pos` elements
    sort_idx = eigvals.argsort()
    eigvals = eigvals[sort_idx]
    eigvals = eigvals[:num_pos]
    eigvecs = eigvecs[:, sort_idx]
    eigvecs = eigvecs[:, :num_pos]

    # Normalize the eigvecs
    eigvecs = eigvecs / np.maximum(np.sqrt(np.sum(eigvecs**2, axis=0, keepdims=True)), 1e-4)

    return eigvals, eigvecs


def normalize_matrix(
    matrix: Union[np.ndarray, spmatrix],
    degree_vector=None,
    normalization: str = None,
) -> Union[np.ndarray, spmatrix]:
    r"""
    Normalize a given matrix using its degree vector

    Parameters
    ---------------

        matrix: torch.tensor(N, N) or scipy.sparse.spmatrix(N, N)
            A square matrix representing either an Adjacency matrix or a Laplacian.

        degree_vector: torch.tensor(N) or np.ndarray(N) or None
            A vector representing the degree of ``matrix``.
            ``None`` is only accepted if ``normalization==None``

        normalization: str or None, Default='none'
            Normalization to use on the eig_matrix

            - 'none' or ``None``: no normalization

            - 'sym': Symmetric normalization ``D^-0.5 L D^-0.5``

            - 'inv': Inverse normalization ``D^-1 L``

    Returns
    -----------
        matrix: torch.tensor(N, N) or scipy.sparse.spmatrix(N, N)
            The normalized matrix

    """

    # Transform the degree vector into a matrix
    if degree_vector is None:
        if not ((normalization is None) or (normalization.lower() == "none")):
            raise ValueError("`degree_vector` cannot be `None` if `normalization` is not `None`")
    else:
        if is_dtype_numpy_array(matrix.dtype):
            with np.errstate(divide="ignore", invalid="ignore"):
                degree_inv = np.expand_dims(degree_vector**-0.5, axis=1)
                degree_inv[np.isinf(degree_inv)] = 0
        elif is_dtype_torch_tensor(matrix.dtype):
            degree_inv = torch.unsqueeze(degree_vector**-0.5, dim=1)
            degree_inv[torch.isinf(degree_inv)] = 0

    # Compute the normalized matrix
    if (normalization is None) or (normalization.lower() == "none"):
        pass
    elif normalization.lower() == "sym":
        matrix = degree_inv * matrix * degree_inv.T
    elif normalization.lower() == "inv":
        matrix = (degree_inv**2) * matrix
    else:
        raise ValueError(
            f'`normalization` should be `None`, `"None"`, `"sym"` or `"inv"`, but `{normalization}` was provided'
        )

    return matrix


# graphium/features/transfer_pos_level.py

from typing import Tuple, Union, List, Dict, Any, Optional

import numpy as np

from scipy.sparse import spmatrix, issparse, coo_matrix

from torch_geometric.utils import from_scipy_sparse_matrix


def transfer_pos_level(
    pe: np.ndarray,
    in_level: str,
    out_level: str,
    adj: Union[np.ndarray, spmatrix],
    num_nodes: int,
    cache: Optional[Dict[str, Any]] = None,
) -> np.ndarray:
    r"""
    Transfer positional encoding between different positional levels (node, edge, nodepair, graph)

    Parameters:
        pe: Input pe with pos_level defined by in_level
        in_level: pos_level of input pe
        out_level: desired pos_level of output pe
        adj [num_nodes, num_nodes]: Adjacency matrix of the graph
        num_nodes: Number of nodes in the graph
        cache: Dictionary of cached objects

    Returns:
        pe: Output pe with pos_level defined by out_level
    """

    if cache is None:
        cache = {}

    if in_level == "node":
        if out_level == "node":
            pass

        elif out_level == "edge":
            pe, cache = node_to_edge(pe, adj, cache)

        elif out_level == "nodepair":
            pe = node_to_nodepair(pe, num_nodes)

        elif out_level == "graph":
            raise NotImplementedError("Transfer function (node -> graph) not yet implemented.")

        else:
            raise ValueError(f"Unknown `pos_level`: {out_level}")

    elif in_level == "edge":
        raise NotImplementedError("Transfer function (edge -> *) not yet implemented.")

    elif in_level == "nodepair":
        if len(pe.shape) == 2:
            pe = np.expand_dims(pe, -1)

        if out_level == "node":
            pe = nodepair_to_node(pe)

        elif out_level == "edge":
            pe, cache = nodepair_to_edge(pe, adj, cache)

        elif out_level == "nodepair":
            pass

        elif out_level == "graph":
            raise NotImplementedError("Transfer function (nodepair -> graph) not yet implemented.")

        else:
            raise ValueError(f"Unknown `pos_level`: {out_level}")

    elif in_level == "graph":
        if out_level == "node":
            pe = graph_to_node(pe, num_nodes, cache)

        elif out_level in ["edge", "nodepair"]:
            raise NotImplementedError("Transfer function (graph -> edge/nodepair) not yet implemented.")

        else:
            raise ValueError(f"Unknown `pos_level`: {out_level}")

    else:
        raise ValueError(f"Unknown `pos_level`: {in_level}")

    return pe


# Transfer functions between different levels, i.e., node, edge, nodepair and graph level.

# TODO:
#   - Implement missing transfer functions below
#   - Are transfer functions graph -> edge/nodepair and edge -> graph needed?


def node_to_edge(
    pe: np.ndarray, adj: Union[np.ndarray, spmatrix], cache: Optional[Dict[str, Any]] = None
) -> Tuple[np.ndarray, Dict[str, Any]]:
    r"""
    Get an edge-level positional encoding from a node-level positional encoding.
     -> For each edge, concatenate the sum and absolute difference of pe of source and destination node.

    Parameters:
        pe [num_nodes, num_feat]: Node-level positional encoding
        adj [num_nodes, num_nodes]: Adjacency matrix of the graph
        cache: Dictionary of cached objects

    Returns:
        edge_pe [2 * num_edges, 2 * num_feat]: Edge-level positional encoding
        cache: Updated dictionary of cached objects
    """

    if cache is None:
        cache = {}

    if not issparse(adj):
        if "coo_adj" in cache:
            adj = cache["coo_adj"]
        elif "csr_adj" in cache:
            adj = cache["csr_adj"]
        else:
            adj = coo_matrix(adj, dtype=np.float64)
            cache["coo_adj"] = adj

    edge_index, _ = from_scipy_sparse_matrix(adj)
    src, dst = edge_index[0], edge_index[1]

    pe_sum = pe[src] + pe[dst]
    pe_abs_diff = np.abs(pe[src] - pe[dst])

    edge_pe = np.concatenate((pe_sum, pe_abs_diff), axis=-1)

    return edge_pe, cache


def node_to_nodepair(pe: np.ndarray, num_nodes: int) -> np.ndarray:
    r"""
    Get a nodepair-level positional encoding from a node-level positional encoding.
     -> For each nodepair (i,j) concatenate the sum and absolute difference of pe at node i and j.

    Parameters:
        pe [num_nodes, num_feat]: Node-level positional encoding
        num_nodes: Number of nodes in the graph

    Returns:
        nodepair_pe [num_nodes, num_nodes, 2 * num_feat]: Nodepair-level positional encoding
    """

    expanded_pe = np.expand_dims(pe, axis=1)
    expanded_pe = np.repeat(expanded_pe, repeats=num_nodes, axis=1)

    pe_sum = expanded_pe + expanded_pe.transpose([1, 0, 2])
    pe_abs_diff = np.abs(expanded_pe - expanded_pe.transpose([1, 0, 2]))

    nodepair_pe = np.concatenate((pe_sum, pe_abs_diff), axis=-1)

    return nodepair_pe


def node_to_graph(pe: np.ndarray, num_nodes: int) -> np.ndarray:
    r"""
    Get a graph-level positional encoding from a node-level positional encoding.
     -> E.g., min/max/mean-pooling of node features.

    Parameters:
        pe [num_nodes, num_feat]: Node-level positional encoding
        num_nodes: Number of nodes in the graph

    Returns:
        graph_pe [1, num_feat]: Graph-level positional encoding
    """

    raise NotImplementedError("Transfer function (node -> graph) not yet implemented.")


def edge_to_node(pe: np.ndarray, adj: Union[np.ndarray, spmatrix]) -> np.ndarray:
    r"""
    Get a node-level positional encoding from an edge-level positional encoding.
     -> E.g., min/max/mean-pooling of information from edges (i,j) that contain node i

    Parameters:
        pe [num_edges, num_feat]: Edge-level positional encoding
        adj [num_nodes, num_nodes]: Adjacency matrix of the graph

    Returns:
        node_pe [num_edges, num_feat]: Node-level positional encoding
    """

    raise NotImplementedError("Transfer function (edge -> node) not yet implemented.")


def edge_to_nodepair(
    pe: np.ndarray, adj: Union[np.ndarray, spmatrix], num_nodes: int, cache: Optional[Dict[str, Any]] = None
) -> np.ndarray:
    r"""
    Get a nodepair-level positional encoding from an edge-level positional encoding.
     -> Zero-padding of non-existing edges.

    Parameters:
        pe [num_edges, num_feat]: Edge-level positional encoding
        adj [num_nodes, num_nodes]: Adjacency matrix of the graph
        num_nodes: Number of nodes in the graph
        cache: Dictionary of cached objects

    Returns:
        nodepair_pe [num_edges, num_edges, num_feat]: Nodepair-level positional encoding
        cache: Updated dictionary of cached objects
    """

    if cache is None:
        cache = {}

    num_feat = pe.shape[-1]

    if not isinstance(adj, coo_matrix):
        if "coo_adj" in cache:
            adj = cache["coo_adj"]
        else:
            adj = coo_matrix(adj, dtype=np.float64)
        cache["coo_adj"] = adj

    dst, src = adj.row, adj.col

    nodepair_pe = np.zeros((num_nodes, num_nodes, num_feat))

    for i in range(len(dst)):
        nodepair_pe[dst[i], src[i], ...] = pe[i, ...]

    return nodepair_pe, cache


def edge_to_graph(pe: np.ndarray) -> np.ndarray:
    r"""
    Get a graph-level positional encoding from an edge-level positional encoding.

    Parameters:
        pe [num_edges, num_feat]: Edge-level positional encoding

    Returns:
        graph_pe [1, num_feat]: Graph-level positional encoding
    """

    raise NotImplementedError("Transfer function (edge -> graph) not yet implemented.")


def nodepair_to_node(pe: np.ndarray, stats_list: List = [np.min, np.mean, np.std]) -> np.ndarray:
    r"""
    Get a node-level positional encoding from a graph-level positional encoding.
     -> Calculate statistics over rows & cols of input positional encoding

    Parameters:
        pe [num_nodes, num_nodes, num_feat]: Nodepair-level positional encoding
        stats_list: List of statistics to calculate per row/col of nodepair-level pe

    Returns:
        node_pe [num_nodes, 2 * len(stats_list) * num_feat]: Node-level positional encoding
    """

    num_feat = pe.shape[-1]

    node_pe_list = []

    for stat in stats_list:
        for i in range(num_feat):
            node_pe_list.append(stat(pe[..., i], axis=0))
            node_pe_list.append(stat(pe[..., i], axis=1))
    node_pe = np.stack(node_pe_list, axis=-1)

    return node_pe


def nodepair_to_edge(
    pe: np.ndarray, adj: Union[np.ndarray, spmatrix], cache: Optional[Dict[str, Any]] = None
) -> np.ndarray:
    r"""
    Get a edge-level positional encoding from a nodepair-level positional encoding.
     -> Mask and sparsify nodepair-level positional encoding

    Parameters:
        pe [num_nodes, num_nodes, num_feat]: Nodepair-level positional encoding
        adj [num_nodes, num_nodes]: Adjacency matrix of the graph
        cache: Dictionary of cached objects

    Returns:
        edge_pe [num_edges, num_feat]: Edge-level positional encoding
        cache: Updated dictionary of cached objects
    """

    if cache is None:
        cache = {}

    num_feat = pe.shape[-1]

    if not isinstance(adj, coo_matrix):
        if "coo_adj" in cache:
            adj = cache["coo_adj"]
        else:
            adj = coo_matrix(adj, dtype=np.float64)
        cache["coo_adj"] = adj

    dst, src = adj.row, adj.col

    edge_pe = np.zeros((len(dst), num_feat))

    for i in range(len(src)):
        edge_pe[i, ...] = pe[dst[i], src[i]]

    return edge_pe, cache


def nodepair_to_graph(pe: np.ndarray, num_nodes: int) -> np.ndarray:
    r"""
    Get a graph-level positional encoding from a nodepair-level positional encoding.
     -> E.g., min/max/mean-pooling of entries of input pe

    Parameters:
        pe [num_nodes, num_nodes, num_feat]: Nodepair-level positional encoding
        num_nodes: Number of nodes in the graph

    Returns:
        graph_pe [1, num_feat]: Graph-level positional encoding
    """

    raise NotImplementedError("Transfer function (nodepair -> graph) not yet implemented.")


def graph_to_node(
    pe: Union[np.ndarray, List], num_nodes: int, cache: Optional[Dict[str, Any]] = None
) -> np.ndarray:
    r"""
    Get a node-level positional encoding from a nodepair-level positional encoding.
     -> E.g., expand dimension of graph-level pe

    Parameters:
        pe [num_feat]: Nodepair-level positional encoding (or list of them if graph disconnected)
        num_nodes: Number of nodes in the graph
        cache: Dictionary of cached objects

    Returns:
        node_pe [num_nodes, num_feat]: Node-level positional encoding
    """

    if cache is None:
        cache = {}

    node_pe = None

    # The key 'components' is only in cache if disconnected_comp == True when computing base pe
    if "components" in cache:
        if len(cache["components"]) > 1:
            node_pe = np.zeros((num_nodes, len(pe)))
            components = cache["components"]

            for i, component in enumerate(components):
                comp = list(component)
                node_pe[comp, :] = np.real(pe[i])

    if node_pe is None:
        node_pe = np.tile(pe, (num_nodes, 1))

    return node_pe


# graphium/ipu/ipu_dataloader.py

# Copyright (c) 2022 Graphcore Ltd. All rights reserved.

from typing import Callable, Iterable, Optional, List, Tuple, Dict, Any, Union
from copy import deepcopy
from dataclasses import dataclass
import numpy as np
from loguru import logger
from torch import Tensor

import torch
from torch_geometric.data import Data, Batch, Dataset
from torch_geometric.transforms import BaseTransform

from graphium.data.utils import get_keys
from graphium.ipu.ipu_utils import import_poptorch
from graphium.utils.packing import (
    fast_packing,
    hybrid_packing,
    get_pack_sizes,
    node_to_pack_indices_mask,
    estimate_max_pack_node_size,
)


@dataclass
class IPUDataloaderOptions:
    r"""
    This data class stores the arguments necessary to instantiate a model for the Predictor.

    Parameters:
        model_class:
            pytorch module used to create a model

        model_kwargs:
            Key-word arguments used to initialize the model from `model_class`.
    """

    batch_size: int
    max_num_nodes: Optional[int] = None
    max_num_nodes_per_graph: Optional[int] = None
    max_num_edges: Optional[int] = None
    max_num_edges_per_graph: Optional[int] = None
    mode: "poptorch.DataLoaderMode" = "Sync"

    def set_kwargs(self):
        # Get the maximum number of nodes
        if self.max_num_nodes is not None:
            assert (
                self.max_num_nodes_per_graph is None
            ), "Cannot use `max_num_nodes` and `max_num_nodes_per_graph` simultaneously"
        elif self.max_num_nodes_per_graph is not None:
            assert (
                self.max_num_nodes is None
            ), "Cannot use `max_num_nodes` and `max_num_nodes_per_graph` simultaneously"
            self.max_num_nodes = self.max_num_nodes_per_graph * self.batch_size
        else:
            raise ValueError("Must provide either `max_num_nodes` or `max_num_nodes_per_graph`")

        # Get the maximum number of edges
        if self.max_num_edges is not None:
            assert (
                self.max_num_edges_per_graph is None
            ), "Cannot use `max_num_edges` and `max_num_edges_per_graph` simultaneously"
        elif self.max_num_edges_per_graph is not None:
            assert (
                self.max_num_edges is None
            ), "Cannot use `max_num_edges` and `max_num_edges_per_graph` simultaneously"
            self.max_num_edges = self.max_num_edges_per_graph * self.batch_size
        else:
            raise ValueError("Must provide either `max_num_nodes` or `max_num_nodes_per_graph`")

        # poptorch mode
        poptorch = import_poptorch()
        if isinstance(self.mode, str):
            if self.mode.lower() == "sync":
                self.mode = poptorch.DataLoaderMode.Sync
            elif self.mode.lower() == "async":
                self.mode = poptorch.DataLoaderMode.Async
            elif self.mode.lower() == "asyncrebatched":
                self.mode = poptorch.DataLoaderMode.AsyncRebatched
            else:
                raise ValueError(f"`{self.mode}` not a valid parameter.")


class CombinedBatchingCollator:
    """
    Collator object that manages the combined batch size defined as:

        combined_batch_size = batch_size * device_iterations
                             * replication_factor * gradient_accumulation

    This is intended to be used in combination with the poptorch.DataLoader
    """

    def __init__(
        self,
        batch_size: int,
        max_num_nodes: int,
        max_num_edges: int,
        dataset_max_nodes_per_graph: int,
        dataset_max_edges_per_graph: int,
        collate_fn: Optional[Callable] = None,
    ):
        """
        Parameters:
            batch_size: mini batch size used by the model
            max_num_nodes: Maximum number of nodes in the batched padded graph
            max_num_edges: Maximum number of edges in the batched padded graph
            dataset_max_nodes_per_graph: Maximum number of nodes per graph in the full dataset
            dataset_max_edges_per_graph: Maximum number of edges per graph in the full dataset
            collate_fn: Function used to collate (or batch) the single data or graphs together
        """
        super().__init__()
        self.batch_size = batch_size
        self.collate_fn = collate_fn
        self.max_num_nodes = max_num_nodes
        self.max_num_edges = max_num_edges
        self.dataset_max_nodes_per_graph = dataset_max_nodes_per_graph
        self.dataset_max_edges_per_graph = dataset_max_edges_per_graph

    def __call__(
        self, batch: List[Dict[str, Union[Data, Dict[str, Tensor]]]]
    ) -> Dict[str, Union[Batch, Dict[str, Tensor], Any]]:
        """
        Stack tensors, batch the pyg graphs, and pad each tensor to be same size.

        Parameters:
            batch: The batch of data, including pyg-graphs `Data` and labels `Dict[str, Tensor]` to be padded

        Returns:
            out_batch: A dictionary where the graphs are batched and the labels or other Tensors are stacked
        """

        # Sort the batch such that large graphs are paired with small graphs
        num_nodes = [b["features"].num_nodes for b in batch]
        packed_indices = hybrid_packing(num_nodes, batch_size=self.batch_size)
        packs = [[batch[idx] for idx in pack] for pack in packed_indices]

        # Loop all mini-batches within the global batch
        all_batches = []
        for pack in packs:
            if self.collate_fn != None:
                local_batch = self.collate_fn(pack)

            transform = Pad(
                max_num_nodes=self.max_num_nodes,
                max_num_edges=self.max_num_edges,
                dataset_max_nodes_per_graph=self.dataset_max_nodes_per_graph,
                dataset_max_edges_per_graph=self.dataset_max_edges_per_graph,
            )

            local_batch["features"] = transform(local_batch["features"])
            local_batch["labels"] = transform(local_batch["labels"])
            all_batches.append(local_batch)

        out_batch = {}

        # Stack tensors in the first dimension to allow IPUs to differentiate between local and global graph
        all_keys = get_keys(all_batches[0]["labels"])
        out_batch["labels"] = {
            key: torch.stack([this_batch["labels"][key] for this_batch in all_batches], 0) for key in all_keys
        }
        out_graphs = [this_batch["features"] for this_batch in all_batches]
        stacked_features = deepcopy(out_graphs[0])
        for key, val in out_graphs[0].items():
            if isinstance(val, torch.Tensor):
                stacked_features[key] = torch.stack([this_graph[key] for this_graph in out_graphs], dim=0)

        out_batch["features"] = stacked_features
        for key in all_batches[0].keys():
            if key not in ("features", "labels"):
                out_batch[key] = [this_batch[key] for this_batch in all_batches]

        #
        for data_key, data_val in out_batch.items():
            if isinstance(data_val, Batch):
                for sub_key, sub_val in data_val.items():
                    if isinstance(sub_val, Tensor) and sub_val.dtype == torch.int64:
                        out_batch[data_key][sub_key] = sub_val.to(torch.int32)

        return out_batch


def create_ipu_dataloader(
    dataset: Dataset,
    ipu_dataloader_options: IPUDataloaderOptions,
    ipu_options: Optional["poptorch.Options"] = None,
    batch_size: Optional[int] = 1,
    collate_fn=None,
    num_workers: Optional[int] = 0,
    **kwargs,
) -> "poptorch.DataLoader":
    """
    Creates a poptorch.DataLoader for graph datasets
    Applies the mini-batching method of concatenating multiple graphs into a
    single graph with multiple disconnected subgraphs. See:
    https://pytorch-geometric.readthedocs.io/en/2.0.2/notes/batching.html

    Parameters:

        dataset: The torch_geometric.data.Dataset instance from which to
            load the graph examples for the IPU.
        ipu_dataloader_options: The options to initialize the Dataloader for IPU
        ipu_options: The poptorch.Options used by the
            poptorch.DataLoader. Will use the default options if not provided.
        batch_size: How many graph examples to load in each batch
            (default: 1).
        collate_fn: The function used to collate batches
        **kwargs (optional): Additional arguments of :class:`poptorch.DataLoader`.

    Returns:
        The dataloader
    """
    poptorch = import_poptorch()

    if ipu_options is None:
        # Create IPU default options
        ipu_options = poptorch.Options()

    # Define the collater function
    collater = CombinedBatchingCollator(
        batch_size,
        collate_fn=collate_fn,
        max_num_nodes=ipu_dataloader_options.max_num_nodes,
        max_num_edges=ipu_dataloader_options.max_num_edges,
        dataset_max_nodes_per_graph=dataset.max_num_nodes_per_graph,
        dataset_max_edges_per_graph=dataset.max_num_edges_per_graph,
    )

    # Get the global batch size
    num_nodes = np.asarray(dataset.num_nodes_list)
    accum = ipu_options.Training.gradient_accumulation
    repli = ipu_options._values["replication_factor"]
    device_iter = ipu_options._values["device_iterations"]
    combined_batch_size = batch_size * accum * repli * device_iter
    num_batches = len(dataset) // combined_batch_size
    num_workers = min(num_batches, num_workers)
    buffer_size = num_batches // num_workers if num_workers > 0 else None
    buffer_size = 3 if buffer_size is None else buffer_size
    async_options = {
        "sharing_strategy": poptorch.SharingStrategy.ForkServer,
        "early_preload": True,
        "buffer_size": buffer_size,
        "load_indefinitely": True,
        "miss_sleep_time_in_ms": 0,
    }

    # Estimate the packing size needed
    max_pack_size, max_pack_size_per_graph = 0, 0
    for _ in range(4):
        this_max_pack_size, this_max_pack_size_per_graph = estimate_max_pack_node_size(
            num_nodes=num_nodes,
            batch_size=batch_size,
            combined_batch_size=combined_batch_size,
        )
        max_pack_size = max(max_pack_size, this_max_pack_size)
        max_pack_size_per_graph = max(max_pack_size_per_graph, this_max_pack_size_per_graph)

    max_num_nodes = collater.max_num_nodes
    # Log the estimated pack size, with warnings if too big or too small
    logger.info(
        f"Estimating pack max_pack_size={max_pack_size} or max_pack_size_per_graph={max_pack_size_per_graph}"
    )
    logger.info(f"Provided `max_num_nodes={max_num_nodes}`")
    if max_pack_size > max_num_nodes - 10:
        logger.warning(
            f"The value of `max_num_nodes={max_num_nodes}` seems to be insufficient compared to `max_pack_size={max_pack_size}` and will likely crash"
        )
    elif max_pack_size < max_num_nodes - 20:
        logger.warning(
            f"The value of `max_num_nodes={max_num_nodes}` seems to be large compared to `max_pack_size={max_pack_size}` and will likely waste memory"
        )

    return poptorch.DataLoader(
        options=deepcopy(ipu_options),
        dataset=dataset,
        batch_size=batch_size,
        num_workers=num_workers,
        collate_fn=collater,
        async_options=async_options,
        **kwargs,
    )


class Pad(BaseTransform):
    """
    Data transform that applies padding to enforce consistent tensor shapes.
    """

    def __init__(
        self,
        max_num_nodes: int,
        dataset_max_nodes_per_graph,
        dataset_max_edges_per_graph,
        max_num_edges: Optional[int] = None,
        node_value: float = 0,
        edge_value: float = 0,
    ):
        """
        Parameters:
            max_num_nodes: The maximum number of nodes for the total padded graph
            dataset_max_nodes_per_graph: the maximum number of nodes per graph in the dataset
            dataset_max_edges_per_graph: the maximum number of edges per graph in the dataset
            max_num_edges: The maximum number of edges for the total padded graph
            node_value: Value to add to the node padding
            edge_value: Value to add to the edge padding
        """
        super().__init__()
        self.max_num_nodes = max_num_nodes
        self.dataset_max_nodes_per_graph = dataset_max_nodes_per_graph
        self.dataset_max_edges_per_graph = dataset_max_edges_per_graph

        if max_num_edges:
            self.max_num_edges = max_num_edges
        else:
            # Assume fully connected graph
            self.max_num_edges = max_num_nodes * (max_num_nodes - 1)

        self.node_value = node_value
        self.edge_value = edge_value

    def validate(self, data):
        """
        Validates that the input graph does not exceed the constraints that:

          * the number of nodes must be <= max_num_nodes
          * the number of edges must be <= max_num_edges

        Returns:
            Tuple containing the number nodes and the number of edges
        """
        num_nodes = data.num_nodes
        num_edges = data.num_edges

        assert num_nodes <= self.max_num_nodes, (
            f"Too many nodes. Graph has {num_nodes} nodes " f"and max_num_nodes is {self.max_num_nodes}."
        )

        assert num_edges <= self.max_num_edges, (
            f"Too many edges. Graph has {num_edges} edges defined "
            f"and max_num_edges is {self.max_num_edges}."
        )

        return num_nodes, num_edges

    def __call__(self, batch: Batch) -> Batch:
        return self._call(batch)

    def forward(self, batch: Batch) -> Batch:
        return self._call(batch)

    def _call(self, batch: Batch) -> Batch:
        """
        Pad the batch with a fake graphs that has the desired
        number of nodes and edges.
        """
        num_nodes, num_edges = self.validate(batch)
        num_pad_nodes = self.max_num_nodes - num_nodes
        num_pad_edges = self.max_num_edges - num_edges
        # Create a copy to update with padded features
        new_batch = deepcopy(batch)

        real_graphs = new_batch.to_data_list()

        for g in real_graphs:
            g.graph_is_true = torch.tensor([1], dtype=bool)
            g.node_is_true = torch.full([g.num_nodes], True, dtype=bool)
            g.edge_is_true = torch.full([g.num_edges], True, dtype=bool)

        # create fake graph with the needed # of nodes and edges
        fake = Data()
        fake.num_nodes = num_pad_nodes
        fake.num_edges = num_pad_edges
        fake.graph_is_true = torch.tensor([False], dtype=bool)
        fake.node_is_true = torch.full([num_pad_nodes], False, dtype=bool)
        fake.edge_is_true = torch.full([num_pad_edges], False, dtype=bool)

        for key, value in real_graphs[0]:
            if not torch.is_tensor(value):
                continue

            if key == "graph_is_true" or key == "node_is_true" or key == "edge_is_true":
                continue

            dim = real_graphs[0].__cat_dim__(key, value)
            pad_shape = list(value.shape)

            if batch.is_node_attr(key):
                pad_shape[dim] = num_pad_nodes
                pad_value = self.node_value
            elif batch.is_edge_attr(key):
                pad_shape[dim] = num_pad_edges
                if key == "edge_index":
                    # Padding edges are self-loops on the first padding node
                    pad_value = 0
                else:
                    pad_value = self.edge_value
            # identify graph attributes, pad nan label for the fake graph
            elif key.startswith("graph_"):
                num_pad_graphs = 1  # we pad with one big fake graph
                pad_shape[dim] = num_pad_graphs
                pad_value = float("nan")
            else:
                continue

            pad_value = value.new_full(pad_shape, pad_value)
            fake[key] = torch.cat([pad_value], dim=dim)
        real_graphs.append(fake)
        new_batch = Batch.from_data_list(real_graphs)

        if "num_nodes" in new_batch:
            new_batch.num_nodes = self.max_num_nodes

        return new_batch

    def __repr__(self) -> str:
        s = f"{self.__class__.__name__}("
        s += f"max_num_nodes={self.max_num_nodes}, "
        s += f"max_num_edges={self.max_num_edges}, "
        s += f"node_value={self.node_value}, "
        s += f"edge_value={self.edge_value})"
        return s


# graphium/ipu/ipu_losses.py

import torch
from torch import Tensor
from torch.nn import BCELoss, BCEWithLogitsLoss, MSELoss, L1Loss
from torch._C import _infer_size
from loguru import logger
from graphium.trainer.losses import HybridCELoss


class BCEWithLogitsLossIPU(BCEWithLogitsLoss):
    """
    A modified version of the `torch.nn.BCEWithLogitsLoss` that can ignore NaNs
    by giving them a weight of `0`. This allows it to work with compilation
    and IPUs since it doesn't modify the tensor's shape.
    """

    def forward(self, input: Tensor, target: Tensor) -> Tensor:
        prev_weight = None

        target = target.clone().to(input.dtype)
        weight = self.weight

        # Get the original weight matrix. If None, set all weights = 1
        if weight is not None:
            prev_weight = self.weight.clone()
            new_size = _infer_size(target.size(), weight.size())
            weight = weight.expand(new_size).clone()
        else:
            weight = torch.ones(target.shape, dtype=input.dtype, device=input.device)

        # Replace the nan-targets by 0 or 1. Take the value closest to the input.
        # Give a weight of 0 where there are nan-targets
        nan_targets = target.isnan()
        nan_targets_0 = (input < 0.5) & nan_targets
        nan_targets_1 = (input >= 0.5) & nan_targets
        target[nan_targets_0] = 0.0
        target[nan_targets_1] = 1.0
        weight[nan_targets] = 0.0

        # Compute the loss, and rescale by the number of nan elements
        self.weight = weight
        loss = super().forward(input, target)

        num_real_targets = (~nan_targets).sum()
        factor1 = torch.where(num_real_targets > 0, 1, 0)
        factor2 = torch.where(num_real_targets > 0, 0, 1)
        loss = factor1 * loss * nan_targets.numel() / (num_real_targets + factor2)

        # Reset the self.weight to its original value
        self.weight = prev_weight

        return loss


class BCELossIPU(BCELoss):
    """
    A modified version of the `torch.nn.BCELoss` that can ignore NaNs
    by giving them a weight of `0`. This allows it to work with compilation
    and IPUs since it doesn't modify the tensor's shape.
    """

    def forward(self, input: Tensor, target: Tensor) -> Tensor:
        prev_weight = None

        target = target.clone().to(input.dtype)
        weight = self.weight

        # Get the original weight matrix. If None, set all weights = 1
        if weight is not None:
            prev_weight = self.weight.clone()
            new_size = _infer_size(target.size(), weight.size())
            weight = weight.expand(new_size).clone()
        else:
            weight = torch.ones(target.shape, dtype=input.dtype, device=input.device)

        # Replace the nan-targets by 0 or 1. Take the value closest to the input.
        # Give a weight of 0 where there are nan-targets
        nan_targets = target.isnan()
        nan_targets_0 = (input < 0.5) & nan_targets
        nan_targets_1 = (input >= 0.5) & nan_targets
        target[nan_targets_0] = 0.0
        target[nan_targets_1] = 1.0
        weight[nan_targets] = 0.0

        # Compute the loss, and rescale by the number of nan elements
        self.weight = weight
        loss = super().forward(input, target)

        num_real_targets = (~nan_targets).sum()
        factor1 = torch.where(num_real_targets > 0, 1, 0)
        factor2 = torch.where(num_real_targets > 0, 0, 1)
        loss = factor1 * loss * nan_targets.numel() / (num_real_targets + factor2)

        # Reset the self.weight to its original value
        self.weight = prev_weight

        return loss


class MSELossIPU(MSELoss):
    """
    A modified version of the `torch.nn.MSELoss` that can ignore NaNs
    by giving them the same value for both `input` and `target`.
    This allows it to work with compilation
    and IPUs since it doesn't modify the tensor's shape.
    """

    def forward(self, input: Tensor, target: Tensor) -> Tensor:
        target = target.clone().to(input.dtype)
        input = input.clone()

        # Replace the nan-targets in the input/target tensors by 0
        nan_targets = target.isnan()
        input[nan_targets] = 0.0
        target[nan_targets] = 0.0

        # Compute the loss, and rescale by the number of nan elements
        loss = super().forward(input, target)

        num_real_targets = (~nan_targets).sum()
        factor1 = torch.where(num_real_targets > 0, 1, 0)
        factor2 = torch.where(num_real_targets > 0, 0, 1)
        loss = factor1 * loss * nan_targets.numel() / (num_real_targets + factor2)

        return loss


class L1LossIPU(L1Loss):
    """
    A modified version of the `torch.nn.L1Loss` that can ignore NaNs
    by giving them the same value for both `input` and `target`.
    This allows it to work with compilation
    and IPUs since it doesn't modify the tensor's shape.
    """

    def forward(self, input: Tensor, target: Tensor) -> Tensor:
        target = target.clone().to(input.dtype)
        input = input.clone()

        # Replace the nan-targets in the input/target tensors by 0
        nan_targets = target.isnan()
        input[nan_targets] = 0.0
        target[nan_targets] = 0.0

        # Compute the loss, and rescale by the number of nan elements
        loss = super().forward(input, target)
        num_real_targets = (~nan_targets).sum()
        factor1 = torch.where(num_real_targets > 0, 1, 0)
        factor2 = torch.where(num_real_targets > 0, 0, 1)
        loss = factor1 * loss * nan_targets.numel() / (num_real_targets + factor2)

        return loss


class HybridCELossIPU(HybridCELoss):
    def __init__(
        self,
        n_brackets,
        alpha: float = 0.5,
    ) -> None:
        """
        Parameters:
            n_brackets: the number of brackets that will be used to group the regression targets.
                Expected to have the same size as the number of classes in the transformed regression task.
        """
        super().__init__(n_brackets=n_brackets, alpha=alpha)

    def forward(self, input: Tensor, target: Tensor) -> Tensor:
        """
        Parameters:
            input: (batch_size x n_classes) tensor of logits predicted for each bracket.
            target: (batch_size) or (batch_size, 1) tensor of target brackets in {0, 1, ..., self.n_brackets}.
        """

        target = target.clone().to(input.dtype)
        input = input.clone()

        # Replace the nan-targets in the input/target tensors by 0
        nan_targets = target.isnan()

        # Compute the loss, and rescale by the number of nan elements
        loss = super().forward(input, target, nan_targets)
        return loss


# graphium/ipu/ipu_metrics.py

from typing import Optional, Tuple, Sequence, Literal

import torch
from torch import BoolTensor, IntTensor, Tensor
from torchmetrics.functional import auroc, average_precision, pearson_corrcoef, r2_score
from torchmetrics.utilities.checks import _input_squeeze
from torchmetrics.functional.classification.accuracy import (
    _mode,
    _check_subset_validity,
    _accuracy_compute,
    _accuracy_update,
)
from torchmetrics.functional.classification.precision_recall import _precision_compute, _recall_compute
from torchmetrics.functional.classification.f_beta import _fbeta_compute
from torchmetrics.functional import mean_squared_error, mean_absolute_error
from torchmetrics.utilities.checks import _input_squeeze
from torchmetrics.utilities.enums import AverageMethod

from graphium.utils.tensor import nan_mean
from graphium.ipu.ipu_utils import import_poptorch


def auroc_ipu(
    preds: Tensor,
    target: Tensor,
    num_classes: Optional[int] = None,
    task: Optional[Literal["binary", "multiclass", "multilabel"]] = None,
    pos_label: Optional[int] = None,
    average: Optional[str] = "macro",
    max_fpr: Optional[float] = None,
    sample_weights: Optional[Sequence] = None,
):
    """
    A modified version of the `torchmetrics.functional.auroc` that can ignore NaNs
    by giving them the same value for both `preds` and `target`.
    This allows it to work with compilation
    and IPUs since it doesn't modify the tensor's shape.
    """

    target = target.clone()
    preds = preds.clone()

    # Replace the nan-targets in the preds/target tensors by 0
    nan_targets = target.isnan()
    preds[nan_targets] = 0.0
    target[nan_targets] = 0.0

    # Get the original weight matrix. If None, set all weights = 1
    if sample_weights is None:
        sample_weights = torch.ones(target.shape[0], dtype=preds.dtype, device=preds.device)
    sample_weights[nan_targets] = 0.0

    # Compute the loss, and rescale by the number of nan elements
    score = auroc(
        preds=preds,
        target=target.to(int),
        num_classes=num_classes,
        task=task,
        pos_label=pos_label,
        average=average,
        max_fpr=max_fpr,
        sample_weights=sample_weights,
    )

    return score


def average_precision_ipu(
    preds: Tensor,
    target: Tensor,
    num_classes: Optional[int] = None,
    task: Optional[Literal["binary", "multiclass", "multilabel"]] = None,
    ignore_index: Optional[int] = None,
    pos_label: Optional[int] = None,
    average: Optional[str] = "macro",
    sample_weights: Optional[Sequence] = None,
):
    """
    A modified version of the `torchmetrics.functional.average_precision` that can ignore NaNs
    by giving them the same value for both `preds` and `target`.
    This allows it to work with compilation
    and IPUs since it doesn't modify the tensor's shape.
    """

    target = target.clone()
    preds = preds.clone()

    # Replace the nan-targets in the preds/target tensors by 0
    # Average precision is not sensitive to true negatives
    nan_targets = target.isnan()
    preds[nan_targets] = 0.0
    target[nan_targets] = 0.0

    # No need to use sample weights (which is no longer supported in torchmetrics >=0.10)
    # # Get the original weight matrix. If None, set all weights = 1
    # if sample_weights is None:
    #     sample_weights = torch.ones(target.shape[0], dtype=preds.dtype, device=preds.device)
    # sample_weights[nan_targets] = 0.0

    # Compute the loss, and rescale by the number of nan elements
    score = average_precision(
        preds=preds,
        target=target,
        num_classes=num_classes,
        task=task,
        ignore_index=ignore_index,
        pos_label=pos_label,
        average=average,
        # sample_weights=sample_weights,
    )

    return score


def precision_ipu(
    preds: Tensor,
    target: Tensor,
    average: Optional[str] = "micro",
    mdmc_average: Optional[str] = None,
    ignore_index: Optional[int] = None,
    num_classes: Optional[int] = None,
    threshold: float = 0.5,
    top_k: Optional[int] = None,
    multiclass: Optional[bool] = None,
):
    """
    A modified version of the `torchmetrics.functional.precision` that can ignore NaNs
    by giving them the same value for both `preds` and `target`.
    This allows it to work with compilation
    and IPUs since it doesn't modify the tensor's shape.
    """

    (tp, fp, tn, fn), mode = get_confusion_matrix(
        preds=preds,
        target=target,
        average=average,
        mdmc_average=mdmc_average,
        threshold=threshold,
        top_k=top_k,
        subset_accuracy=False,
        num_classes=num_classes,
        multiclass=multiclass,
        ignore_index=ignore_index,
    )

    return _precision_compute(tp, fp, fn, average, mdmc_average)


def recall_ipu(
    preds: Tensor,
    target: Tensor,
    average: Optional[str] = "micro",
    mdmc_average: Optional[str] = None,
    ignore_index: Optional[int] = None,
    num_classes: Optional[int] = None,
    threshold: float = 0.5,
    top_k: Optional[int] = None,
    multiclass: Optional[bool] = None,
):
    """
    A modified version of the `torchmetrics.functional.recall` that can ignore NaNs
    by giving them the same value for both `preds` and `target`.
    This allows it to work with compilation
    and IPUs since it doesn't modify the tensor's shape.
    """

    (tp, fp, tn, fn), mode = get_confusion_matrix(
        preds=preds,
        target=target,
        average=average,
        mdmc_average=mdmc_average,
        threshold=threshold,
        top_k=top_k,
        num_classes=num_classes,
        multiclass=multiclass,
        ignore_index=ignore_index,
    )

    return _recall_compute(tp, fp, fn, average, mdmc_average)


def accuracy_ipu(
    preds: Tensor,
    target: Tensor,
    average: Optional[str] = "micro",
    mdmc_average: Optional[str] = "global",
    threshold: float = 0.5,
    top_k: Optional[int] = None,
    subset_accuracy: bool = False,
    num_classes: Optional[int] = None,
    multiclass: Optional[bool] = None,
    ignore_index: Optional[int] = None,
) -> Tensor:
    """
    A modified version of the `torchmetrics.functional.accuracy` that can ignore NaNs
    by giving them the same value for both `preds` and `target`.
    This allows it to work with compilation
    and IPUs since it doesn't modify the tensor's shape.

    Args:
        preds: Predictions from model (probabilities, logits or labels)
        target: Ground truth labels
        average:
            Defines the reduction that is applied. Should be one of the following:

            - ``'micro'`` [default]: Calculate the metric globally, across all samples and classes.
            - ``'macro'``: Calculate the metric for each class separately, and average the
              metrics across classes (with equal weights for each class).
            - ``'weighted'``: Calculate the metric for each class separately, and average the
              metrics across classes, weighting each class by its support (``tp + fn``).
            - ``'none'`` or ``None``: Calculate the metric for each class separately, and return
              the metric for every class.
            - ``'samples'``: Calculate the metric for each sample, and average the metrics
              across samples (with equal weights for each sample).

            .. note:: What is considered a sample in the multi-dimensional multi-class case
                depends on the value of ``mdmc_average``.

            .. note:: If ``'none'`` and a given class doesn't occur in the ``preds`` or ``target``,
                the value for the class will be ``nan``.

        mdmc_average:
            Defines how averaging is done for multi-dimensional multi-class inputs (on top of the
            ``average`` parameter). Should be one of the following:

            - ``None`` [default]: Should be left unchanged if your data is not multi-dimensional multi-class.

            - ``'samplewise'``: In this case, the statistics are computed separately for each
              sample on the ``N`` axis, and then averaged over samples.
              The computation for each sample is done by treating the flattened extra axes ``...``
              (see :ref:`pages/classification:input types`) as the ``N`` dimension within the sample,
              and computing the metric for the sample based on that.

            - ``'global'``: In this case the ``N`` and ``...`` dimensions of the inputs
              (see :ref:`pages/classification:input types`)
              are flattened into a new ``N_X`` sample axis, i.e. the inputs are treated as if they
              were ``(N_X, C)``. From here on the ``average`` parameter applies as usual.

        num_classes:
            Number of classes. Necessary for ``'macro'``, ``'weighted'`` and ``None`` average methods.

        threshold:
            Threshold for transforming probability or logit predictions to binary (0,1) predictions, in the case
            of binary or multi-label inputs. Default value of 0.5 corresponds to input being probabilities.
        top_k:
            Number of the highest probability or logit score predictions considered finding the correct label,
            relevant only for (multi-dimensional) multi-class inputs. The
            default value (``None``) will be interpreted as 1 for these inputs.

            Should be left at default (``None``) for all other types of inputs.
        multiclass:
            Used only in certain special cases, where you want to treat inputs as a different type
            than what they appear to be. See the parameter's
            :ref:`documentation section <pages/classification:using the multiclass parameter>`
            for a more detailed explanation and examples.
        ignore_index:
            Integer specifying a target class to ignore. If given, this class index does not contribute
            to the returned score, regardless of reduction method. If an index is ignored, and ``average=None``
            or ``'none'``, the score for the ignored class will be returned as ``nan``.
        subset_accuracy:
            Whether to compute subset accuracy for multi-label and multi-dimensional
            multi-class inputs (has no effect for other input types).

            - For multi-label inputs, if the parameter is set to ``True``, then all labels for
              each sample must be correctly predicted for the sample to count as correct. If it
              is set to ``False``, then all labels are counted separately - this is equivalent to
              flattening inputs beforehand (i.e. ``preds = preds.flatten()`` and same for ``target``).

            - For multi-dimensional multi-class inputs, if the parameter is set to ``True``, then all
              sub-sample (on the extra axis) must be correct for the sample to be counted as correct.
              If it is set to ``False``, then all sub-samples are counter separately - this is equivalent,
              in the case of label predictions, to flattening the inputs beforehand (i.e.
              ``preds = preds.flatten()`` and same for ``target``). Note that the ``top_k`` parameter
              still applies in both cases, if set.

    Raises:
        ValueError:
            If ``top_k`` parameter is set for ``multi-label`` inputs.
        ValueError:
            If ``average`` is none of ``"micro"``, ``"macro"``, ``"weighted"``, ``"samples"``, ``"none"``, ``None``.
        ValueError:
            If ``mdmc_average`` is not one of ``None``, ``"samplewise"``, ``"global"``.
        ValueError:
            If ``average`` is set but ``num_classes`` is not provided.
        ValueError:
            If ``num_classes`` is set
            and ``ignore_index`` is not in the range ``[0, num_classes)``.
        ValueError:
            If ``top_k`` is not an ``integer`` larger than ``0``.
    """

    (tp, fp, tn, fn), mode = get_confusion_matrix(
        preds=preds,
        target=target,
        average=average,
        mdmc_average=mdmc_average,
        threshold=threshold,
        top_k=top_k,
        subset_accuracy=subset_accuracy,
        num_classes=num_classes,
        multiclass=multiclass,
        ignore_index=ignore_index,
    )

    return _accuracy_compute(tp, fp, tn, fn, average, mdmc_average, mode)


def get_confusion_matrix(
    preds: Tensor,
    target: Tensor,
    average: Optional[str] = "micro",
    mdmc_average: Optional[str] = "global",
    threshold: float = 0.5,
    top_k: Optional[int] = None,
    subset_accuracy: bool = False,
    num_classes: Optional[int] = None,
    multiclass: Optional[bool] = None,
    ignore_index: Optional[int] = None,
) -> Tuple[Tuple[Tensor], Tensor]:
    """
    Calculates the confusion matrix according to the specified average method.

    Args:
        preds: Predictions from model (probabilities, logits or labels)
        target: Ground truth labels
        average:
            Defines the reduction that is applied. Should be one of the following:

            - ``'micro'`` [default]: Calculate the metric globally, across all samples and classes.
            - ``'macro'``: Calculate the metric for each class separately, and average the
              metrics across classes (with equal weights for each class).
            - ``'weighted'``: Calculate the metric for each class separately, and average the
              metrics across classes, weighting each class by its support (``tp + fn``).
            - ``'none'`` or ``None``: Calculate the metric for each class separately, and return
              the metric for every class.
            - ``'samples'``: Calculate the metric for each sample, and average the metrics
              across samples (with equal weights for each sample).

            .. note:: What is considered a sample in the multi-dimensional multi-class case
                depends on the value of ``mdmc_average``.

            .. note:: If ``'none'`` and a given class doesn't occur in the ``preds`` or ``target``,
                the value for the class will be ``nan``.

        mdmc_average:
            Defines how averaging is done for multi-dimensional multi-class inputs (on top of the
            ``average`` parameter). Should be one of the following:

            - ``None`` [default]: Should be left unchanged if your data is not multi-dimensional multi-class.

            - ``'samplewise'``: In this case, the statistics are computed separately for each
              sample on the ``N`` axis, and then averaged over samples.
              The computation for each sample is done by treating the flattened extra axes ``...``
              (see :ref:`pages/classification:input types`) as the ``N`` dimension within the sample,
              and computing the metric for the sample based on that.

            - ``'global'``: In this case the ``N`` and ``...`` dimensions of the inputs
              (see :ref:`pages/classification:input types`)
              are flattened into a new ``N_X`` sample axis, i.e. the inputs are treated as if they
              were ``(N_X, C)``. From here on the ``average`` parameter applies as usual.

        num_classes:
            Number of classes. Necessary for ``'macro'``, ``'weighted'`` and ``None`` average methods.

        threshold:
            Threshold for transforming probability or logit predictions to binary (0,1) predictions, in the case
            of binary or multi-label inputs. Default value of 0.5 corresponds to input being probabilities.
        top_k:
            Number of the highest probability or logit score predictions considered finding the correct label,
            relevant only for (multi-dimensional) multi-class inputs. The
            default value (``None``) will be interpreted as 1 for these inputs.

            Should be left at default (``None``) for all other types of inputs.
        multiclass:
            Used only in certain special cases, where you want to treat inputs as a different type
            than what they appear to be. See the parameter's
            :ref:`documentation section <pages/classification:using the multiclass parameter>`
            for a more detailed explanation and examples.
        ignore_index:
            Integer specifying a target class to ignore. If given, this class index does not contribute
            to the returned score, regardless of reduction method. If an index is ignored, and ``average=None``
    """
    allowed_average = ["micro", "macro", "weighted", "samples", "none", None]
    if average not in allowed_average:
        raise ValueError(f"The `average` has to be one of {allowed_average}, got {average}.")

    if average in ["macro", "weighted", "none", None] and (not num_classes or num_classes < 1):
        raise ValueError(f"When you set `average` as {average}, you have to provide the number of classes.")

    allowed_mdmc_average = [None, "samplewise", "global"]
    if mdmc_average not in allowed_mdmc_average:
        raise ValueError(f"The `mdmc_average` has to be one of {allowed_mdmc_average}, got {mdmc_average}.")

    if num_classes and ignore_index is not None and (not ignore_index < num_classes or num_classes == 1):
        raise ValueError(
            f"The `ignore_index` {ignore_index} is not valid for inputs with {num_classes} classes"
        )

    if top_k is not None and (not isinstance(top_k, int) or top_k <= 0):
        raise ValueError(f"The `top_k` should be an integer larger than 0, got {top_k}")

    #### ADDED ####
    # Put all the NaNs as the 0-class
    nans = torch.isnan(target)
    target[nans] = 0
    preds[nans] = 0
    if (preds.ndim > 1) and (preds.shape[1] > 1):
        preds[nans, 0] = 1
    target = target.to(int)
    #### END ADDED ####

    preds, target = _input_squeeze(preds, target)
    mode = _mode(preds, target, threshold, top_k, num_classes, multiclass, ignore_index)
    reduce = "macro" if average in ["weighted", "none", None] else average

    if subset_accuracy and _check_subset_validity(mode):
        # correct, total = _subset_accuracy_update(preds, target, threshold, top_k, ignore_index)
        # return _subset_accuracy_compute(correct, total)
        raise NotImplementedError("subset_accuracy not implemented")
    tp, fp, tn, fn = _accuracy_update(
        preds, target, reduce, mdmc_average, threshold, num_classes, top_k, multiclass, ignore_index, mode
    )

    #### ADDED ####
    num_nans = nans.sum(0)
    if tp.numel() > 1:
        tp[0] = tp[0] - num_nans
        tn[1:] = tn[1:] - num_nans
    else:
        tn = tn - num_nans
        if (preds.ndim > 1) and (preds.shape[1] > 1):
            tp = tp - num_nans
    #### END ADDED ####

    return (tp, fp, tn, fn), mode


class NaNTensor(Tensor):
    """
    Class to create and manage a NaN tensor along it's properties

    The goal of the class is to override the regular tensor such that the basic
    operations (sum, mean, max, etc) ignore the NaNs in the input.
    It also supports NaNs in integer tensors (as the lowest integer possible).
    """

    @property
    def get_nans(self) -> BoolTensor:
        """
        Gets the boolean Tensor containing the location of NaNs.
        In the case of an integer tensor, this returns where the tensor is equal to its minimal value
        In the case of a boolean tensor, this returns a Tensor filled with `False`
        """
        if self.is_floating_point():
            return self.isnan()
        elif self.is_signed():
            return self == torch.iinfo(self.dtype).min
        else:
            return torch.zeros(self.shape, device=self.device, dtype=bool)

    def sum(self, *args, **kwargs) -> Tensor:
        """
        Overloads the traditional sum to ignore the NaNs
        """
        tensor = self.to(float)
        tensor[self.get_nans] = float("nan")
        if self.is_floating_point():
            dtype = self.dtype
        else:
            dtype = torch.int64
        return tensor.nansum(*args, **kwargs).to(dtype)

    def mean(self, *args, **kwargs) -> Tensor:
        """
        Overloads the traditional mean to ignore the NaNs
        """
        tensor = self.to(float)
        tensor[self.get_nans] = float("nan")
        return nan_mean(tensor, *args, **kwargs).to(self.dtype)

    def numel(self) -> int:
        """
        Returns the number of non-NaN elements.
        """
        return super(NaNTensor, ~self.get_nans).sum()

    def min(self, *args, **kwargs) -> Tensor:
        """
        Returns the min vale of a tensor whitout NaNs
        """
        tensor = self
        tensor = tensor[~self.get_nans]
        return super(NaNTensor, tensor).min(*args, **kwargs)

    def max(self, *args, **kwargs) -> Tensor:
        """
        Returns the max vale of a tensor whitout NaNs
        """
        tensor = self
        tensor = tensor[~self.get_nans]
        return super(NaNTensor, tensor).max(*args, **kwargs)

    def argsort(self, dim=-1, descending=False) -> IntTensor:
        """
        Return the indices that sort the tensor, while putting all the NaNs to the end of the sorting.
        """
        tensor = self
        if descending:
            tensor[tensor.get_nans] = float("-inf")
        else:
            tensor[tensor.get_nans] = float("inf")
        return super(NaNTensor, tensor).argsort(dim=dim, descending=descending)

    def size(self, dim) -> Tensor:
        """
        Instead of returning the size, return the number of non-NaN elements in
        a specific dimension. Useful for the `r2_score` metric.
        """
        return (~self.get_nans).sum(dim=dim)

    def __lt__(self, other) -> Tensor:
        """
        Stupid fix that allows the code to work with `r2_score`,
        since it requires the size to be > 2. But since `self.size` now returns
        a Tensor instead of a value, we check that all elements are > 2.
        """
        if (not isinstance(other, Tensor)) and (other == 2):
            return super().__lt__(other).all()
        else:
            return super().__lt__(other)

    @classmethod
    def __torch_function__(cls, func, types, args=(), kwargs=None):
        """
        This __torch_function__ implementation wraps subclasses such that
        methods called on subclasses return a subclass instance instead of
        a ``torch.Tensor`` instance.

        One corollary to this is that you need coverage for torch.Tensor
        methods if implementing __torch_function__ for subclasses.

        Affects the call torch.sum() as to behave the same way as NaNTensor.sum()

        We recommend always calling ``super().__torch_function__`` as the base
        case when doing the above.

        While not mandatory, we recommend making `__torch_function__` a classmethod.
        """
        if func.__name__ == "sum":
            kwargs = {} if kwargs is None else kwargs
            return args[0].sum(*args[1:], **kwargs)
        else:
            return super().__torch_function__(func, types, args=args, kwargs=kwargs)


def pearson_ipu(preds, target):
    """Computes pearson correlation coefficient.

    Handles NaNs in the target without reshaping tensors in order to work on IPU.

    Args:
        preds: estimated scores
        target: ground truth scores
    """
    preds = NaNTensor(preds)
    target = NaNTensor(target)
    preds[target.get_nans] = float("nan")
    pearson = pearson_corrcoef(preds, target.to(preds.dtype))
    return Tensor(pearson)


def spearman_ipu(preds, target):
    """Computes spearman rank correlation coefficient.

    Handles NaNs in the target without reshaping tensors in order to work on IPU.

    Args:
        preds: estimated scores
        target: ground truth scores
    """
    nans = target.isnan()
    dtype = preds.dtype
    preds[nans] = float("inf")
    target[nans] = float("inf")
    preds_sort = _rank_data(preds).to(dtype=dtype)
    target_sort = _rank_data(target).to(dtype=dtype)
    target_sort[nans] = float("nan")
    spearman = pearson_ipu(preds_sort, target_sort)
    return Tensor(spearman)


def _rank_data(data: Tensor) -> Tensor:
    """Calculate the rank for each element of a tensor.

    The rank refers to the indices of an element in the corresponding sorted tensor (starting from 1).
    Duplicates of the same value will be assigned the mean of their rank.

    Adopted from `Rank of element tensor`_
    """
    n = data.numel()
    rank = torch.empty_like(data)
    idx = data.argsort()
    rank[idx] = torch.arange(1, n + 1, dtype=data.dtype, device=data.device)

    # TODO: Repeats not yet supported
    # repeats = _find_repeats(data)
    # for r in repeats:
    #     condition = data == r
    #     rank[condition] = rank[condition].mean()
    return rank


def r2_score_ipu(preds, target, *args, **kwargs) -> Tensor:
    """
    Computes r2 score also known as `R2 Score_Coefficient Determination`_:

    .. math:: R^2 = 1 - \frac{SS_{res}}{SS_{tot}}

    where :math:`SS_{res}=\sum_i (y_i - f(x_i))^2` is the sum of residual squares, and
    :math:`SS_{tot}=\sum_i (y_i - \bar{y})^2` is total sum of squares. Can also calculate
    adjusted r2 score given by

    .. math:: R^2_{adj} = 1 - \frac{(1-R^2)(n-1)}{n-k-1}

    where the parameter :math:`k` (the number of independent regressors) should
    be provided as the ``adjusted`` argument.
    Handles NaNs without reshaping tensors in order to work on IPU.

    Args:
        preds: estimated labels
        target: ground truth labels
        adjusted: number of independent regressors for calculating adjusted r2 score.
        multioutput: Defines aggregation in the case of multiple output scores. Can be one of the following strings:

            * ``'raw_values'`` returns full set of scores
            * ``'uniform_average'`` scores are uniformly averaged
            * ``'variance_weighted'`` scores are weighted by their individual variances
    """
    preds = NaNTensor(preds)
    target = NaNTensor(target)
    preds[target.get_nans] = float("nan")
    score = r2_score(preds, target, *args, **kwargs)
    return Tensor(score)


def fbeta_score_ipu(
    preds: Tensor,
    target: Tensor,
    beta: float = 1.0,
    average: Optional[str] = "micro",
    mdmc_average: Optional[str] = None,
    ignore_index: Optional[int] = None,
    num_classes: Optional[int] = None,
    threshold: float = 0.5,
    top_k: Optional[int] = None,
    multiclass: Optional[bool] = None,
):
    """
    A modified version of the `torchmetrics.functional.classification.f_beta._fbeta_compute`
    that can ignore NaNs by giving them the same value for both `preds` and `target`.
    This allows it to work with compilation
    and IPUs since it doesn't modify the tensor's shape.

    Args:
        preds: Predictions from model (probabilities, logits or labels)
        target: Ground truth labels
        average:
            Defines the reduction that is applied. Should be one of the following:

            - ``'micro'`` [default]: Calculate the metric globally, across all samples and classes.
            - ``'macro'``: Calculate the metric for each class separately, and average the
              metrics across classes (with equal weights for each class).
            - ``'weighted'``: Calculate the metric for each class separately, and average the
              metrics across classes, weighting each class by its support (``tp + fn``).
            - ``'none'`` or ``None``: Calculate the metric for each class separately, and return
              the metric for every class.
            - ``'samples'``: Calculate the metric for each sample, and average the metrics
              across samples (with equal weights for each sample).

            .. note:: What is considered a sample in the multi-dimensional multi-class case
                depends on the value of ``mdmc_average``.

            .. note:: If ``'none'`` and a given class doesn't occur in the ``preds`` or ``target``,
                the value for the class will be ``nan``.

        mdmc_average:
            Defines how averaging is done for multi-dimensional multi-class inputs (on top of the
            ``average`` parameter). Should be one of the following:

            - ``None`` [default]: Should be left unchanged if your data is not multi-dimensional multi-class.

            - ``'samplewise'``: In this case, the statistics are computed separately for each
              sample on the ``N`` axis, and then averaged over samples.
              The computation for each sample is done by treating the flattened extra axes ``...``
              (see :ref:`pages/classification:input types`) as the ``N`` dimension within the sample,
              and computing the metric for the sample based on that.

            - ``'global'``: In this case the ``N`` and ``...`` dimensions of the inputs
              (see :ref:`pages/classification:input types`)
              are flattened into a new ``N_X`` sample axis, i.e. the inputs are treated as if they
              were ``(N_X, C)``. From here on the ``average`` parameter applies as usual.

        num_classes:
            Number of classes. Necessary for ``'macro'``, ``'weighted'`` and ``None`` average methods.

        threshold:
            Threshold for transforming probability or logit predictions to binary (0,1) predictions, in the case
            of binary or multi-label inputs. Default value of 0.5 corresponds to input being probabilities.
        top_k:
            Number of the highest probability or logit score predictions considered finding the correct label,
            relevant only for (multi-dimensional) multi-class inputs. The
            default value (``None``) will be interpreted as 1 for these inputs.

            Should be left at default (``None``) for all other types of inputs.
        multiclass:
            Used only in certain special cases, where you want to treat inputs as a different type
            than what they appear to be. See the parameter's
            :ref:`documentation section <pages/classification:using the multiclass parameter>`
            for a more detailed explanation and examples.
        ignore_index:
            Integer specifying a target class to ignore. If given, this class index does not contribute
            to the returned score, regardless of reduction method. If an index is ignored, and ``average=None``
            or ``'none'``, the score for the ignored class will be returned as ``nan``.
        subset_accuracy:
            Whether to compute subset accuracy for multi-label and multi-dimensional
            multi-class inputs (has no effect for other input types).

            - For multi-label inputs, if the parameter is set to ``True``, then all labels for
              each sample must be correctly predicted for the sample to count as correct. If it
              is set to ``False``, then all labels are counted separately - this is equivalent to
              flattening inputs beforehand (i.e. ``preds = preds.flatten()`` and same for ``target``).

            - For multi-dimensional multi-class inputs, if the parameter is set to ``True``, then all
              sub-sample (on the extra axis) must be correct for the sample to be counted as correct.
              If it is set to ``False``, then all sub-samples are counter separately - this is equivalent,
              in the case of label predictions, to flattening the inputs beforehand (i.e.
              ``preds = preds.flatten()`` and same for ``target``). Note that the ``top_k`` parameter
              still applies in both cases, if set.

    Raises:
        ValueError:
            If ``top_k`` parameter is set for ``multi-label`` inputs.
        ValueError:
            If ``average`` is none of ``"micro"``, ``"macro"``, ``"weighted"``, ``"samples"``, ``"none"``, ``None``.
        ValueError:
            If ``mdmc_average`` is not one of ``None``, ``"samplewise"``, ``"global"``.
        ValueError:
            If ``average`` is set but ``num_classes`` is not provided.
        ValueError:
            If ``num_classes`` is set
            and ``ignore_index`` is not in the range ``[0, num_classes)``.
        ValueError:
            If ``top_k`` is not an ``integer`` larger than ``0``.
    """

    (tp, fp, tn, fn), mode = get_confusion_matrix(
        preds=preds,
        target=target,
        average=average,
        mdmc_average=mdmc_average,
        ignore_index=ignore_index,
        num_classes=num_classes,
        threshold=threshold,
        top_k=top_k,
        multiclass=multiclass,
    )

    b2 = beta**2
    fbeta = ((1 + b2) * tp) / ((1 + b2) * tp + b2 * fn + fp)

    if average in (None, "none", AverageMethod.NONE):
        pass
    elif average == AverageMethod.MICRO:
        pass
    elif average == AverageMethod.MACRO:
        fbeta = fbeta.mean()
    elif average == AverageMethod.WEIGHTED:
        weights = tp + fn
        fbeta = (weights * fbeta).sum() / weights.sum()
    else:
        raise ValueError(
            f"`average={average}` not yet supported. Chose between None, Micro, Macro, or Weighted"
        )

    return fbeta


def f1_score_ipu(
    preds: Tensor,
    target: Tensor,
    beta: float = 1.0,
    average: Optional[str] = "micro",
    mdmc_average: Optional[str] = None,
    ignore_index: Optional[int] = None,
    num_classes: Optional[int] = None,
    threshold: float = 0.5,
    top_k: Optional[int] = None,
    multiclass: Optional[bool] = None,
):
    """
    A modified version of the `torchmetrics.functional.classification.f_beta._fbeta_compute`
    that can ignore NaNs by giving them the same value for both `preds` and `target`.
    Used to calculate the f1_score on IPU with beta parameter equal to 1.0
    This allows it to work with compilation and IPUs since it doesn't modify the tensor's shape.

    Computes f_beta metric from stat scores: true positives, false positives, true negatives, false negatives.

    Args:
        tp: True positives
        fp: False positives
        tn: True negatives
        fn: False negatives
        beta: The parameter `beta` (which determines the weight of recall in the combined score)
        ignore_index: Integer specifying a target class to ignore. If given, this class index does not contribute
            to the returned score, regardless of reduction method
        average: Defines the reduction that is applied
        mdmc_average: Defines how averaging is done for multi-dimensional multi-class inputs (on top of the
            ``average`` parameter)
    """

    return fbeta_score_ipu(
        preds,
        target,
        beta=beta,
        average=average,
        mdmc_average=mdmc_average,
        ignore_index=ignore_index,
        num_classes=num_classes,
        threshold=threshold,
        top_k=top_k,
        multiclass=multiclass,
    )


def mean_squared_error_ipu(preds: Tensor, target: Tensor, squared: bool) -> Tensor:
    """Computes mean squared error.

    Handles NaNs without reshaping tensors in order to work on IPU.

    Args:
        preds: estimated labels
        target: ground truth labels
        squared: returns RMSE value if set to False

    Return:
        Tensor with MSE
    """
    target = target.clone()
    preds = preds.clone()

    # Replace the nan-targets in the preds/target tensors by 0
    nan_targets = target.isnan()
    preds[nan_targets] = 0.0
    target[nan_targets] = 0.0

    # Compute the loss, and rescale by the number of nan elements
    loss = mean_squared_error(preds, target, squared)

    if squared:
        factor = nan_targets.numel() / ((~nan_targets).sum())
    else:
        factor = (nan_targets.numel() / ((~nan_targets).sum())).sqrt()

    loss = loss * factor

    return loss


def mean_absolute_error_ipu(preds: Tensor, target: Tensor) -> Tensor:
    """Computes mean absolute error.

    Handles NaNs without reshaping tensors in order to work on IPU.

    Args:
        preds: estimated labels
        target: ground truth labels

    Return:
        Tensor with MAE
    """
    target = target.clone()
    preds = preds.clone()

    # Replace the nan-targets in the preds/target tensors by 0
    nan_targets = target.isnan()
    preds[nan_targets] = 0.0
    target[nan_targets] = 0.0

    # Compute the loss, and rescale by the number of nan elements
    loss = mean_absolute_error(preds, target)
    loss = loss * nan_targets.numel() / ((~nan_targets).sum())

    return loss


# graphium/ipu/ipu_simple_lightning.py

# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
import lightning
from lightning_graphcore import IPUStrategy
from lightning.pytorch.loggers import WandbLogger

import torch
from torch import nn

import torchvision
import torchvision.transforms as transforms

import mup

from graphium.nn.base_layers import FCLayer
from graphium.utils.mup import set_base_shapes


ON_IPU = True  # Change this line to run on CPU
SEED = 42


# The simple PyTorch model used in each of these examples
class SimpleTorchModel(torch.nn.Module):
    def __init__(self, in_dim, hidden_dim, kernel_size, num_classes):
        super().__init__()
        self.in_dim = in_dim
        self.hidden_dim = hidden_dim
        self.kernel_size = kernel_size
        self.num_classes = num_classes

        conv_block = nn.Sequential(
            nn.Conv2d(in_channels=in_dim, out_channels=hidden_dim, kernel_size=kernel_size),
            nn.BatchNorm2d(hidden_dim),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size),
            nn.MaxPool2d(kernel_size),
        )

        self.the_network = nn.Sequential(
            conv_block,
            torch.nn.Flatten(),
            FCLayer(4 * hidden_dim, hidden_dim),
            FCLayer(hidden_dim, hidden_dim),
            FCLayer(hidden_dim, num_classes, activation=None, is_readout_layer=True),
            nn.LogSoftmax(1),
        )

    def make_mup_base_kwargs(self, divide_factor: float = 2.0):
        return dict(
            in_dim=self.in_dim,
            hidden_dim=round(self.hidden_dim / divide_factor),
            kernel_size=self.kernel_size,
            num_classes=self.num_classes,
        )

    def forward(self, x):
        return self.the_network(x)


# This class shows a minimal lightning example. This example uses our own
# SimpleTorchModel which is a basic 2 conv, 2 FC torch network. It can be
# found in simple_torch_model.py.
class SimpleLightning(lightning.LightningModule):
    def __init__(self, in_dim, hidden_dim, kernel_size, num_classes, on_ipu):
        super().__init__()
        self.model = SimpleTorchModel(
            in_dim=in_dim, hidden_dim=hidden_dim, kernel_size=kernel_size, num_classes=num_classes
        )
        self.on_ipu = on_ipu

    def training_step(self, batch, _):
        x, label = batch
        prediction = self.model(x)
        loss = torch.nn.functional.nll_loss(prediction, label)
        return loss

    def validation_step(self, batch, _):
        x, label = batch
        prediction = self.model(x)
        preds = torch.argmax(prediction, dim=1)
        acc = torch.sum(preds == label).float() / len(label)
        loss = torch.nn.functional.nll_loss(prediction, label)
        return loss, acc

    # PopTorch doesn't currently support logging within steps. Use the Lightning
    # callback hooks instead.
    def on_train_batch_end(self, outputs, batch, batch_idx):
        self.log("StepLoss", outputs["loss"])

    def validation_epoch_end(self, outputs):
        loss = [out[0] for out in outputs]
        self.log("val_loss", torch.stack(loss).mean(), prog_bar=True)

        acc = [out[1] for out in outputs]
        self.log("val_acc", torch.stack(acc).mean(), prog_bar=True)

    def configure_optimizers(self):
        adam = torch.optim.Adam

        if self.on_ipu:
            import poptorch

            adam = poptorch.optim.Adam

        optimizer = mup.MuAdam(self.parameters(), lr=0.01, impl=adam)
        return optimizer


if __name__ == "__main__":
    torch.manual_seed(SEED)

    # Create the model as usual.
    predictor = SimpleLightning(in_dim=1, hidden_dim=32, kernel_size=3, num_classes=10, on_ipu=ON_IPU)
    model = predictor.model
    base = model.__class__(**model.make_mup_base_kwargs(divide_factor=2))
    predictor.model = set_base_shapes(model, base, rescale_params=False)

    torch.manual_seed(SEED)
    # Normal PyTorch dataset.
    train_set = torchvision.datasets.FashionMNIST(
        "out/FashionMNIST", train=True, download=True, transform=transforms.Compose([transforms.ToTensor()])
    )
    val_set = torchvision.datasets.FashionMNIST(
        "out/FashionMNIST", train=False, download=True, transform=transforms.Compose([transforms.ToTensor()])
    )

    # Normal PyTorch dataloader.
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=16, shuffle=True)
    val_loader = torch.utils.data.DataLoader(val_set, batch_size=16, shuffle=False)

    torch.manual_seed(SEED)

    ipus = None
    plugins = None
    if ON_IPU:
        import poptorch

        training_opts = poptorch.Options()
        inference_opts = poptorch.Options()

        # Set the seeds
        training_opts.randomSeed(SEED)
        inference_opts.randomSeed(SEED)
        ipus = 1
        strategy = IPUStrategy(training_opts=training_opts, inference_opts=inference_opts)

    trainer = lightning.Trainer(
        logger=WandbLogger(),
        ipus=ipus,
        max_epochs=3,
        log_every_n_steps=1,
        plugins=plugins,
    )

    # When fit is called the model will be compiled for IPU and will run on the available IPU devices.
    trainer.fit(predictor, train_dataloaders=train_loader, val_dataloaders=val_loader)


# graphium/ipu/ipu_utils.py

import os
import tempfile
from datetime import datetime
from copy import deepcopy
from types import ModuleType
from typing import Optional, Tuple, List
import torch


def import_poptorch(raise_error=True) -> Optional[ModuleType]:
    """
    Import poptorch and returns it.
    It is wrapped in a function to avoid breaking the code
    for non-IPU devices which did not install poptorch.

    Parameters:
        raise_error: Whether to raise an error if poptorch is unavailable.
            If `False`, return `None`

    Returns:
        The poptorch module

    """
    try:
        import poptorch

        return poptorch
    except ImportError as e:
        if raise_error:
            raise e
        return


def is_running_on_ipu() -> bool:
    """
    Returns whether the current module is running on ipu.
    Needs to be used in the `forward` or `backward` pass.
    """
    poptorch = import_poptorch(raise_error=False)
    on_ipu = (poptorch is not None) and (poptorch.isRunningOnIpu())
    return on_ipu


def load_ipu_options(
    ipu_opts: List[str],
    seed: Optional[int] = None,
    model_name: Optional[str] = None,
    gradient_accumulation: Optional[int] = None,
    precision: Optional[int] = None,
    ipu_inference_opts: Optional[List[str]] = None,
) -> Tuple["poptorch.Options", "poptorch.Options"]:
    """
    Load the IPU options from the config file.

    Parameters:
        ipu_cfg: The list  configurations for the IPU, written as a list of strings to make use of `poptorch.Options.loadFromFile`

            write a temporary config gile, and read it. See `Options.loadFromFile`
            #? see the tutorial for IPU options here
            # https://github.com/graphcore/tutorials/tree/sdk-release-2.6/tutorials/pytorch/efficient_data_loading
            #? see the full documentation for ipu options here
            # https://docs.graphcore.ai/projects/poptorch-user-guide/en/latest/reference.html?highlight=options#poptorch.Options

            ***minibatch size***: The number of samples processed by one simple fwd/bwd pass.
            = # of samples in a minibatch

            ***device iterations***: A device iteration corresponds to one iteration of the training loop executed on the IPU, starting with data-loading and ending with a weight update.
            In this simple case, when we set n deviceIterations, the host will prepare n mini-batches in an infeed queue so the IPU can perform efficiently n iterations.
            = # of minibatches to be processed at a time
            = # of training / backward pass in this call

            ***gradient accumulation factor***: After each backward pass the gradients are accumulated together for K mini-batches. set K in the argument
            = # of minibatches to accumulate gradients from

            ***replication factor***: Replication describes the process of running multiple instances of the same model simultaneously on different IPUs to achieve data parallelism.
            If the model requires N IPUs and the replication factor is M, N x M IPUs will be necessary.
            = # of times the model is copied to speed up computation, each replica of the model is sent a different subset of the dataset

            ***global batch size***: In a single device iteration, many mini-batches may be processed and the resulting gradients accumulated.
            We call this total number of samples processed for one optimiser step the global batch size.
            = total number of samples processed for *one optimiser step*
            = (minibatch size x Gradient accumulation factor) x Number of replicas

        seed: random seed for the IPU
        model_name: Name of the model, to be used for ipu profiling
        ipu_inference_opts: optional IPU configuration overrides for inference.
            If this is provided, options in this file override those in `ipu_file` for inference.

    Returns:

        training_opts: IPU options for the training set.

        inference_opts: IPU options for inference.
            It differs from the `training_opts` by enforcing `gradientAccumulation` to 1

    """

    poptorch = import_poptorch()
    ipu_options = poptorch.Options()
    ipu_opts_file = ipu_options_list_to_file(ipu_opts)
    ipu_options.loadFromFile(ipu_opts_file.name)
    ipu_opts_file.close()

    ipu_options.outputMode(poptorch.OutputMode.All)
    if seed is not None:
        ipu_options.randomSeed(seed)
    if model_name is not None:
        ipu_options.modelName(f"{model_name}_train")
    if gradient_accumulation is not None:
        current = ipu_options.Training.gradient_accumulation
        assert (current == 1) or (
            current == gradient_accumulation
        ), f"Received inconsistent gradient accumulation `{current}` and `{gradient_accumulation}"
        ipu_options.Training.gradientAccumulation(gradient_accumulation)

    if precision == "16-true":
        # IPUOptions.loadFromFile currently doesn't support setting half partials, doing it here
        ipu_options.Precision.setPartialsType(torch.half)
    training_opts = ipu_options

    # Change the inference options to remove gradient accumulation
    inference_opts = deepcopy(ipu_options)
    inference_opts.Training.gradientAccumulation(1)
    if ipu_inference_opts is not None:
        ipu_inference_opts_file = ipu_options_list_to_file(ipu_inference_opts)
        inference_opts.loadFromFile(ipu_inference_opts_file.name)
        ipu_inference_opts_file.close()

    return training_opts, inference_opts


def ipu_options_list_to_file(ipu_opts: Optional[List[str]]) -> tempfile._TemporaryFileWrapper:
    """
    Create a temporary file from a list of ipu configs, such that it can be read by `poptorch.Options.loadFromFile`

    Parameters:
        ipu_opts: The list  configurations for the IPU, written as a list of strings to make use of `poptorch.Options.loadFromFile`
    Returns:
        tmp_file: The temporary file of ipu configs
    """
    if ipu_opts is None:
        return

    tmp_file = tempfile.NamedTemporaryFile("w", delete=True)
    for s in ipu_opts:
        tmp_file.write(s + "\n")
    tmp_file.flush()
    return tmp_file


# graphium/ipu/ipu_wrapper.py

from typing import Dict, Any, Optional, Callable, Union, Type, Tuple, Iterable

from torch_geometric.data import Batch
from torch import Tensor
from lightning_graphcore import IPUStrategy
from lightning.pytorch.utilities.types import STEP_OUTPUT
from lightning.pytorch.trainer.states import RunningStage

from graphium.trainer.predictor import PredictorModule
from graphium.ipu.ipu_utils import import_poptorch

import torch
from torch_geometric.data import Data, Batch
from torch_geometric.data.data import BaseData
from loguru import logger
import functools
import collections
from graphium.data.utils import get_keys

poptorch = import_poptorch()


class PyGArgsParser(poptorch.ICustomArgParser):
    """
    This class is responsible for converting a PyG Batch from and to
    a tensor of tuples. This allows PyG Batch to be used as inputs to
    IPU programs. Copied from poppyg repo, in the future import from
    the repo directly.
    """

    @staticmethod
    def sortedTensorKeys(struct: BaseData) -> Iterable[str]:
        """
        Find all the keys that map to a tensor value in struct. The keys
        are returned in sorted order.
        """
        all_keys = sorted(get_keys(struct))

        def isTensor(k: str) -> bool:
            return isinstance(struct[k], torch.Tensor)

        return filter(isTensor, all_keys)

    def yieldTensors(self, struct: BaseData):
        """
        yield every torch.Tensor in struct in sorted order
        """
        for k in self.sortedTensorKeys(struct):
            yield struct[k]

    def reconstruct(self, original_structure: BaseData, tensor_iterator: Iterable[Tensor]):
        """
        Create a new instance with the same class type as the
        original_structure. This new instance will be initialized with tensors
        from the provided iterator and uses the same sorted keys from the
        yieldTensors() implementation.
        """
        tensor_keys = self.sortedTensorKeys(original_structure)
        kwargs = {k: next(tensor_iterator) for k in tensor_keys}

        for k in get_keys(original_structure):
            if k not in kwargs:
                # copy non-tensor properties to the new instance
                kwargs[k] = original_structure[k]

        cls = original_structure.__class__

        if issubclass(cls, Batch):
            kwargs["_base_cls"] = Data
            return Batch(**kwargs)

        return cls(**kwargs)


# PyG uses the BaseData object as the root for data and batch objects
poptorch.registerCustomArgParser(BaseData, PyGArgsParser())


class PredictorModuleIPU(PredictorModule):
    """
    This class wraps around the `PredictorModule` to make it work with IPU and the `IPUPluginGraphium`.
    """

    def __init__(self, *args, **kwargs):
        # Import poptorch in a safe way that will work when working with cpu/gpu
        self.poptorch = import_poptorch()
        super().__init__(*args, **kwargs)

    @staticmethod
    def compute_loss(
        preds: Dict[str, Tensor],
        targets: Dict[str, Tensor],
        weights: Optional[Tensor],
        loss_fun: Dict[str, Callable],
        target_nan_mask: Union[Type, str] = "ignore",
        multitask_handling: Optional[str] = None,
    ) -> Tuple[Tensor, Dict[str, Tensor]]:
        return PredictorModule.compute_loss(
            preds, targets, weights, loss_fun, target_nan_mask, multitask_handling
        )

    def on_train_batch_end(self, outputs, batch, batch_idx):
        outputs = self.convert_from_fp16(outputs)
        outputs["loss"] = outputs["loss"][outputs["loss"] != 0].mean()
        super().on_train_batch_end(outputs, batch, batch_idx)

    def training_step(self, batch, batch_idx) -> Dict[str, Any]:
        features, labels = batch["features"], batch["labels"]
        features, labels = self.squeeze_input_dims(features, labels)
        dict_input = {"features": features, "labels": labels}
        step_dict = super().training_step(dict_input, to_cpu=False)

        loss = step_dict.pop("loss")
        step_dict["loss"] = self.poptorch.identity_loss(loss, reduction="mean")
        return step_dict

    def validation_step(self, batch, batch_idx) -> Dict[str, Any]:
        features, labels = batch["features"], batch["labels"]
        features, labels = self.squeeze_input_dims(features, labels)
        dict_input = {"features": features, "labels": labels}
        step_dict = super().validation_step(dict_input, to_cpu=False)

        return step_dict

    def test_step(self, batch, batch_idx) -> Dict[str, Any]:
        # Build a dictionary from the tuples
        features, labels = batch["features"], batch["labels"]
        features, labels = self.squeeze_input_dims(features, labels)
        dict_input = {"features": features, "labels": labels}
        step_dict = super().test_step(dict_input, to_cpu=False)

        return step_dict

    def predict_step(self, **inputs) -> Dict[str, Any]:
        # Build a dictionary from the tuples
        dict_input = inputs
        step_dict = super().predict_step(dict_input, to_cpu=False)

        return step_dict

    def on_validation_batch_end(self, outputs: Any, batch: Any, batch_idx: int) -> None:
        # convert data that will be tracked
        outputs = self.convert_from_fp16(outputs)
        super().on_validation_batch_end(outputs, batch, batch_idx)

    def evaluation_epoch_end(self, outputs: Any):
        outputs = self.convert_from_fp16(outputs)
        super().evaluation_epoch_end(outputs)

    def on_test_batch_end(self, outputs: Any, batch: Any, batch_idx: int) -> None:
        outputs = self.convert_from_fp16(outputs)
        super().on_test_batch_end(outputs, batch, batch_idx)

    def configure_optimizers(self, impl=None):
        if impl is None:
            dtype = self.precision_to_dtype(self.trainer.precision)
            impl = functools.partial(
                self.poptorch.optim.Adam,
                accum_type=dtype,
                first_order_momentum_accum_type=dtype,
                second_order_momentum_accum_type=torch.float,
            )
        return super().configure_optimizers(impl=impl)

    def squeeze_input_dims(self, features, labels):
        for key, tensor in features:
            if isinstance(tensor, torch.Tensor):
                features[key] = features[key].squeeze(0)

        for key in labels:
            labels[key] = labels[key].squeeze(0)

        return features, labels

    def convert_from_fp16(self, data: Any) -> Any:
        """
        Converts tensors from FP16 to FP32. Useful to convert the IPU program output data
        """
        if isinstance(data, collections.Sequence):
            for idx in range(len(data)):
                data[idx] = self.convert_from_fp16(data[idx])
        elif isinstance(data, collections.Mapping):
            for key in data:
                data[key] = self.convert_from_fp16(data[key])
        elif isinstance(data, torch.Tensor) and data.dtype == torch.float16:
            data = data.float()
        return data

    def _convert_features_dtype(self, feats):
        """
        Converts features to trainer precision rather than model precision.
        Necessary to run IPU on FP16.
        """
        dtype = self.precision_to_dtype(self.trainer.precision)

        # Convert features to dtype
        if isinstance(feats, torch.Tensor):
            feats = feats.to(dtype)
        elif isinstance(feats, (Data, Batch, dict)):
            for key, val in feats.items():
                if isinstance(val, torch.Tensor) and (val.is_floating_point()):
                    feats[key] = val.to(dtype=dtype)
        else:
            raise ValueError(f"Unsupported feats type `{type(feats)}` : {feats}")
        return feats

    def precision_to_dtype(self, precision):
        return torch.half if precision == "16-true" else torch.float

    def get_num_graphs(self, data: Batch):
        """
        IPU specific method to compute the number of graphs in a Batch,
        that considers gradient accumulation, multiple IPUs and multiple
        device iterations. Essential to estimate throughput in graphs/s.
        """
        num_graphs = torch.max(data.batch, dim=-1).values
        num_graphs = torch.sum(num_graphs)

        return num_graphs


# graphium/ipu/to_dense_batch.py

from typing import Optional, Tuple

import torch
from torch import Tensor
from torch_scatter import scatter_add


def to_sparse_batch(x: Tensor, mask_idx: Tensor):
    """
    Reverse function of `to_dense_batch`
    """
    return torch.index_select(x.reshape(-1, x.shape[-1]), 0, mask_idx)


def to_sparse_batch_from_packed(x: Tensor, pack_from_node_idx: Tensor):
    """
    Reverse function of `to_packed_dense_batch`
    """
    return x[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]]


def to_dense_batch(
    x: Tensor,
    batch: Optional[Tensor] = None,
    fill_value: float = 0.0,
    max_num_nodes_per_graph: Optional[int] = None,
    batch_size: Optional[int] = None,
    drop_nodes_last_graph=False,
) -> Tuple[Tensor, Tensor]:
    r"""Given a sparse batch of node features
    :math:`\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}` (with
    :math:`N_i` indicating the number of nodes in graph :math:`i`), creates a
    dense node feature tensor
    :math:`\mathbf{X} \in \mathbb{R}^{B \times N_{\max} \times F}` (with
    :math:`N_{\max} = \max_i^B N_i`).
    In addition, a mask of shape :math:`\mathbf{M} \in \{ 0, 1 \}^{B \times
    N_{\max}}` is returned, holding information about the existence of
    fake-nodes in the dense representation.

    Parameters:
        x: Node feature matrix
            :math:`\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}`.
        batch: Batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
            node to a specific example. Must be ordered. (default: :obj:`None`)
        fill_value: The value for invalid entries in the
            resulting dense output tensor. (default: :obj:`0`)
        max_num_nodes_per_graph: The size of the output node dimension.
            (default: :obj:`None`)
        batch_size: The batch size. (default: :obj:`None`)
        drop_nodes_last_graph: Whether to drop the nodes of the last graphs that exceed
            the `max_num_nodes_per_graph`. Useful when the last graph is a padding.

    :rtype: (:class:`Tensor`, :class:`BoolTensor`)
    """
    if batch is None and max_num_nodes_per_graph is None:
        mask = torch.ones(1, x.size(0), dtype=torch.bool, device=x.device)
        return x.unsqueeze(0), mask

    if batch is None:
        batch = x.new_zeros(x.size(0), dtype=torch.long)

    if batch_size is None:
        assert x.device.type != "ipu", (
            "When using the IPU the batch size must be "
            "provided during compilation instead of determined at runtime"
        )
        batch_size = int(batch.max()) + 1
    if x.device not in ["ipu", "xla"]:
        num_nodes = scatter_add(batch.new_ones(x.size(0)), batch, dim=0, dim_size=batch_size)
    else:
        # Can't use scatter_add here due to PopTorch bug, will be fixed in SDK 3.3
        arange = torch.arange(batch_size).unsqueeze(-1)
        num_nodes = batch.eq(arange).sum(dim=-1)
    cum_nodes = torch.cat([batch.new_zeros(1), num_nodes.cumsum(dim=0)])

    if max_num_nodes_per_graph is None:  # Must be provided on IPU
        max_num_nodes_per_graph = int(num_nodes.max())

    idx = torch.arange(batch.size(0), dtype=torch.long, device=x.device)
    idx = (idx - cum_nodes[batch]) + (batch * max_num_nodes_per_graph)

    size = [batch_size * max_num_nodes_per_graph] + list(x.size())[1:]

    out = x.new_full(size, fill_value)

    ##### CHANGES FROM PYG #####

    # In case the last graph represents padding. Drop the overflowing nodes.
    if drop_nodes_last_graph:
        num_nodes = num_nodes[:-1]
        idx[idx >= size[0]] = size[0] - 1

    # Raise error if num_nodes > max_num_nodes
    if x.device.type != "ipu":
        assert (
            num_nodes <= max_num_nodes_per_graph
        ).all(), f"Encountered graphs with {num_nodes.max()} nodes, greater than `max_num_nodes = {max_num_nodes_per_graph}`"

    out[idx] = x
    out = out.view([batch_size, max_num_nodes_per_graph] + list(x.size())[1:])

    # Create a zero-mask on the right device
    mask_sz = batch_size * max_num_nodes_per_graph
    if x.device.type in ("ipu", "xla"):
        mask = torch.zeros(mask_sz, dtype=torch.int32, device="cpu")
        mask = mask.to(x.device)
        # Can't use mask[idx] here due to PopTorch bug, will be fixed in SDK 3.3
        # mask[idx] = 1
        # mask = mask.bool()
        if drop_nodes_last_graph:
            num_nodes_with_padding = torch.cat((num_nodes, torch.tensor([0], dtype=torch.int32)), dim=0)
        else:
            num_nodes_with_padding = num_nodes

        arange = torch.arange(max_num_nodes_per_graph)
        mask = num_nodes_with_padding.unsqueeze(-1).gt(arange).flatten()

    else:
        mask = torch.zeros(mask_sz, dtype=torch.bool, device=x.device)
        mask[idx] = 1

    ##### END CHANGES FROM PYG #####

    mask = mask.view(batch_size, max_num_nodes_per_graph)

    return out, mask, idx  # Added `idx` as a return


def to_packed_dense_batch(
    x: Tensor,
    pack_from_node_idx: Tensor,
    pack_attn_mask: Tensor,
    fill_value: float = 0.0,
    max_num_nodes_per_pack: Optional[int] = None,
) -> Tuple[Tensor, Tensor]:
    r"""Given a sparse batch of node features
    :math:`\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}` (with
    :math:`N_i` indicating the number of nodes in graph :math:`i`), creates a
    dense node feature tensor
    :math:`\mathbf{X} \in \mathbb{R}^{B \times N_{\max} \times F}` (with
    :math:`N_{\max} = \max_i^B N_i`).
    In addition, a mask of shape :math:`\mathbf{M} \in \{ 0, 1 \}^{B \times
    N_{\max}}` is returned, holding information about the existence of
    fake-nodes in the dense representation.

    Parameters: # TODO: Update docstring
        x: Node feature matrix
            :math:`\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}`.
        batch: Batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
            node to a specific example. Must be ordered. (default: :obj:`None`)
        fill_value: The value for invalid entries in the
            resulting dense output tensor. (default: :obj:`0`)
        max_num_nodes_per_graph: The size of the output node dimension.
            (default: :obj:`None`)
        batch_size: The batch size. (default: :obj:`None`)
        drop_nodes_last_graph: Whether to drop the nodes of the last graphs that exceed
            the `max_num_nodes_per_graph`. Useful when the last graph is a padding.

    :rtype: (:class:`Tensor`, :class:`BoolTensor`)
    """

    if max_num_nodes_per_pack is None:  # Must be provided on IPU
        max_num_nodes_per_pack = pack_attn_mask.shape[-1]

    size = [pack_attn_mask[0], max_num_nodes_per_pack] + list(x.size())[1:]

    out = x.new_full(size, fill_value)
    out[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]] = x

    return out


# graphium/nn/base_graph_layer.py

import abc
from typing import Union, Callable, List, Optional, Mapping
from copy import deepcopy

import torch
import torch.nn as nn
from torch import Tensor, IntTensor
from torch_sparse import SparseTensor

from graphium.nn.base_layers import get_activation, DropPath
from graphium.utils.decorators import classproperty
import torch_geometric as pyg


class BaseGraphStructure:
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        activation: Union[str, Callable] = "relu",
        dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
        layer_idx: Optional[int] = None,
        layer_depth: Optional[int] = None,
        droppath_rate: float = 0.0,
    ):
        r"""
        Abstract class used to standardize the implementation of Pyg layers
        in the current library. It will allow a network to seemlesly swap between
        different GNN layers by better understanding the expected inputs
        and outputs.

        Parameters:

            in_dim:
                Input feature dimensions of the layer

            out_dim:
                Output feature dimensions of the layer

            activation:
                activation function to use in the layer

            dropout:
                The ratio of units to dropout. Must be between 0 and 1

            normalization:
                Normalization to use. Choices:

                - "none" or `None`: No normalization
                - "batch_norm": Batch normalization
                - "layer_norm": Layer normalization
                - `Callable`: Any callable function

            layer_idx:
                The index of the current layer

            layer_depth:
                The total depth (number of layers) associated to this specific layer

            droppath_rate:
                stochastic depth drop rate, between 0 and 1, see https://arxiv.org/abs/1603.09382
        """

        super().__init__()

        # Basic attributes
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.normalization = normalization
        self.dropout = dropout
        self.activation = activation
        self.layer_idx = layer_idx
        self.layer_depth = layer_depth
        self.droppath_rate = droppath_rate
        self._max_num_nodes_per_graph = None
        self._max_num_edges_per_graph = None

    def _initialize_activation_dropout_norm(self):
        if not isinstance(self, nn.Module):
            raise TypeError(
                "This function requires the current object to be an `nn.Module`. Use multi-inheritance or the class `BaseGraphModule` instead"
            )

        # Build the layers
        self.activation_layer = get_activation(self.activation)
        self.dropout_layer = self._parse_dropout(self.dropout)

        self.norm_layer = self._parse_norm(self.normalization)
        self.droppath_layer = self._parse_droppath(self.droppath_rate)

    def _parse_dropout(self, dropout):
        if callable(dropout):
            return deepcopy(dropout)
        elif dropout > 0:
            return nn.Dropout(p=dropout)
        return

    def _parse_droppath(self, droppath_rate):
        if droppath_rate == 0:
            return

        droppath_rate = DropPath.get_stochastic_drop_rate(droppath_rate, self.layer_idx, self.layer_depth)
        return DropPath(drop_rate=droppath_rate)

    def _parse_norm(self, normalization, dim=None):
        if dim is None:
            dim = self.out_dim * self.out_dim_factor
        if normalization is None or normalization == "none":
            parsed_norm = None
        elif callable(normalization):
            parsed_norm = deepcopy(normalization)
        elif normalization == "batch_norm":
            parsed_norm = nn.BatchNorm1d(dim)
        elif normalization == "layer_norm":
            parsed_norm = nn.LayerNorm(dim)
        else:
            raise ValueError(
                f"Undefined normalization `{normalization}`, must be `None`, `Callable`, 'batch_norm', 'layer_norm', 'none'"
            )
        return parsed_norm

    def apply_norm_activation_dropout(
        self,
        feat: Tensor,
        normalization: bool = True,
        activation: bool = True,
        dropout: bool = True,
        droppath: bool = True,
        batch_idx: Optional[IntTensor] = None,
        batch_size: Optional[int] = None,
    ):
        r"""
        Apply the different normalization and the dropout to the
        output layer.

        Parameters:

            feat:
                Feature tensor, to be normalized

            batch_idx

            normalization:
                Whether to apply the normalization

            activation:
                Whether to apply the activation layer

            dropout:
                Whether to apply the dropout layer

            droppath:
                Whether to apply the DropPath layer

        Returns:

            feat:
                Normalized and dropped-out features

        """

        if normalization and (self.norm_layer is not None):
            feat = self.norm_layer(feat)

        if activation and (self.activation_layer is not None):
            feat = self.activation_layer(feat)

        if dropout and (self.dropout_layer is not None):
            feat = self.dropout_layer(feat)

        if droppath and (self.droppath_layer is not None):
            feat = self.droppath_layer(feat, batch_idx=batch_idx, batch_size=batch_size)

        return feat

    @classproperty
    def layer_supports_edges(cls) -> bool:
        r"""
        Abstract method. Return a boolean specifying if the layer type
        supports output edges edges or not.

        Returns:

            bool:
                Whether the layer supports the use of edges
        """
        ...

    @property
    @abc.abstractmethod
    def layer_inputs_edges(self) -> bool:
        r"""
        Abstract method. Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_input_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:

            bool:
                Whether the layer uses input edges in the forward pass
        """
        ...

    @property
    @abc.abstractmethod
    def layer_outputs_edges(self) -> bool:
        r"""
        Abstract method. Return a boolean specifying if the layer type
        uses edges as input or not.
        It is different from ``layer_supports_output_edges`` since a layer that
        supports edges can decide to not use them.

        Returns:

            bool:
                Whether the layer outputs edges in the forward pass
        """
        ...

    @property
    @abc.abstractmethod
    def out_dim_factor(self) -> int:
        r"""
        Abstract method.
        Get the factor by which the output dimension is multiplied for
        the next layer.

        For standard layers, this will return ``1``.

        But for others, such as ``GatLayer``, the output is the concatenation
        of the outputs from each head, so the out_dim gets multiplied by
        the number of heads, and this function should return the number
        of heads.

        Returns:

            int:
                The factor that multiplies the output dimensions
        """
        ...

    @property
    def max_num_nodes_per_graph(self) -> Optional[int]:
        """
        Get the maximum number of nodes per graph. Useful for reshaping a compiled model (IPU)
        """
        return self._max_num_nodes_per_graph

    @max_num_nodes_per_graph.setter
    def max_num_nodes_per_graph(self, value: Optional[int]):
        """
        Set the maximum number of nodes per graph. Useful for reshaping a compiled model (IPU)
        """
        if value is not None:
            assert isinstance(value, int) and (
                value > 0
            ), f"Value should be a positive integer, provided f{value} of type {type(value)}"
        self._max_num_nodes_per_graph = value

    @property
    def max_num_edges_per_graph(self) -> Optional[int]:
        """
        Get the maximum number of nodes per graph. Useful for reshaping a compiled model (IPU)
        """
        return self._max_num_edges_per_graph

    @max_num_edges_per_graph.setter
    def max_num_edges_per_graph(self, value: Optional[int]):
        """
        Set the maximum number of nodes per graph. Useful for reshaping a compiled model (IPU)
        """
        if value is not None:
            assert isinstance(value, int) and (
                value > 0
            ), f"Value should be a positive integer, provided f{value} of type {type(value)}"
        self._max_num_edges_per_graph = value

    def __repr__(self):
        r"""
        Controls how the class is printed
        """
        f = self.out_dim_factor
        out_dim_f_print = "" if f == 1 else f" * {f}"
        return f"{self.__class__.__name__}({self.in_dim} -> {self.out_dim}{out_dim_f_print}, activation={self.activation})"


class BaseGraphModule(BaseGraphStructure, nn.Module):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        activation: Union[str, Callable] = "relu",
        dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
        layer_idx: Optional[int] = None,
        layer_depth: Optional[int] = None,
        droppath_rate: float = 0.0,
    ):
        r"""
        Abstract class used to standardize the implementation of Pyg layers
        in the current library. It will allow a network to seemlesly swap between
        different GNN layers by better understanding the expected inputs
        and outputs.

        Parameters:

            in_dim:
                Input feature dimensions of the layer

            out_dim:
                Output feature dimensions of the layer

            activation:
                activation function to use in the layer

            dropout:
                The ratio of units to dropout. Must be between 0 and 1

            normalization:
                Normalization to use. Choices:

                - "none" or `None`: No normalization
                - "batch_norm": Batch normalization
                - "layer_norm": Layer normalization
                - `Callable`: Any callable function

            layer_idx:
                The index of the current layer

            layer_depth:
                The total depth (number of layers) associated to this specific layer

            droppath_rate:
                stochastic depth drop rate, between 0 and 1, see https://arxiv.org/abs/1603.09382
        """

        super().__init__(
            in_dim=in_dim,
            out_dim=out_dim,
            normalization=normalization,
            dropout=dropout,
            activation=activation,
            layer_idx=layer_idx,
            layer_depth=layer_depth,
            droppath_rate=droppath_rate,
        )

        self._initialize_activation_dropout_norm()


def check_intpus_allow_int(obj, edge_index, size):
    """
    Overwrite the __check_input__ to allow for int32 and int16
    TODO: Remove when PyG and pytorch supports int32.
    """
    the_size: List[Optional[int]] = [None, None]

    if isinstance(edge_index, Tensor):
        # These 3 lines are different. They check for more int types and avoid overflow
        assert edge_index.dtype in (torch.long, torch.int64, torch.int32, torch.int16)
        # assert edge_index.min() >= 0
        # assert edge_index.max() < torch.iinfo(edge_index.dtype).max

        assert edge_index.dim() == 2
        assert edge_index.size(0) == 2
        if size is not None:
            the_size[0] = size[0]
            the_size[1] = size[1]
        return the_size

    elif isinstance(edge_index, SparseTensor):
        if obj.flow == "target_to_source":
            raise ValueError(
                (
                    'Flow direction "target_to_source" is invalid for '
                    "message propagation via `torch_sparse.SparseTensor`. If "
                    "you really want to make use of a reverse message "
                    "passing flow, pass in the transposed sparse tensor to "
                    "the message passing module, e.g., `adj_t.t()`."
                )
            )
        the_size[0] = edge_index.sparse_size(1)
        the_size[1] = edge_index.sparse_size(0)
        return the_size

    raise ValueError(
        (
            "`MessagePassing.propagate` only supports `torch.LongTensor` of "
            "shape `[2, num_messages]` or `torch_sparse.SparseTensor` for "
            "argument `edge_index`."
        )
    )


# graphium/nn/base_layers.py

from typing import Union, Callable, Optional, Type, Tuple, Iterable
from copy import deepcopy
from loguru import logger

import inspect

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor, IntTensor
import mup.init as mupi
from mup import set_base_shapes, MuReadout
from torch.nn.functional import linear

from graphium.ipu.ipu_utils import is_running_on_ipu

SUPPORTED_ACTIVATION_MAP = {
    "ReLU",
    "Sigmoid",
    "Tanh",
    "ELU",
    "SELU",
    "GLU",
    "GELU",
    "LeakyReLU",
    "Softplus",
    "None",
}


def get_activation(activation: Union[type(None), str, Callable]) -> Optional[Callable]:
    r"""
    returns the activation function represented by the input string

    Parameters:
        activation: Callable, `None`, or string with value:
            "none", "ReLU", "Sigmoid", "Tanh", "ELU", "SELU", "GLU", "GELU", "LeakyReLU", "Softplus"

    Returns:
        Callable or None: The activation function
    """
    if (activation is not None) and callable(activation):
        # activation is already a function
        return activation

    if (activation is None) or (activation.lower() == "none"):
        return None

    # search in SUPPORTED_ACTIVATION_MAP a torch.nn.modules.activation
    activation = [x for x in SUPPORTED_ACTIVATION_MAP if activation.lower() == x.lower()]
    assert len(activation) == 1 and isinstance(
        activation[0], str
    ), f"Unhandled activation function {activation} of type {type(activation)}"
    activation = activation[0]

    return vars(torch.nn.modules.activation)[activation]()


def get_activation_str(activation: Union[type(None), str, Callable]) -> str:
    r"""
    returns the string related to the activation function

    Parameters:
        activation: Callable, `None`, or string with value:
            "none", "ReLU", "Sigmoid", "Tanh", "ELU", "SELU", "GLU", "LeakyReLU", "Softplus"

    Returns:
        The name of the activation function
    """

    if isinstance(activation, str):
        return activation

    if activation is None:
        return "None"

    if isinstance(activation, Callable):
        return activation.__class__._get_name(activation)
    else:
        raise ValueError(f"Unhandled activation function {activation} of type {type(activation)}")


def get_norm(normalization: Union[Type[None], str, Callable], dim: Optional[int] = None):
    r"""
    returns the normalization function represented by the input string

    Parameters:
        normalization: Callable, `None`, or string with value:
            "none", "batch_norm", "layer_norm"
        dim: Dimension where to apply the norm. Mandatory for 'batch_norm' and 'layer_norm'

    Returns:
        Callable or None: The normalization function
    """
    parsed_norm = None
    if (normalization is None) or (normalization in ["none", "NoneType"]):
        pass
    elif callable(normalization):
        parsed_norm = normalization
    elif normalization in ["batch_norm", "BatchNorm1d"]:
        parsed_norm = nn.BatchNorm1d(dim)
    elif normalization in ["layer_norm", "LayerNorm"]:
        parsed_norm = nn.LayerNorm(dim)
    else:
        raise ValueError(
            f"Undefined normalization `{normalization}`, must be `None`, `Callable`, 'batch_norm', 'layer_norm', 'none'"
        )
    return deepcopy(parsed_norm)


class MultiheadAttentionMup(nn.MultiheadAttention):
    """
    Modifying the MultiheadAttention to work with the muTransfer paradigm.
    The layers are initialized using the mup package.
    The `_scaled_dot_product_attention` normalizes the attention matrix with `1/d` instead of `1/sqrt(d)`
    The biased self-attention option is added to have 3D attention bias.
    """

    def __init__(self, biased_attention, **kwargs):
        super().__init__(**kwargs)
        self.biased_attention = biased_attention

    def _reset_parameters(self):
        set_base_shapes(self, None, rescale_params=False)  # Set the shapes of the tensors, useful for mup
        if self._qkv_same_embed_dim:
            mupi.xavier_uniform_(self.in_proj_weight)
        else:
            mupi.xavier_uniform_(self.q_proj_weight)
            mupi.xavier_uniform_(self.k_proj_weight)
            mupi.xavier_uniform_(self.v_proj_weight)

        if self.in_proj_bias is not None:
            nn.init.constant_(self.in_proj_bias, 0.0)
            nn.init.constant_(self.out_proj.bias, 0.0)
        if self.bias_k is not None:
            mupi.xavier_normal_(self.bias_k)
        if self.bias_v is not None:
            mupi.xavier_normal_(self.bias_v)

    def forward(
        self,
        query: Tensor,
        key: Tensor,
        value: Tensor,
        key_padding_mask: Optional[Tensor] = None,
        attn_bias: Optional[Tensor] = None,
        precision: Optional[str] = "32",
        *args,
        **kwargs,
    ) -> Tuple[Tensor, Optional[Tensor]]:
        # attn_bias [batch, num_heads, nodes, nodes]
        if not self.biased_attention or attn_bias is None:
            attn_bias = 0.0
        # assuming source and target have the same sequence length (homogeneous graph attention)
        batch, nodes, hidden = query.size()
        assert (
            hidden == self.embed_dim
        ), f"query hidden dimension {hidden} != embed_dim {self.embed_dim} in class"
        head_dim = self.embed_dim // self.num_heads
        assert head_dim * self.num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
        scaling_factor = 1 / head_dim  # use head_dim instead of (head_dim**0.5) for mup
        b_q, b_k, b_v = self.in_proj_bias.chunk(3)
        q_proj_weight, k_proj_weight, v_proj_weight = self.in_proj_weight.chunk(3)
        # [batch, num_heads, nodes, head_size]
        q = linear(query, q_proj_weight, b_q).view(batch, nodes, self.num_heads, -1).transpose(1, 2)
        # [batch, num_heads, nodes, head_size]
        k = linear(key, k_proj_weight, b_k).view(batch, nodes, self.num_heads, -1).transpose(1, 2)
        # [batch, num_heads, nodes, head_size]
        v = linear(value, v_proj_weight, b_v).view(batch, nodes, self.num_heads, -1).transpose(1, 2)
        q = q * scaling_factor
        # [batch, num_heads, nodes, nodes]
        attn_weights = q @ k.transpose(-1, -2)
        # [batch, num_heads, nodes, nodes]
        attn_weights += attn_bias
        key_padding_mask_value = float("-inf") if precision == "32" else -10000
        # key_padding_mask: [batch, 1, 1, nodes]
        if key_padding_mask is not None:
            masked_attn_weights = attn_weights.masked_fill(
                key_padding_mask.unsqueeze(1).unsqueeze(2),
                key_padding_mask_value,
            )
        else:
            masked_attn_weights = attn_weights
        masked_attn_weights = F.softmax(masked_attn_weights, dim=-1)
        attn_probs = F.dropout(masked_attn_weights, p=self.dropout, training=self.training)
        # [batch, num_heads, nodes, nodes] * [batch, num_heads, nodes, head_size] -> [batch, num_heads, nodes, head_size]
        attn = attn_probs @ v
        # [batch, nodes, embd_dim]
        attn = attn.transpose(1, 2).contiguous().view(batch, nodes, self.embed_dim)
        # [batch, nodes, embd_dim]
        out = (self.out_proj(attn), None)
        return out


class TransformerEncoderLayerMup(nn.TransformerEncoderLayer):
    r"""
    Modified version of ``torch.nn.TransformerEncoderLayer`` that uses :math:`1/n`-scaled attention
    for compatibility with muP (as opposed to the original :math:`1/\sqrt{n}` scaling factor)
    Arguments are the same as ``torch.nn.TransformerEncoderLayer``.
    """

    def __init__(self, biased_attention, *args, **kwargs) -> None:
        super(TransformerEncoderLayerMup, self).__init__(*args, **kwargs)

        # Extract arguments passed to __init__ as a dictionary
        signature = inspect.signature(nn.TransformerEncoderLayer.__init__)

        # `self` needs to passed, which makes things tricky, but using this object seems fine for now
        bound_signature = signature.bind(self, *args, **kwargs)
        bound_signature.apply_defaults()

        mha_names = ["embed_dim", "num_heads", "dropout", "batch_first", "device", "dtype"]
        transformer_names = ["d_model", "nhead", "dropout", "batch_first", "device", "dtype"]

        # Override self attention to use muP
        self.self_attn = MultiheadAttentionMup(
            biased_attention,
            **{
                mha_name: bound_signature.arguments[transformer_name]
                for mha_name, transformer_name in zip(mha_names, transformer_names)
            },
        )


class MuReadoutGraphium(MuReadout):
    """
    PopTorch-compatible replacement for `mup.MuReadout`

    Not quite a drop-in replacement for `mup.MuReadout` - you need to specify
    `base_width`.

    Set `base_width` to width of base model passed to `mup.set_base_shapes`
    to get same results on IPU and CPU. Should still "work" with any other
    value, but won't give the same results as CPU
    """

    def __init__(self, in_features, *args, **kwargs):
        super().__init__(in_features, *args, **kwargs)
        self.base_width = in_features

    @property
    def absolute_width(self):
        return float(self.in_features)

    @property
    def base_width(self):
        return self._base_width

    @base_width.setter
    def base_width(self, val):
        if val is None:
            return
        assert isinstance(
            val, (int, torch.int, torch.long)
        ), f"`base_width` must be None, int or long, provided {val} of type {type(val)}"
        self._base_width = val

    def width_mult(self):
        return self.absolute_width / self.base_width


class MuReadoutGraphium(MuReadout):
    """
    PopTorch-compatible replacement for `mup.MuReadout`

    Not quite a drop-in replacement for `mup.MuReadout` - you need to specify
    `base_width`.

    Set `base_width` to width of base model passed to `mup.set_base_shapes`
    to get same results on IPU and CPU. Should still "work" with any other
    value, but won't give the same results as CPU
    """

    def __init__(self, in_features, *args, **kwargs):
        super().__init__(in_features, *args, **kwargs)
        self.base_width = in_features

    @property
    def absolute_width(self):
        return float(self.in_features)

    @property
    def base_width(self):
        return self._base_width

    @base_width.setter
    def base_width(self, val):
        if val is None:
            return
        assert isinstance(
            val, (int, torch.int, torch.long)
        ), f"`base_width` must be None, int or long, provided {val} of type {type(val)}"
        self._base_width = val

    def width_mult(self):
        return self.absolute_width / self.base_width


class FCLayer(nn.Module):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        activation: Union[str, Callable] = "relu",
        dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
        bias: bool = True,
        init_fn: Optional[Callable] = None,
        is_readout_layer: bool = False,
        droppath_rate: float = 0.0,
    ):
        r"""
        A simple fully connected and customizable layer. This layer is centered around a `torch.nn.Linear` module.
        The order in which transformations are applied is:

        - Dense Layer
        - Activation
        - Dropout (if applicable)
        - Batch Normalization (if applicable)

        Parameters:
            in_dim:
                Input dimension of the layer (the `torch.nn.Linear`)
            out_dim:
                Output dimension of the layer.
            dropout:
                The ratio of units to dropout. No dropout by default.
            activation:
                Activation function to use.
            normalization:
                Normalization to use. Choices:

                - "none" or `None`: No normalization
                - "batch_norm": Batch normalization
                - "layer_norm": Layer normalization
                - `Callable`: Any callable function
            bias:
                Whether to enable bias in for the linear layer.
            init_fn:
                Initialization function to use for the weight of the layer. Default is
                $$\mathcal{U}(-\sqrt{k}, \sqrt{k})$$ with $$k=\frac{1}{ \text{in_dim}}$$
            is_readout_layer: Whether the layer should be treated as a readout layer by replacing of `torch.nn.Linear`
                by `mup.MuReadout` from the muTransfer method https://github.com/microsoft/mup

            droppath_rate:
                stochastic depth drop rate, between 0 and 1, see https://arxiv.org/abs/1603.09382
        Attributes:
            dropout (int):
                The ratio of units to dropout.
            normalization (None or Callable):
                Normalization layer
            linear (`torch.nn.Linear`):
                The linear layer
            activation (`torch.nn.Module`):
                The activation layer
            init_fn (Callable):
                Initialization function used for the weight of the layer
            in_dim (int):
                Input dimension of the linear layer
            out_dim (int):
                Output dimension of the linear layer
        """

        super().__init__()

        self.__params = locals()
        del self.__params["__class__"]
        del self.__params["self"]

        # Basic parameters
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.bias = bias
        self.dropout = None
        self.normalization = get_norm(normalization, dim=out_dim)

        # Dropout and activation
        if dropout:
            self.dropout = nn.Dropout(p=dropout)
        self.activation = get_activation(activation)

        self.drop_path = None
        if droppath_rate > 0:
            self.drop_path = DropPath(droppath_rate)

        # Linear layer, or MuReadout layer
        if not is_readout_layer:
            self.linear = nn.Linear(in_dim, out_dim, bias=bias)
        else:
            self.linear = MuReadoutGraphium(in_dim, out_dim, bias=bias)

            # Warn user in case of weird parameters
            if self.normalization is not None:
                logger.warning(
                    f"Normalization is not `None` for the readout layer. Provided {self.normalization}"
                )
            if (self.dropout is not None) and (self.dropout.p > 0):
                logger.warning(f"Dropout is not `None` or `0` for the readout layer. Provided {self.dropout}")

        # Define the initialization function based on `muTransfer`, and reset the parameters
        self.init_fn = init_fn if init_fn is not None else mupi.xavier_uniform_
        self.reset_parameters()

    def reset_parameters(self, init_fn=None):
        """
        Reset the parameters of the linear layer using the `init_fn`.
        """
        set_base_shapes(self, None, rescale_params=False)  # Set the shapes of the tensors, useful for mup
        init_fn = init_fn or self.init_fn
        if init_fn is not None:
            init_fn(self.linear.weight)
        if self.bias:
            self.linear.bias.data.zero_()

    def forward(self, h: torch.Tensor) -> torch.Tensor:
        r"""
        Apply the FC layer on the input features.

        Parameters:

            h: `torch.Tensor[..., Din]`:
                Input feature tensor, before the FC.
                `Din` is the number of input features

        Returns:

            `torch.Tensor[..., Dout]`:
                Output feature tensor, after the FC.
                `Dout` is the number of output features

        """

        if torch.prod(torch.as_tensor(h.shape[:-1])) == 0:
            h = torch.zeros(
                list(h.shape[:-1]) + [self.linear.out_features],
                device=h.device,
                dtype=h.dtype,
            )
            return h

        h = self.linear(h)

        if self.normalization is not None:
            if h.shape[1] != self.out_dim:
                h = self.normalization(h.transpose(1, 2)).transpose(1, 2)
            else:
                h = self.normalization(h)

        if self.dropout is not None:
            h = self.dropout(h)
        if self.activation is not None:
            h = self.activation(h)
        if self.drop_path is not None:
            h = self.drop_path(h)

        return h

    @property
    def in_channels(self) -> int:
        r"""
        Get the input channel size. For compatibility with PyG.
        """
        return self.in_dim

    @property
    def out_channels(self) -> int:
        r"""
        Get the output channel size. For compatibility with PyG.
        """
        return self.out_dim

    def __repr__(self):
        return f"{self.__class__.__name__}({self.in_dim} -> {self.out_dim}, activation={self.activation})"


class MLP(nn.Module):
    def __init__(
        self,
        in_dim: int,
        hidden_dims: Union[Iterable[int], int],
        out_dim: int,
        depth: int,
        activation: Union[str, Callable] = "relu",
        last_activation: Union[str, Callable] = "none",
        dropout: float = 0.0,
        last_dropout: float = 0.0,
        normalization: Union[Type[None], str, Callable] = "none",
        last_normalization: Union[Type[None], str, Callable] = "none",
        first_normalization: Union[Type[None], str, Callable] = "none",
        last_layer_is_readout: bool = False,
        droppath_rate: float = 0.0,
        constant_droppath_rate: bool = True,
    ):
        r"""
        Simple multi-layer perceptron, built of a series of FCLayers

        Parameters:
            in_dim:
                Input dimension of the MLP
            hidden_dims:
                Either an integer specifying all the hidden dimensions,
                or a list of dimensions in the hidden layers.
            out_dim:
                Output dimension of the MLP.
            depth:
                If `hidden_dims` is an integer, `depth` is 1 + the number of
                hidden layers to use.
                If `hidden_dims` is a list, then
                `depth` must be `None` or equal to `len(hidden_dims) + 1`
            activation:
                Activation function to use in all the layers except the last.
                if `layers==1`, this parameter is ignored
            last_activation:
                Activation function to use in the last layer.
            dropout:
                The ratio of units to dropout. Must be between 0 and 1
            normalization:
                Normalization to use. Choices:

                - "none" or `None`: No normalization
                - "batch_norm": Batch normalization
                - "layer_norm": Layer normalization in the hidden layers.
                - `Callable`: Any callable function

                if `layers==1`, this parameter is ignored
            last_normalization:
                Norrmalization to use **after the last layer**. Same options as `normalization`.
            first_normalization:
                Norrmalization to use in **before the first layer**. Same options as `normalization`.
            last_dropout:
                The ratio of units to dropout at the last layer.
            last_layer_is_readout: Whether the last layer should be treated as a readout layer.
                Allows to use the `mup.MuReadout` from the muTransfer method https://github.com/microsoft/mup
            droppath_rate:
                stochastic depth drop rate, between 0 and 1.
                See https://arxiv.org/abs/1603.09382
            constant_droppath_rate:
                If `True`, drop rates will remain constant accross layers.
                Otherwise, drop rates will vary stochastically.
                See `DropPath.get_stochastic_drop_rate`
        """

        super().__init__()

        self.in_dim = in_dim
        self.out_dim = out_dim

        # Parse the hidden dimensions and depth
        if isinstance(hidden_dims, int):
            self.hidden_dims = [hidden_dims] * (depth - 1)
        else:
            self.hidden_dims = list(hidden_dims)
            assert (depth is None) or (
                depth == len(self.hidden_dims) + 1
            ), "Mismatch between the provided network depth from `hidden_dims` and `depth`"
        self.depth = len(self.hidden_dims) + 1

        # Parse the normalization
        self.first_normalization = get_norm(first_normalization, dim=in_dim)

        all_dims = [in_dim] + self.hidden_dims + [out_dim]
        fully_connected = []
        if depth == 0:
            self.fully_connected = None
            return
        else:
            for ii in range(depth):
                if ii < (depth - 1):
                    # Define the parameters for all intermediate layers
                    this_activation = activation
                    this_normalization = normalization
                    this_dropout = dropout
                    is_readout_layer = False
                else:
                    # Define the parameters for the last layer
                    this_activation = last_activation
                    this_normalization = last_normalization
                    this_dropout = last_dropout
                    is_readout_layer = last_layer_is_readout

                if constant_droppath_rate:
                    this_drop_rate = droppath_rate
                else:
                    this_drop_rate = DropPath.get_stochastic_drop_rate(droppath_rate, ii, depth)

                # Add a fully-connected layer
                fully_connected.append(
                    FCLayer(
                        all_dims[ii],
                        all_dims[ii + 1],
                        activation=this_activation,
                        normalization=this_normalization,
                        dropout=this_dropout,
                        is_readout_layer=is_readout_layer,
                        droppath_rate=this_drop_rate,
                    )
                )

        self.fully_connected = nn.Sequential(*fully_connected)

    def forward(self, h: torch.Tensor) -> torch.Tensor:
        r"""
        Apply the MLP on the input features.

        Parameters:

            h: `torch.Tensor[..., Din]`:
                Input feature tensor, before the MLP.
                `Din` is the number of input features

        Returns:

            `torch.Tensor[..., Dout]`:
                Output feature tensor, after the MLP.
                `Dout` is the number of output features

        """
        if self.first_normalization is not None:
            h = self.first_normalization(h)
        if self.fully_connected is not None:
            h = self.fully_connected(h)
        return h

    @property
    def in_features(self):
        return self.in_dim

    def __getitem__(self, idx: int) -> nn.Module:
        return self.fully_connected[idx]

    def __repr__(self):
        r"""
        Controls how the class is printed
        """
        return self.__class__.__name__ + " (" + str(self.in_dim) + " -> " + str(self.out_dim) + ")"


class GRU(nn.Module):
    def __init__(self, in_dim: int, hidden_dim: int):
        r"""
        Wrapper class for the GRU used by the GNN framework, nn.GRU is used for the Gated Recurrent Unit itself

        Parameters:
            in_dim:
                Input dimension of the GRU layer
            hidden_dim:
                Hidden dimension of the GRU layer.
        """

        super().__init__()
        self.in_dim = in_dim
        self.hidden_dim = hidden_dim
        self.gru = nn.GRU(in_dim=in_dim, hidden_dim=hidden_dim)

    def forward(self, x, y):
        r"""
        Parameters:
            x:  `torch.Tensor[B, N, Din]`
                where Din <= in_dim (difference is padded)
            y:  `torch.Tensor[B, N, Dh]`
                where Dh <= hidden_dim (difference is padded)

        Returns:
            torch.Tensor: `torch.Tensor[B, N, Dh]`

        """
        assert x.shape[-1] <= self.in_dim and y.shape[-1] <= self.hidden_dim

        (B, N, _) = x.shape
        x = x.reshape(1, B * N, -1).contiguous()
        y = y.reshape(1, B * N, -1).contiguous()

        # padding if necessary
        if x.shape[-1] < self.in_dim:
            x = F.pad(input=x, pad=[0, self.in_dim - x.shape[-1]], mode="constant", value=0)
        if y.shape[-1] < self.hidden_dim:
            y = F.pad(
                input=y,
                pad=[0, self.hidden_dim - y.shape[-1]],
                mode="constant",
                value=0,
            )

        x = self.gru(x, y)[1]
        x = x.reshape(B, N, -1)
        return x


class DropPath(nn.Module):
    def __init__(self, drop_rate: float):
        r"""
        DropPath class for stochastic depth
        Deep Networks with Stochastic Depth
        Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra and Kilian Weinberger
        https://arxiv.org/abs/1603.09382

        Parameters:
            drop_rate:
                Drop out probability
        """

        super().__init__()
        self.drop_rate = drop_rate

    @staticmethod
    def get_stochastic_drop_rate(
        drop_rate: float, layer_idx: Optional[int] = None, layer_depth: Optional[int] = None
    ):
        """
        Get the stochastic drop rate from the nominal drop rate, the layer index, and the layer depth.

        `return drop_rate * (layer_idx / (layer_depth - 1))`

        Parameters:
            drop_rate:
                Drop out nominal probability

            layer_idx:
                The index of the current layer

            layer_depth:
                The total depth (number of layers) associated to this specific layer

        """
        if drop_rate == 0:
            return 0
        else:
            assert (layer_idx is not None) and (
                layer_depth is not None
            ), f"layer_idx={layer_idx} and layer_depth={layer_depth} should be integers when `droppath_rate>0`"
            return drop_rate * (layer_idx / (layer_depth - 1))

    def forward(
        self,
        input: Tensor,
        batch_idx: IntTensor,
        batch_size: Optional[int] = None,
    ) -> Tensor:
        r"""
        Parameters:
            input:  `torch.Tensor[total_num_nodes, hidden]`
            batch: batch attribute of the batch object, batch.batch
            batch_size: The batch size. Must be provided when working on IPU

        Returns:
            torch.Tensor: `torch.Tensor[total_num_nodes, hidde]`

        """
        on_ipu = is_running_on_ipu()

        if self.drop_rate > 0:
            keep_prob = 1 - self.drop_rate

            # Parse the batch size
            if batch_size is None:
                if on_ipu:
                    raise ValueError(
                        "When using the IPU the batch size must be "
                        "provided during compilation instead of determined at runtime"
                    )
                else:
                    batch_size = int(batch_idx.max()) + 1

            # mask shape: [num_graphs, 1]
            mask = input.new_empty(batch_size, 1).bernoulli_(keep_prob)
            # if on_ipu, the last graph is a padded fake graph
            if on_ipu:
                mask[-1] = 0
            # using gather to extend mask to [total_num_nodes, 1]
            node_mask = mask[batch_idx]
            if keep_prob == 0:
                # avoid dividing by 0
                input_scaled = input
            else:
                input_scaled = input / keep_prob
            out = input_scaled * node_mask
        else:
            out = input
        return out


# graphium/nn/residual_connections.py

"""
Different types of residual connections, including None, Simple (ResNet-like),
Concat and DenseNet
"""

import abc
import torch
import torch.nn as nn
from typing import List, Union, Callable

from graphium.nn.base_layers import FCLayer
from graphium.utils.decorators import classproperty


class ResidualConnectionBase(nn.Module):
    def __init__(self, skip_steps: int = 1):
        r"""
        Abstract class for the residual connections. Using this class,
        we implement different types of residual connections, such as
        the ResNet, weighted-ResNet, skip-concat and DensNet.

        The following methods must be implemented in a children class

        - ``h_dim_increase_type()``
        - ``has_weights()``

        Parameters:

            skip_steps: int
                The number of steps to skip between the residual connections.
                If `1`, all the layers are connected. If `2`, half of the
                layers are connected.
        """

        super().__init__()
        self.skip_steps = skip_steps

    def _bool_apply_skip_step(self, step_idx: int):
        r"""
        Whether to apply the skip connection, depending on the
        ``step_idx`` and ``self.skip_steps``.

        Parameters:

            step_idx: int
                The current layer step index.

        """
        return (self.skip_steps != 0) and ((step_idx % self.skip_steps) == 0)

    def __repr__(self):
        r"""
        Controls how the class is printed
        """
        return f"{self.__class__.__name__}(skip_steps={self.skip_steps})"

    @classproperty
    @abc.abstractmethod
    def h_dim_increase_type(cls):
        r"""
        How does the dimension of the output features increases after each layer?

        Returns:

            h_dim_increase_type: None or str
                - ``None``: The dimension of the output features do not change at each layer.
                E.g. ResNet.

                - "previous": The dimension of the output features is the concatenation of
                the previous layer with the new layer.

                - "cumulative": The dimension of the output features is the concatenation
                of all previous layers.

        """
        ...

    def get_true_out_dims(self, out_dims: List) -> List:
        r"""
        find the true output dimensions
        Parameters:
            out_dims: List
        Returns:
            true_out_dims: List
        """
        true_out_dims = [out_dims[0]]
        out_dims_at_skip = [out_dims[0]]
        for ii in range(1, len(out_dims) - 1):
            # For the `None` type, don't change the output dims
            if self.h_dim_increase_type is None:
                true_out_dims.append(out_dims[ii])

            # For the "previous" type, add the previous layers when the skip connection applies
            elif self.h_dim_increase_type == "previous":
                if self._bool_apply_skip_step(step_idx=ii):
                    true_out_dims.append(out_dims[ii] + out_dims_at_skip[-1])
                    out_dims_at_skip.append(out_dims[ii])
                else:
                    true_out_dims.append(out_dims[ii])

            # For the "cumulative" type, add all previous layers when the skip connection applies
            elif self.h_dim_increase_type == "cumulative":
                if self._bool_apply_skip_step(step_idx=ii):
                    true_out_dims.append(out_dims[ii] + out_dims_at_skip[-1])
                    out_dims_at_skip.append(true_out_dims[ii])
                else:
                    true_out_dims.append(out_dims[ii])
            else:
                raise ValueError(f"undefined value: {self.h_dim_increase_type}")

        return true_out_dims

    @classproperty
    @abc.abstractmethod
    def has_weights(cls):
        r"""
        Returns:

            has_weights: bool
                Whether the residual connection uses weights

        """
        ...


class ResidualConnectionNone(ResidualConnectionBase):
    r"""
    No residual connection.
    This class is only used for simpler code compatibility
    """

    def __init__(self, skip_steps: int = 1):
        super().__init__(skip_steps=skip_steps)

    def __repr__(self):
        r"""
        Controls how the class is printed
        """
        return f"{self.__class__.__name__}"

    @classproperty
    def h_dim_increase_type(cls):
        r"""
        Returns:

            None:
                The dimension of the output features do not change at each layer.
        """

        return None

    @classproperty
    def has_weights(cls):
        r"""
        Returns:

            False
                The current class does not use weights

        """
        return False

    def forward(self, h: torch.Tensor, h_prev: torch.Tensor, step_idx: int):
        r"""
        Ignore the skip connection.

        Returns:

            h: torch.Tensor(..., m)
                Return same as input.

            h_prev: torch.Tensor(..., m)
                Return same as input.

        """
        return h, h_prev


class ResidualConnectionSimple(ResidualConnectionBase):
    def __init__(self, skip_steps: int = 1):
        r"""
        Class for the simple residual connections proposed by ResNet,
        where the current layer output is summed to a
        previous layer output.

        Parameters:

            skip_steps: int
                The number of steps to skip between the residual connections.
                If `1`, all the layers are connected. If `2`, half of the
                layers are connected.
        """
        super().__init__(skip_steps=skip_steps)

    @classproperty
    def h_dim_increase_type(cls):
        r"""
        Returns:

            None:
                The dimension of the output features do not change at each layer.
        """

        return None

    @classproperty
    def has_weights(cls):
        r"""
        Returns:

            False
                The current class does not use weights

        """
        return False

    def forward(self, h: torch.Tensor, h_prev: torch.Tensor, step_idx: int):
        r"""
        Add ``h`` with the previous layers with skip connection ``h_prev``,
        similar to ResNet.

        Parameters:

            h: torch.Tensor(..., m)
                The current layer features

            h_prev: torch.Tensor(..., m), None
                The features from the previous layer with a skip connection.
                At ``step_idx==0``, ``h_prev`` can be set to ``None``.

            step_idx: int
                Current layer index or step index in the forward loop of the architecture.

        Returns:

            h: torch.Tensor(..., m)
                Either return ``h`` unchanged, or the sum with
                on ``h_prev``, depending on the ``step_idx`` and ``self.skip_steps``.

            h_prev: torch.Tensor(..., m)
                Either return ``h_prev`` unchanged, or the same value as ``h``,
                depending on the ``step_idx`` and ``self.skip_steps``.

        """
        if self._bool_apply_skip_step(step_idx):
            if step_idx > 0:
                h = h + h_prev
            h_prev = h

        return h, h_prev


class ResidualConnectionWeighted(ResidualConnectionBase):
    def __init__(
        self,
        out_dims,
        skip_steps: int = 1,
        dropout=0.0,
        activation: Union[str, Callable] = "none",
        normalization="none",
        bias=False,
    ):
        r"""
        Class for the simple residual connections proposed by ResNet,
        with an added layer in the residual connection itself.
        The layer output is summed to a a non-linear transformation
        of a previous layer output.

        Parameters:

            skip_steps: int
                The number of steps to skip between the residual connections.
                If `1`, all the layers are connected. If `2`, half of the
                layers are connected.

            out_dims: list(int)
                list of all output dimensions for the network
                that will use this residual connection.
                E.g. ``out_dims = [4, 8, 8, 8, 2]``.

            dropout: float
                value between 0 and 1.0 representing the percentage of dropout
                to use in the weights

            activation: str, Callable
                The activation function to use after the skip weights

            normalization:
                Normalization to use. Choices:

                - "none" or `None`: No normalization
                - "batch_norm": Batch normalization
                - "layer_norm": Layer normalization in the hidden layers.
                - `Callable`: Any callable function

            bias: bool
                Whether to apply add a bias after the weights

        """

        super().__init__(skip_steps=skip_steps)

        self.residual_list = nn.ModuleList()
        self.skip_count = 0
        self.out_dims = out_dims

        for ii in range(0, len(self.out_dims) - 1, self.skip_steps):
            this_dim = self.out_dims[ii]
            self.residual_list.append(
                FCLayer(
                    this_dim,
                    this_dim,
                    activation=activation,
                    dropout=dropout,
                    normalization=normalization,
                    bias=False,
                )
            )

    @classproperty
    def h_dim_increase_type(cls):
        r"""
        Returns:

            None:
                The dimension of the output features do not change at each layer.
        """
        return None

    @classproperty
    def has_weights(cls):
        r"""
        Returns:

            True
                The current class uses weights

        """
        return True

    def forward(self, h: torch.Tensor, h_prev: torch.Tensor, step_idx: int):
        r"""
        Add ``h`` with the previous layers with skip connection ``h_prev``, after
        a feed-forward layer.

        Parameters:

            h: torch.Tensor(..., m)
                The current layer features

            h_prev: torch.Tensor(..., m), None
                The features from the previous layer with a skip connection.
                At ``step_idx==0``, ``h_prev`` can be set to ``None``.

            step_idx: int
                Current layer index or step index in the forward loop of the architecture.

        Returns:

            h: torch.Tensor(..., m)
                Either return ``h`` unchanged, or the sum with the output of a NN layer
                on ``h_prev``, depending on the ``step_idx`` and ``self.skip_steps``.

            h_prev: torch.Tensor(..., m)
                Either return ``h_prev`` unchanged, or the same value as ``h``,
                depending on the ``step_idx`` and ``self.skip_steps``.

        """

        if self._bool_apply_skip_step(step_idx):
            if step_idx > 0:
                h = h + self.residual_list[self.skip_count].forward(h_prev)
                self.skip_count += 1
            h_prev = h

        return h, h_prev

    def _bool_apply_skip_step(self, step_idx: int):
        return super()._bool_apply_skip_step(step_idx) and self.skip_count < len(self.residual_list)


class ResidualConnectionConcat(ResidualConnectionBase):
    def __init__(self, skip_steps: int = 1):
        r"""
        Class for the simple residual connections proposed but where
        the skip connection features are concatenated to the current
        layer features.

        Parameters:

            skip_steps: int
                The number of steps to skip between the residual connections.
                If `1`, all the layers are connected. If `2`, half of the
                layers are connected.
        """

        super().__init__(skip_steps=skip_steps)

    @classproperty
    def h_dim_increase_type(cls):
        r"""
        Returns:

            "previous":
                The dimension of the output layer is the concatenation with the previous layer.
        """

        return "previous"

    @classproperty
    def has_weights(cls):
        r"""
        Returns:

            False
                The current class does not use weights

        """
        return False

    def forward(self, h: torch.Tensor, h_prev: torch.Tensor, step_idx: int):
        r"""
        Concatenate ``h`` with the previous layers with skip connection ``h_prev``.

        Parameters:

            h: torch.Tensor(..., m)
                The current layer features

            h_prev: torch.Tensor(..., n), None
                The features from the previous layer with a skip connection.
                Usually, we have ``n`` equal to ``m``.
                At ``step_idx==0``, ``h_prev`` can be set to ``None``.

            step_idx: int
                Current layer index or step index in the forward loop of the architecture.

        Returns:

            h: torch.Tensor(..., m) or torch.Tensor(..., m + n)
                Either return ``h`` unchanged, or the concatenation
                with ``h_prev``, depending on the ``step_idx`` and ``self.skip_steps``.

            h_prev: torch.Tensor(..., m) or torch.Tensor(..., m + n)
                Either return ``h_prev`` unchanged, or the same value as ``h``,
                depending on the ``step_idx`` and ``self.skip_steps``.

        """

        if self._bool_apply_skip_step(step_idx):
            h_in = h
            if step_idx > 0:
                h = torch.cat([h, h_prev], dim=-1)
            h_prev = h_in

        return h, h_prev


class ResidualConnectionDenseNet(ResidualConnectionBase):
    def __init__(self, skip_steps: int = 1):
        r"""
        Class for the residual connections proposed by DenseNet, where
        all previous skip connection features are concatenated to the current
        layer features.

        Parameters:

            skip_steps: int
                The number of steps to skip between the residual connections.
                If `1`, all the layers are connected. If `2`, half of the
                layers are connected.
        """

        super().__init__(skip_steps=skip_steps)

    @classproperty
    def h_dim_increase_type(cls):
        r"""
        Returns:

            "cumulative":
                The dimension of the output layer is the concatenation of all the previous layer.
        """

        return "cumulative"

    @classproperty
    def has_weights(cls):
        r"""
        Returns:

            False
                The current class does not use weights

        """
        return False

    def forward(self, h: torch.Tensor, h_prev: torch.Tensor, step_idx: int):
        r"""
        Concatenate ``h`` with all the previous layers with skip connection ``h_prev``.

        Parameters:

            h: torch.Tensor(..., m)
                The current layer features

            h_prev: torch.Tensor(..., n), None
                The features from the previous layers.
                n = ((step_idx // self.skip_steps) + 1) * m

                At ``step_idx==0``, ``h_prev`` can be set to ``None``.

            step_idx: int
                Current layer index or step index in the forward loop of the architecture.

        Returns:

            h: torch.Tensor(..., m) or torch.Tensor(..., m + n)
                Either return ``h`` unchanged, or the concatenation
                with ``h_prev``, depending on the ``step_idx`` and ``self.skip_steps``.

            h_prev: torch.Tensor(..., m) or torch.Tensor(..., m + n)
                Either return ``h_prev`` unchanged, or the same value as ``h``,
                depending on the ``step_idx`` and ``self.skip_steps``.

        """

        if self._bool_apply_skip_step(step_idx):
            if step_idx > 0:
                h = torch.cat([h, h_prev], dim=-1)
            h_prev = h

        return h, h_prev


class ResidualConnectionRandom(ResidualConnectionBase):
    def __init__(self, skip_steps=1, out_dims: List[int] = None, num_layers: int = None):
        r"""
        Class for the random residual connection, where each layer is connected
        to each following layer with a random weight between 0 and 1.
        Parameters:
            skip_steps:
                Parameter only there for compatibility with other classes of the same parent.
            out_dims:
                The list of output dimensions. Only required to get the number
                of layers. Must be provided if `num_layers` is None.
            num_layers:
                The number of layers. Must be provided if `out_dims` is None.
        """
        if skip_steps != 1:
            raise ValueError("Only `skip_step=1` is implemented")
        super().__init__(skip_steps=skip_steps)

        if out_dims is not None:
            if num_layers is not None:
                assert num_layers == len(out_dims)
            num_layers = len(out_dims)
        if num_layers is None:
            raise ValueError("Either `out_dims` or `num_layers` must be provided")
        self.num_layers = num_layers

        self.random_dict_weights = {}
        for ii in range(1, self.num_layers):
            random_weights = torch.rand(ii)
            self.random_dict_weights[ii] = random_weights

    @classproperty
    def h_dim_increase_type(cls):
        r"""
        Returns:
            None:
                The dimension of the output features do not change at each layer.
        """

        return None

    @classproperty
    def has_weights(cls):
        r"""
        Returns:
            False
                The current class does not use weights
        """
        return False

    def forward(self, h: torch.Tensor, h_prev: torch.Tensor, step_idx: int):
        r"""
        Add ``h`` with the previous layers with skip connection ``h_prev``,
        similar to ResNet.
        Parameters:
            h: torch.Tensor(..., m)
                The current layer features
            h_prev: torch.Tensor(..., m), None
                The features from the previous layer with a skip connection.
                At ``step_idx==0``, ``h_prev`` can be set to ``None``.
            step_idx: int
                Current layer index or step index in the forward loop of the architecture.
        Returns:
            h: torch.Tensor(..., m)
                Either return ``h`` unchanged, or the sum with
                on ``h_prev``, depending on the ``step_idx`` and ``self.skip_steps``.
            h_prev: torch.Tensor(..., m)
                Either return ``h_prev`` unchanged, or the same value as ``h``,
                depending on the ``step_idx`` and ``self.skip_steps``.
        """

        if self._bool_apply_skip_step(step_idx):
            for i in range(0, step_idx):
                h += (
                    self.random_dict_weights[step_idx][i].to(dtype=h_prev[i].dtype, device=h_prev[i].device)
                    * h_prev[i]
                )
            if h_prev is None:
                h_prev = [h]
            else:
                h_prev.append(h)

        return h, h_prev


# graphium/nn/utils.py

import abc
import inspect
from numbers import Real
from typing import Optional


class MupMixin(abc.ABC):
    @abc.abstractmethod
    def make_mup_base_kwargs(self, divide_factor: float = 2.0, factor_in_dim: Optional[bool] = None):
        """
        Create a 'base' model to be used by the `mup` or `muTransfer` scaling of the model.
        The base model is usually identical to the regular model, but with the
        layers width divided by a given factor (2 by default)

        This is done using the `scale_kwargs()` method with `scale_factor = 1 / divide_factor`.

        Parameter:
            divide_factor: Factor by which to divide the width.
            factor_in_dim: Whether to factor the input dimension for the nodes. If None, the default for scale_kwargs is used
        Returns:
            kwargs: Dictionary of parameters to be used to instanciate the base model divided by the factor
        """
        ...

    def scale_kwargs(self, scale_factor: Real, scale_in_dim: bool = False):
        """
        Create a "scaled" version of the module where the hidden dims are scaled as in muTransfer.

        This can be used with `scale_factor` < 1 to create a "base" model to extract shape
        information as in the `mup` package or with `scale_factor` > 1 to create a scaled model
        to which optimal hyperparameters can be "muTransferred" from the original model

        Parameters:
            scale_factor: Factor by which to scale the width.
            scale_in_dim: Whether to factor the input dimension for the nodes

        Returns:
            kwargs: Dictionary of parameters to be used to instantiate the scaled model
        """

        divide_factor = 1 / scale_factor

        if scale_in_dim is None:
            return self.make_mup_base_kwargs(divide_factor=divide_factor)

        # If scale_in_dim passed, need to check it can be forwarded
        try:
            return self.make_mup_base_kwargs(divide_factor=divide_factor, factor_in_dim=scale_in_dim)
        except TypeError as e:
            raise "This error may have been caused by passing scale_in_dim to scale_kwargs for a class that does not support passing factor_in_dim to make_mup_base_kwargs, which cannot be done" from e


# graphium/trainer/__init__.py

from . import predictor
from . import metrics

from .predictor import PredictorModule


# graphium/trainer/losses.py

from typing import Optional

import torch
from torch import Tensor
from torch.nn import functional as F
from torch.nn.modules.loss import _WeightedLoss


class HybridCELoss(_WeightedLoss):
    def __init__(
        self,
        n_brackets,
        regression_loss: str = "mse",
        alpha: float = 0.5,
        weight: Optional[Tensor] = None,
        reduction: str = "mean",
    ) -> None:
        """
        A hybrid between the regression loss (either MAE or MSE) and the cross entropy loss. Intended
        to be used with noisy regression datasets, for which the targets are assigned to binary brackets,
        and the task is transformed into a multi-class classification.

        Note that it assumes that the brackets are consecutive integers starting at 0 up to n_brackets,
        which has an impact on the scale of the regression loss component.

        Parameters:
            n_brackets: the number of brackets that will be used to group the regression targets.
                Expected to have the same size as the number of classes in the transformed regression task.
            regression_loss: type of regression loss, either 'mse' or 'mae'.
            alpha: weight assigned to the CE loss component. Must be a value in [0, 1] range.
            weight: a manual rescaling weight given to each class in the CE loss component.
                If given, has to be a Tensor of the same size as the number of classes.
            reduction: specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'.
                'none': no reduction will be applied, 'mean': the sum of the output will be divided
                by the number of elements in the output, 'sum': the output will be summed.
        """
        super().__init__(weight=weight, reduction=reduction)

        if regression_loss == "mae":
            self.regression_loss = F.l1_loss
        elif regression_loss == "mse":
            self.regression_loss = F.mse_loss
        else:
            raise ValueError(
                f"Expected regression_loss to be in {{'mae', 'mse'}}, received {regression_loss}."
            )

        if alpha < 0 or alpha > 1:
            raise ValueError(f"Expected alpha to be in the [0, 1] range, received {alpha}.")

        self.brackets = Tensor(range(n_brackets))
        self.alpha = alpha
        self.softmax = torch.nn.Softmax(dim=1)

    def forward(self, input: Tensor, target: Tensor, nan_targets: Tensor = None) -> Tensor:
        """
        Parameters:
            input: (batch_size x n_classes) tensor of logits predicted for each bracket.
            target: (batch_size) or (batch_size, 1) tensor of target brackets in {0, 1, ..., self.n_brackets}.
        """

        target = target.flatten()

        # set input and target with nans to 0s for regression loss
        if nan_targets is not None:
            input = torch.masked_fill(input, nan_targets.unsqueeze(1), 0)
            target = torch.masked_fill(target, nan_targets, 0)
        # regression loss needs normalized logits to probability as input to do inner product with self.brackets
        # we apply softmax on the raw logits first
        softmax_input = self.softmax(input)
        # the softmax of a tensor of 0s would not be 0s anymore, so need to apply nan_targets here again to filter out
        if nan_targets is not None:
            softmax_input = torch.masked_fill(softmax_input, nan_targets.unsqueeze(1), 0.0)
        # [batch_size, n_classes] * [n_classes] ([0, 1, 2...n_brakets-1]) -> [batch_size]
        regression_input = torch.inner(softmax_input.to(self.brackets.dtype), self.brackets.to(input.device))
        regression_loss = self.regression_loss(regression_input, target.float(), reduction=self.reduction)
        # regression_loss needs some scaling by total_targets/num_real_targets
        if nan_targets is not None:
            num_real_targets = (~nan_targets).sum()
            factor1 = torch.where(num_real_targets > 0, 1, 0)
            factor2 = torch.where(num_real_targets > 0, 0, 1)
            regression_loss = factor1 * regression_loss * nan_targets.numel() / (num_real_targets + factor2)

            # set input and target with nans to -1000s for ce loss
            input = torch.masked_fill(input, nan_targets.unsqueeze(1), -1000)
            target = torch.masked_fill(target, nan_targets, -1000)
        # cross_entropy loss needs raw logits as input
        # ce_loss does not need scaling as it already ignores -1000 masked nan values
        ce_loss = F.cross_entropy(
            input, target.long(), weight=self.weight, ignore_index=-1000, reduction=self.reduction
        )

        return self.alpha * ce_loss + (1 - self.alpha) * regression_loss


# graphium/trainer/metrics.py

from typing import Union, Callable, Optional, Dict, Any

import sys

import torch
from torch import Tensor
import operator as op

from torchmetrics.utilities.distributed import reduce
import torchmetrics.functional.regression.mae

from graphium.utils.tensor import nan_mean

# NOTE(hadim): the below is a fix to be able to import previously saved Graphium model that are incompatible
# with the current version of torchmetrics.
# In the future, we should NOT save any torchmetrics objects during serialization.
# See https://github.com/valence-discovery/graphium/issues/106
sys.modules["torchmetrics.functional.regression.mean_absolute_error"] = torchmetrics.functional.regression.mae

EPS = 1e-5


class Thresholder:
    def __init__(
        self,
        threshold: float,
        operator: Union[str, Callable] = "greater",
        th_on_preds: bool = True,
        th_on_target: bool = False,
    ):
        # Basic params
        self.threshold = threshold
        self.th_on_target = th_on_target
        self.th_on_preds = th_on_preds
        self.operator, self.op_str = self._get_operator(operator)

    def compute(self, preds: Tensor, target: Tensor):
        # Apply the threshold on the predictions
        if self.th_on_preds:
            preds = self.operator(preds, self.threshold)

        # Apply the threshold on the targets
        if self.th_on_target:
            target = self.operator(target, self.threshold)

        return preds, target

    def __call__(self, preds: Tensor, target: Tensor):
        return self.compute(preds, target)

    def __repr__(self):
        r"""
        Control how the class is printed
        """

        return f"{self.op_str}{self.threshold}"

    @staticmethod
    def _get_operator(operator):
        """Operator can either be a string, or a callable"""
        if isinstance(operator, str):
            op_name = operator.lower()
            if op_name in ["greater", "gt"]:
                op_str = ">"
                operator = op.gt
            elif op_name in ["lower", "lt"]:
                op_str = "<"
                operator = op.lt
            else:
                raise ValueError(f"operator `{op_name}` not supported")
        elif callable(operator):
            op_str = operator.__name__
            if op_str == "lt":
                op_str = "<"
            elif op_str == "gt":
                op_str = ">"
        elif operator is None:
            pass
        else:
            raise TypeError(f"operator must be either `str` or `callable`, provided: `{type(operator)}`")

        return operator, op_str

    def __getstate__(self):
        """Serialize the class for pickling."""
        state = {}
        state["threshold"] = self.threshold
        state["th_on_target"] = self.th_on_target
        state["th_on_preds"] = self.th_on_preds

        # Set the operator state.
        # If it's a callable, it's up to the user to ensure it unserializes well
        operator = self.operator
        if operator == op.lt:
            operator = "lower"
        elif operator == op.gt:
            operator = "greater"
        state["operator"] = operator

        return state

    def __setstate__(self, state: dict):
        state["operator"], state["op_str"] = self._get_operator(state["operator"])
        self.__dict__.update(state)
        self.__init__(**state)

    def __eq__(self, obj) -> bool:
        is_eq = [
            self.threshold == obj.threshold,
            self.th_on_target == obj.th_on_target,
            self.th_on_preds == obj.th_on_preds,
            self.operator == obj.operator,
            self.op_str == obj.op_str,
        ]
        return all(is_eq)


class MetricWrapper:
    r"""
    Allows to initialize a metric from a name or Callable, and initialize the
    `Thresholder` in case the metric requires a threshold.
    """

    def __init__(
        self,
        metric: Union[str, Callable],
        threshold_kwargs: Optional[Dict[str, Any]] = None,
        target_nan_mask: Optional[Union[str, int]] = None,
        multitask_handling: Optional[str] = None,
        squeeze_targets: bool = False,
        target_to_int: bool = False,
        **kwargs,
    ):
        r"""
        Parameters
            metric:
                The metric to use. See `METRICS_DICT`

            threshold_kwargs:
                If `None`, no threshold is applied.
                Otherwise, we use the class `Thresholder` is initialized with the
                provided argument, and called before the `compute`

            target_nan_mask:

                - None: Do not change behaviour if there are NaNs

                - int, float: Value used to replace NaNs. For example, if `target_nan_mask==0`, then
                  all NaNs will be replaced by zeros

                - 'ignore': The NaN values will be removed from the tensor before computing the metrics.
                  Must be coupled with the `multitask_handling='flatten'` or `multitask_handling='mean-per-label'`.

            multitask_handling:
                - None: Do not process the tensor before passing it to the metric.
                  Cannot use the option `multitask_handling=None` when `target_nan_mask=ignore`.
                  Use either 'flatten' or 'mean-per-label'.

                - 'flatten': Flatten the tensor to produce the equivalent of a single task

                - 'mean-per-label': Loop all the labels columns, process them as a single task,
                    and average the results over each task
                  *This option might slow down the computation if there are too many labels*

            squeeze_targets:
                If true, targets will be squeezed prior to computing the metric.
                Required in classifigression task.

            target_to_int:
                If true, targets will be converted to integers prior to computing the metric.

            kwargs:
                Other arguments to call with the metric
        """

        self.metric, self.metric_name = self._get_metric(metric)
        self.thresholder = None
        if threshold_kwargs is not None:
            self.thresholder = Thresholder(**threshold_kwargs)

        self.target_nan_mask = self._parse_target_nan_mask(target_nan_mask)
        self.multitask_handling = self._parse_multitask_handling(multitask_handling, self.target_nan_mask)
        self.squeeze_targets = squeeze_targets
        self.target_to_int = target_to_int
        self.kwargs = kwargs

    @staticmethod
    def _parse_target_nan_mask(target_nan_mask):
        """
        Parse the `target_nan_mask` parameter
        """

        if (target_nan_mask is None) or isinstance(target_nan_mask, (int, float)):
            # None, int, float, are accepted
            pass
        elif isinstance(target_nan_mask, Tensor) and (target_nan_mask.numel() == 1):
            # Only single element tensors are accepted
            target_nan_mask = target_nan_mask.flatten()[0]
        elif isinstance(target_nan_mask, str):
            # Only a few str options are accepted
            target_nan_mask = target_nan_mask.lower()
            accepted_str = ["ignore", "none"]
            assert (
                target_nan_mask in accepted_str
            ), f"Provided {target_nan_mask} not in accepted_str={accepted_str}"

            if target_nan_mask == "none":
                target_nan_mask = None
        else:
            raise ValueError(f"Unrecognized option `target_nan_mask={target_nan_mask}`")
        return target_nan_mask

    @staticmethod
    def _parse_multitask_handling(multitask_handling, target_nan_mask):
        """
        Parse the `multitask_handling` parameter
        """

        if multitask_handling is None:
            # None is accepted
            pass
        elif isinstance(multitask_handling, str):
            # Only a few str options are accepted
            multitask_handling = multitask_handling.lower()
            accepted_str = ["flatten", "mean-per-label", "none"]
            assert (
                multitask_handling in accepted_str
            ), f"Provided {multitask_handling} not in accepted_str={accepted_str}"

            if multitask_handling == "none":
                multitask_handling = None
        else:
            raise ValueError(f"Unrecognized option `multitask_handling={multitask_handling}`")

        if (target_nan_mask == "ignore") and (multitask_handling is None):
            raise ValueError(
                "Cannot use the option `multitask_handling=None` when `target_nan_mask=ignore`. Use either 'flatten' or 'mean-per-label'"
            )

        return multitask_handling

    @staticmethod
    def _get_metric(metric):
        from graphium.utils.spaces import METRICS_DICT

        if isinstance(metric, str):
            metric_name = metric
            metric = METRICS_DICT[metric]
        else:
            metric_name = None
            metric = metric
        return metric, metric_name

    def compute(self, preds: Tensor, target: Tensor) -> Tensor:
        r"""
        Compute the metric, apply the thresholder if provided, and manage the NaNs
        """
        if preds.ndim == 1:
            preds = preds.unsqueeze(-1)

        if target.ndim == 1:
            target = target.unsqueeze(-1)

        # Threshold the prediction
        if self.thresholder is not None:
            preds, target = self.thresholder(preds, target)

        target_nans = torch.isnan(target)

        # for the classifigression task, cast predictions from
        # (batch_size, n_targets * n_brackets) to (batch_size, n_targets, n_brackets)
        # TODO: make this more flexible to the target shape in the future
        if preds.shape[1] != target.shape[1]:
            preds = preds.view(target.shape[0], target.shape[1], -1)
            classifigression = True
        else:
            classifigression = False

        if self.multitask_handling is None:
            # In case of no multi-task handling, apply the nan filtering, then compute the metrics
            assert (
                self.target_nan_mask != "ignore"
            ), f"Cannot use the option `multitask_handling=None` when `target_nan_mask=ignore`. Use either 'flatten' or 'mean-per-label'"
            preds, target = self._filter_nans(preds, target)
            if self.squeeze_targets:
                target = target.squeeze()
            if self.target_to_int:
                target = target.to(int)
            metric_val = self.metric(preds, target, **self.kwargs)
        elif self.multitask_handling == "flatten":
            # Flatten the tensors, apply the nan filtering, then compute the metrics
            if classifigression:
                preds = preds.view(-1, preds.shape[-1])
                target = target.flatten()
            else:
                preds, target = preds.flatten(), target.flatten()
            preds, target = self._filter_nans(preds, target)
            if self.squeeze_targets:
                target = target.squeeze()
            if self.target_to_int:
                target = target.to(int)
            metric_val = self.metric(preds, target, **self.kwargs)
        elif self.multitask_handling == "mean-per-label":
            # Loop the columns (last dim) of the tensors, apply the nan filtering, compute the metrics per column, then average the metrics
            target_list = [target[..., ii][~target_nans[..., ii]] for ii in range(target.shape[-1])]
            # TODO: make this more flexible to the target shape in the future
            if classifigression:
                preds_list = [preds[..., i, :][~target_nans[..., i]] for i in range(preds.shape[1])]
            else:
                preds_list = [preds[..., ii][~target_nans[..., ii]] for ii in range(preds.shape[-1])]
            metric_val = []
            for ii in range(len(target_list)):
                try:
                    this_preds, this_target = self._filter_nans(preds_list[ii], target_list[ii])
                    if self.squeeze_targets:
                        this_target = this_target.squeeze()
                    if self.target_to_int:
                        this_target = this_target.to(int)
                    metric_val.append(self.metric(this_preds, this_target, **self.kwargs))
                except:
                    pass
            # Average the metric
            metric_val = nan_mean(torch.stack(metric_val))
        else:
            # Wrong option
            raise ValueError(f"Invalid option `self.multitask_handling={self.multitask_handling}`")

        return metric_val

    def _filter_nans(self, preds: Tensor, target: Tensor):
        """Handle the NaNs according to the chosen options"""
        target_nans = torch.isnan(target)

        if self.target_nan_mask is None:
            pass
        elif isinstance(self.target_nan_mask, (int, float)):
            target = target.clone()
            target[torch.isnan(target)] = self.target_nan_mask
        elif self.target_nan_mask == "ignore":
            target = target[~target_nans]
            preds = preds[~target_nans]
        else:
            raise ValueError(f"Invalid option `{self.target_nan_mask}`")
        return preds, target

    def __call__(self, preds: Tensor, target: Tensor) -> Tensor:
        r"""
        Compute the metric with the method `self.compute`
        """
        return self.compute(preds, target)

    def __repr__(self):
        r"""
        Control how the class is printed
        """
        full_str = f"{self.metric.__name__}"
        if self.thresholder is not None:
            full_str += f"({self.thresholder})"

        return full_str

    def __eq__(self, obj) -> bool:
        is_eq = [
            self.metric == obj.metric,
            self.metric_name == obj.metric_name,
            self.thresholder == obj.thresholder,
            self.target_nan_mask == obj.target_nan_mask,
            self.multitask_handling == obj.multitask_handling,
            self.squeeze_targets == obj.squeeze_targets,
            self.target_to_int == obj.target_to_int,
            self.kwargs == obj.kwargs,
        ]
        return all(is_eq)

    def __getstate__(self):
        """Serialize the class for pickling."""
        state = {}
        if self.metric_name is None:
            state["metric"] = self.metric
        else:
            state["metric"] = self.metric_name
        state["target_nan_mask"] = self.target_nan_mask
        state["multitask_handling"] = self.multitask_handling
        state["squeeze_targets"] = self.squeeze_targets
        state["target_to_int"] = self.target_to_int
        state["kwargs"] = self.kwargs
        state["threshold_kwargs"] = None
        if self.thresholder is not None:
            state["threshold_kwargs"] = self.thresholder.__getstate__()
        return state

    def __setstate__(self, state: dict):
        """Reload the class from pickling."""
        state["metric"], state["metric_name"] = self._get_metric(state["metric"])
        thresholder = state.pop("threshold_kwargs", None)
        if thresholder is not None:
            thresholder = Thresholder(**thresholder)
        state["thresholder"] = thresholder

        self.__dict__.update(state)


# graphium/trainer/predictor.py

from graphium.trainer.metrics import MetricWrapper
from typing import Dict, List, Any, Union, Any, Callable, Tuple, Type, Optional
import numpy as np
from copy import deepcopy
import time
from loguru import logger

import torch
from torch import nn, Tensor
import lightning
from torch_geometric.data import Data, Batch
from mup.optim import MuAdam

from graphium.config.config_convert import recursive_config_reformating
from graphium.trainer.predictor_options import EvalOptions, FlagOptions, ModelOptions, OptimOptions
from graphium.trainer.predictor_summaries import TaskSummaries
from graphium.data.datamodule import BaseDataModule
from graphium.utils.moving_average_tracker import MovingAverageTracker

GRAPHIUM_PRETRAINED_MODELS = {
    "graphium-zinc-micro-dummy-test": "gcs://graphium-public/pretrained-models/graphium-zinc-micro-dummy-test/model.ckpt",
    "dummy-pretrained-model": "saved_models/pretrained_models/dummy-pretrained-model.ckpt"  # dummy model (to be deleted later)
}


class PredictorModule(lightning.LightningModule):
    def __init__(
        self,
        model_class: Type[nn.Module],
        model_kwargs: Dict[str, Any],
        loss_fun: Dict[str, Union[str, Callable]],
        random_seed: int = 42,
        optim_kwargs: Optional[Dict[str, Any]] = None,
        torch_scheduler_kwargs: Optional[Dict[str, Any]] = None,
        scheduler_kwargs: Optional[Dict[str, Any]] = None,
        target_nan_mask: Optional[Union[str, int]] = None,
        multitask_handling: Optional[str] = None,
        metrics: Dict[str, Callable] = None,
        metrics_on_progress_bar: Dict[str, List[str]] = [],
        metrics_on_training_set: Optional[Dict[str, List[str]]] = None,
        flag_kwargs: Dict[str, Any] = None,
        task_norms: Optional[Dict[Callable, Any]] = None,
        metrics_every_n_train_steps: Optional[int] = None,
    ):
        """
        The Lightning module responsible for handling the predictions, losses, metrics, optimization, etc.
        It works in a multi-task setting, with different losses and metrics per class

        Parameters:
            model_class: The torch Module containing the main forward function
            model_kwargs: The arguments to initialize the model from `model_class`
            loss_fun: A `dict[str, fun]`, where `str` is the task name and `fun` the loss function
            random_seed: The seed for random initialization
            optim_kwargs: The optimization arguments. See class `OptimOptions`
            torch_scheduler_kwargs: The torch scheduler arguments. See class `OptimOptions`
            scheduler_kwargs: The lightning scheduler arguments. See class `OptimOptions`
            target_nan_mask: How to handle the NaNs. See `MetricsWrapper` for options
            metrics: A `dict[str, fun]`, where `str` is the task name and `fun` the metric function
            metrics_on_progress_bar: A `dict[str, list[str2]`, where `str` is the task name and `str2` the metrics to include on the progress bar
            metrics_on_training_set: A `dict[str, list[str2]`, where `str` is the task name and `str2` the metrics to include on the training set
            flag_kwargs: Arguments related to using the FLAG adversarial augmentation
            task_norms: the normalization for each task
            metrics_every_n_train_steps: Compute and log metrics every n training steps.
                Set to `None` to never log the training metrics and statistics (the default). Set to `1` to log at every step.
        """
        self.save_hyperparameters()

        self.random_seed = random_seed
        torch.random.manual_seed(self.random_seed)
        np.random.seed(self.random_seed)

        self.target_nan_mask = target_nan_mask
        self.multitask_handling = multitask_handling
        self.task_norms = task_norms

        super().__init__()

        # Setting the model options
        self.model_kwargs = model_kwargs
        self._model_options = ModelOptions(model_class=model_class, model_kwargs=model_kwargs)
        # Setting the optimizer options
        self.optim_options = OptimOptions(
            optim_kwargs=optim_kwargs,
            torch_scheduler_kwargs=torch_scheduler_kwargs,
            scheduler_kwargs=scheduler_kwargs,
        )
        # Setting the evaluation options
        eval_options = {}
        for task in loss_fun:
            eval_options[task] = EvalOptions(
                loss_fun=loss_fun[task],
                metrics=metrics[task],
                metrics_on_progress_bar=metrics_on_progress_bar[task],
                metrics_on_training_set=metrics_on_training_set[task]
                if metrics_on_training_set is not None
                else None,
            )
            eval_options[task].check_metrics_validity()

        self._eval_options_dict: Dict[str, EvalOptions] = eval_options
        self._eval_options_dict = {
            self._get_task_key(
                task_level=model_kwargs["task_heads_kwargs"][key]["task_level"], task=key
            ): value
            for key, value in self._eval_options_dict.items()
        }
        # Setting the flag options
        self._flag_options = FlagOptions(flag_kwargs=flag_kwargs)

        self.model = self._model_options.model_class(**self._model_options.model_kwargs)
        loss_fun = {
            self._get_task_key(
                task_level=model_kwargs["task_heads_kwargs"][key]["task_level"], task=key
            ): value
            for key, value in loss_fun.items()
        }
        self.tasks = list(loss_fun.keys())

        # Task-specific evalutation attributes
        self.loss_fun = {}
        self.metrics = {}
        self.metrics_on_progress_bar = {}
        self.metrics_on_training_set = {}
        for task in self.tasks:
            self.loss_fun[task] = EvalOptions.parse_loss_fun(loss_fun[task])
            self.metrics[task] = (
                self._eval_options_dict[task].metrics
                if self._eval_options_dict[task].metrics is not None
                else {}
            )
            self.metrics_on_progress_bar[task] = self._eval_options_dict[task].metrics_on_progress_bar
            self.metrics_on_training_set[task] = (
                list(self.metrics[task].keys())
                if self._eval_options_dict[task].metrics_on_training_set is None
                else self._eval_options_dict[task].metrics_on_training_set
            )
        self.n_params = sum(p.numel() for p in self.parameters() if p.requires_grad)

        # Set the parameters and default values for the FLAG adversarial augmentation, and check values
        self._flag_options.set_kwargs()
        self.flag_kwargs = self._flag_options.flag_kwargs

        # Set the parameters for optimizer options
        self.optim_options.set_kwargs()

        # Initialize the epoch summary
        monitor = self.optim_options.scheduler_kwargs["monitor"].split("/")[0]
        mode = self.optim_options.scheduler_kwargs["mode"]

        self.task_epoch_summary = TaskSummaries(
            task_loss_fun=self.loss_fun,
            task_metrics=self.metrics,
            task_metrics_on_training_set=self.metrics_on_training_set,
            task_metrics_on_progress_bar=self.metrics_on_progress_bar,
            monitor=monitor,
            mode=mode,
        )

        # This helps avoid a bug when saving hparams to yaml with different dict or str formats
        self._set_hparams(recursive_config_reformating(self.hparams))

        # throughput estimation
        self.mean_val_time_tracker = MovingAverageTracker()
        self.mean_val_tput_tracker = MovingAverageTracker()
        self.validation_step_outputs = []
        self.test_step_outputs = []
        self.epoch_start_time = None

        # Decide whether to log every step or once at the end
        # of the epoch.
        self.metrics_every_n_train_steps = metrics_every_n_train_steps
        # Wether save preds and targets for each training step.

    def forward(
        self, inputs: Dict
    ) -> Dict[str, Union[Tensor, Dict[str, Tensor], Dict[str, Dict[str, Tensor]]]]:
        r"""
        Returns the result of `self.model.forward(*inputs)` on the inputs.
        If the output of `out = self.model.forward` is a dictionary with a `"preds"` key,
        it is returned directly. Otherwise, a new dictionary is created and
        returns `{"preds": out}`.

        Returns:
            A dict with a key `"preds"` representing the prediction of the network.

        """
        # Convert to the right dtype and run the model
        feats = self._convert_features_dtype(inputs["features"])
        # *check for nan in model output
        out = self.model.forward(feats)
        if isinstance(out, dict) and ("preds" in out.keys()):
            out_dict = out
        else:
            out_dict = {"preds": out}

        return out_dict

    def _convert_features_dtype(self, feats):
        # Convert features to dtype
        if isinstance(feats, torch.Tensor):
            feats = feats.to(self.dtype)
        elif isinstance(feats, (Data, Batch, dict)):
            for key, val in feats.items():
                if isinstance(val, torch.Tensor) and (val.is_floating_point()):
                    feats[key] = val.to(dtype=self.dtype)
        return feats

    def _get_task_key(self, task_level: str, task: str):
        task_prefix = f"{task_level}_"
        if not task.startswith(task_prefix):
            task = task_prefix + task
        return task

    def configure_optimizers(self, impl=None):
        if impl is None:
            impl = torch.optim.Adam

        # Define the optimizer and schedulers
        optimiser = MuAdam(self.parameters(), **self.optim_options.optim_kwargs, impl=impl)
        torch_scheduler = self.optim_options.scheduler_class(
            optimizer=optimiser, **self.optim_options.torch_scheduler_kwargs
        )
        scheduler = {
            "scheduler": torch_scheduler,
            **self.optim_options.scheduler_kwargs,
        }
        return [optimiser], [scheduler]

    @staticmethod
    def compute_loss(
        preds: Dict[str, Tensor],
        targets: Dict[str, Tensor],
        weights: Optional[Tensor],
        loss_fun: Dict[str, Callable],
        target_nan_mask: Optional[Union[str, int]] = None,
        multitask_handling: Optional[str] = None,
    ) -> Tuple[Tensor, Dict[str, Tensor]]:
        r"""
        Compute the loss using the specified loss function, and dealing with
        the nans in the `targets`.

        Parameters:
            preds:
                Predicted values

            targets:
                Target values

            weights:
                No longer supported, will raise an error.

            target_nan_mask:

                - None: Do not change behaviour if there are NaNs

                - int, float: Value used to replace NaNs. For example, if `target_nan_mask==0`, then
                  all NaNs will be replaced by zeros

                - 'ignore': The NaN values will be removed from the tensor before computing the metrics.
                  Must be coupled with the `multitask_handling='flatten'` or `multitask_handling='mean-per-label'`.

            multitask_handling:
                - None: Do not process the tensor before passing it to the metric.
                  Cannot use the option `multitask_handling=None` when `target_nan_mask=ignore`.
                  Use either 'flatten' or 'mean-per-label'.

                - 'flatten': Flatten the tensor to produce the equivalent of a single task

                - 'mean-per-label': Loop all the labels columns, process them as a single task,
                    and average the results over each task
                  *This option might slowdown the computation if there are too many labels*

            loss_fun:
                Loss function to use

        Returns:
            Tensor:
                weighted_loss: Resulting weighted loss
                all_task_losses: Loss per task
        """

        wrapped_loss_fun_dict = {
            task: MetricWrapper(
                metric=loss,
                threshold_kwargs=None,
                target_nan_mask=target_nan_mask,
                multitask_handling=multitask_handling,
            )
            for task, loss in loss_fun.items()
        }

        if weights is not None:
            raise NotImplementedError("Weights are no longer supported in the loss")
        all_task_losses = {
            task: wrapped(preds=preds[task], target=targets[task])
            for task, wrapped in wrapped_loss_fun_dict.items()
        }
        total_loss = torch.sum(torch.stack(list(all_task_losses.values())), dim=0)
        num_tasks = len(all_task_losses.keys())
        weighted_loss = total_loss / num_tasks
        return weighted_loss, all_task_losses

    def _general_step(self, batch: Dict[str, Tensor], step_name: str, to_cpu: bool) -> Dict[str, Any]:
        r"""Common code for training_step, validation_step and testing_step"""
        preds = self.forward(batch)  # The dictionary of predictions

        # * check for nan in model output
        targets_dict = batch.get("labels")

        # Different type of preds can be return by the forward
        if isinstance(preds, dict) and ("preds" in preds.keys()):
            preds = preds["preds"]
        elif isinstance(preds, Tensor):
            preds = {k: preds[ii] for ii, k in enumerate(targets_dict.keys())}

        preds = {
            self._get_task_key(
                task_level=self.model_kwargs["task_heads_kwargs"][key]["task_level"], task=key
            ): value
            for key, value in preds.items()
        }
        # preds = {k: preds[ii] for ii, k in enumerate(targets_dict.keys())}
        for task, pred in preds.items():
            task_specific_norm = self.task_norms[task] if self.task_norms is not None else None
            if hasattr(task_specific_norm, "normalize_val_test"):
                normalize_val_test = task_specific_norm.normalize_val_test
            else:
                normalize_val_test = False
            if step_name != "train" and not normalize_val_test:
                # apply denormalization for val and test predictions for correct loss and metrics evaluation
                # if normalize_val_test is not true, only train loss will stay as the normalized version
                # if normalize_val_test is true, no denormalization is applied, all losses and metrics are normalized version
                preds[task] = task_specific_norm.denormalize(pred)
            targets_dict[task] = targets_dict[task].to(dtype=pred.dtype)
        weights = batch.get("weights", None)

        loss, task_losses = self.compute_loss(
            preds=preds,
            targets=targets_dict,
            weights=weights,
            loss_fun=self.loss_fun,
            target_nan_mask=self.target_nan_mask,
            multitask_handling=self.multitask_handling,
        )

        device = "cpu" if to_cpu else None
        for task in preds:
            task_specific_norm = self.task_norms[task] if self.task_norms is not None else None
            if hasattr(task_specific_norm, "normalize_val_test"):
                normalize_val_test = task_specific_norm.normalize_val_test
            else:
                normalize_val_test = False
            if step_name == "train" and not normalize_val_test:
                # apply denormalization for targets and predictions for the evaluation of training metrics (excluding loss)
                # if normalize_val_test is not true, train loss will stay as the normalized version
                # if normalize_val_test is true, no denormalization is applied, all losses and metrics are normalized version
                preds[task] = task_specific_norm.denormalize(preds[task])
                targets_dict[task] = task_specific_norm.denormalize(targets_dict[task])
            preds[task] = preds[task].detach().to(device=device)
            targets_dict[task] = targets_dict[task].detach().to(device=device)
        if weights is not None:
            weights = weights.detach().to(device=device)

        step_dict = {"preds": preds, "targets": targets_dict, "weights": weights}
        # step_dict[f"{self.loss_fun._get_name()}/{step_name}"] = loss.detach().cpu()            original

        # step_dict[f"weighted_loss/{step_name}"] = loss.detach().cpu()
        # step_dict[f"loss/{step_name}"] = loss.detach().cpu()
        for task in self.tasks:
            step_dict[
                self.task_epoch_summary.metric_log_name(task, self.loss_fun[task]._get_name(), step_name)
            ] = loss.detach()

        step_dict["loss"] = loss
        # print("loss ", self.global_step, self.current_epoch, loss)
        step_dict["task_losses"] = task_losses
        step_dict["gradient_norm"] = self.get_gradient_norm()
        return step_dict

    def flag_step(self, batch: Dict[str, Tensor], step_name: str, to_cpu: bool) -> Dict[str, Any]:
        r"""
        Perform adversarial data agumentation during one training step using FLAG.
        Paper: https://arxiv.org/abs/2010.09891
        Github: https://github.com/devnkong/FLAG
        """

        alpha, n_steps = self.flag_kwargs["alpha"], self.flag_kwargs["n_steps"]

        X = self._convert_features_dtype(batch["features"])
        X_shape = X["feat"].shape

        pert = torch.FloatTensor(X_shape).uniform_(-alpha, alpha).to(device=X["feat"].device)
        pert.requires_grad = True

        # Perturb the features
        pert_batch = deepcopy(batch)
        features = pert_batch["features"]
        features["feat"] = features["feat"] + pert

        preds = self.forward(pert_batch)["preds"]
        targets = batch.pop("labels")
        for key in targets.keys():
            targets[key] = targets[key].to(dtype=preds[key].dtype)
        weights = batch.pop("weights", None)
        loss, task_losses = self.compute_loss(
            preds=preds,
            targets=targets,
            weights=weights,
            target_nan_mask=self.target_nan_mask,
            multitask_handling=self.multitask_handling,
            loss_fun=self.loss_fun,
        )

        loss = loss / n_steps

        # Iteratively augment data by applying perturbations
        # Accumulate the gradients to be applied to the weights of the network later on
        for _ in range(n_steps - 1):
            loss.backward()
            pert_data = pert.detach() + alpha * torch.sign(pert.grad.detach())
            pert.data = pert_data.data
            pert.grad[:] = 0
            features["feat"] = features["feat"] + pert
            pert_batch["features"] = features
            preds = self.forward(pert_batch)["preds"]
            loss, _ = self.compute_loss(
                preds=preds,
                targets=targets,
                weights=weights,
                target_nan_mask=self.target_nan_mask,
                multitask_handling=self.multitask_handling,
                loss_fun=self.loss_fun,
            )
            loss = loss / n_steps

        device = "cpu" if to_cpu else None
        for key in preds.keys():
            preds[key] = preds[key].detach().to(device=device)
            targets[key] = targets[key].detach().to(device=device)
        if weights is not None:
            weights = weights.detach().to(device=device)

        step_dict = {"preds": preds, "targets": targets, "weights": weights}
        step_dict[f"loss/{step_name}"] = loss.detach().cpu()
        step_dict["loss"] = loss
        step_dict["task_losses"] = task_losses
        return step_dict

    def on_train_batch_start(self, batch: Any, batch_idx: int) -> Optional[int]:
        self.train_batch_start_time = time.time()
        self.skip_log_train_metrics = (self.metrics_every_n_train_steps is None) or (
            (batch_idx % self.metrics_every_n_train_steps) != 0
        )
        return super().on_train_batch_start(batch, batch_idx)

    def on_train_batch_end(self, outputs, batch: Any, batch_idx: int) -> None:
        train_batch_time = time.time() - self.train_batch_start_time  # To be used for throughput calculation

        # Get the metrics that are logged at every step (loss, grad_norm, batch_time, batch_tput)
        concatenated_metrics_logs = {}
        concatenated_metrics_logs["train/loss"] = outputs["loss"]

        # report the training loss for each individual tasks
        for task in self.tasks:
            concatenated_metrics_logs[f"train/loss/{task}"] = outputs["task_losses"][task]

        # get the mean loss value for individual tasks as they are a tensor of size --> gradient accumulation * replication * device_iter
        for key in concatenated_metrics_logs:
            if isinstance(concatenated_metrics_logs[key], torch.Tensor):
                if concatenated_metrics_logs[key].numel() > 1:
                    concatenated_metrics_logs[key] = concatenated_metrics_logs[key][
                        concatenated_metrics_logs[key] != 0
                    ].mean()

        # If logging is skipped for this step, then log the important metrics anyway and return
        if self.skip_log_train_metrics:
            if self.logger is not None:
                self.logger.log_metrics(
                    concatenated_metrics_logs, step=self.global_step
                )  # This is a pytorch lightning function call
            return

        ### The code below is not executed if the logging is skipped for this step ###

        # Get the throughput of the batch
        num_graphs = self.get_num_graphs(batch["features"])
        tput = num_graphs / train_batch_time
        concatenated_metrics_logs["train/batch_time"] = train_batch_time
        concatenated_metrics_logs["train/batch_tput"] = tput

        # Compute all the metrics for the training set
        self.task_epoch_summary.update_predictor_state(
            step_name="train",
            targets=outputs["targets"],
            predictions=outputs["preds"],
            loss=outputs["loss"],  # This is the weighted loss for now, but change to task-specific loss
            task_losses=outputs["task_losses"],
            n_epochs=self.current_epoch,
        )
        metrics_logs = self.task_epoch_summary.get_metrics_logs()  # Dict[task, metric_logs]
        metrics_logs["_global"]["grad_norm"] = self.get_gradient_norm()
        concatenated_metrics_logs.update(metrics_logs)

        # Log the metrics
        if self.logger is not None:
            self.logger.log_metrics(
                concatenated_metrics_logs, step=self.global_step
            )  # This is a pytorch lightning function call

    def training_step(self, batch: Dict[str, Tensor], to_cpu: bool = True) -> Dict[str, Any]:
        step_dict = None

        # Train using FLAG
        if self.flag_kwargs["n_steps"] > 0:
            step_dict = self.flag_step(batch=batch, step_name="train", to_cpu=to_cpu)
        # Train normally, without using FLAG
        elif self.flag_kwargs["n_steps"] == 0:
            # step_dict = self._general_step(batch=batch, step_name="train", to_cpu=True)
            step_dict = self._general_step(batch=batch, step_name="train", to_cpu=to_cpu)

        # Remove the preds and targets if no logging is required
        if self.skip_log_train_metrics:
            step_dict.pop("preds")
            step_dict.pop("targets")
        return step_dict  # Returning the metrics_logs with the loss

    def get_gradient_norm(self):
        # compute the norm
        total_norm = torch.tensor(0.0)
        for p in self.parameters():
            if p.grad is not None:
                param_norm = p.grad.detach().data.norm(2)
                total_norm += param_norm.detach().cpu() ** 2
        total_norm = total_norm**0.5
        return total_norm

    def validation_step(self, batch: Dict[str, Tensor], to_cpu: bool = True) -> Dict[str, Any]:
        return self._general_step(batch=batch, step_name="val", to_cpu=to_cpu)

    def test_step(self, batch: Dict[str, Tensor], to_cpu: bool = True) -> Dict[str, Any]:
        return self._general_step(batch=batch, step_name="test", to_cpu=to_cpu)

    def _general_epoch_end(self, outputs: Dict[str, Any], step_name: str) -> None:
        r"""Common code for training_epoch_end, validation_epoch_end and testing_epoch_end"""
        # Transform the list of dict of dict, into a dict of list of dict
        preds = {}
        targets = {}
        device = device = outputs[0]["preds"][self.tasks[0]].device  # should be better way to do this
        # device = 0
        for task in self.tasks:
            preds[task] = torch.cat([out["preds"][task].to(device=device) for out in outputs], dim=0)
            targets[task] = torch.cat([out["targets"][task].to(device=device) for out in outputs], dim=0)
        if ("weights" in outputs[0].keys()) and (outputs[0]["weights"] is not None):
            weights = torch.cat([out["weights"] for out in outputs], dim=0)
        else:
            weights = None
        loss, task_losses = self.compute_loss(
            preds=preds,
            targets=targets,
            weights=weights,
            target_nan_mask=self.target_nan_mask,
            multitask_handling=self.multitask_handling,
            loss_fun=self.loss_fun,
        )

        self.task_epoch_summary.update_predictor_state(
            step_name=step_name,
            predictions=preds,
            targets=targets,
            loss=loss,
            task_losses=task_losses,
            n_epochs=self.current_epoch,
        )
        metrics_logs = self.task_epoch_summary.get_metrics_logs()
        self.task_epoch_summary.set_results(task_metrics=metrics_logs)

        return metrics_logs  # Consider returning concatenated dict for logging

    def on_train_epoch_start(self) -> None:
        self.epoch_start_time = time.time()

    def on_train_epoch_end(self) -> None:
        if self.epoch_start_time is None:
            logger.warning("epoch timer not initialized")
        else:
            epoch_time = time.time() - self.epoch_start_time
            self.epoch_start_time = None
            self.log("epoch_time", torch.tensor(epoch_time))

    def on_validation_epoch_start(self) -> None:
        self.mean_val_time_tracker.reset()
        self.mean_val_tput_tracker.reset()
        return super().on_validation_epoch_start()

    def on_validation_batch_start(self, batch: Any, batch_idx: int) -> None:
        self.validation_batch_start_time = time.time()
        return super().on_validation_batch_start(batch, batch_idx)

    def on_validation_batch_end(self, outputs: Any, batch: Any, batch_idx: int) -> None:
        val_batch_time = time.time() - self.validation_batch_start_time
        self.validation_step_outputs.append(outputs)
        self.mean_val_time_tracker.update(val_batch_time)
        num_graphs = self.get_num_graphs(batch["features"])
        self.mean_val_tput_tracker.update(num_graphs / val_batch_time)
        return super().on_validation_batch_end(outputs, batch, batch_idx)

    def on_validation_epoch_end(self) -> None:
        metrics_logs = self._general_epoch_end(outputs=self.validation_step_outputs, step_name="val")
        self.validation_step_outputs.clear()
        concatenated_metrics_logs = self.task_epoch_summary.concatenate_metrics_logs(metrics_logs)
        concatenated_metrics_logs["val/mean_time"] = torch.tensor(self.mean_val_time_tracker.mean_value)
        concatenated_metrics_logs["val/mean_tput"] = self.mean_val_tput_tracker.mean_value

        if hasattr(self.optimizers(), "param_groups"):
            lr = self.optimizers().param_groups[0]["lr"]
            concatenated_metrics_logs["lr"] = torch.tensor(lr)
        concatenated_metrics_logs["n_epochs"] = torch.tensor(self.current_epoch, dtype=torch.float32)
        self.log_dict(concatenated_metrics_logs)

        # Save yaml file with the per-task metrics summaries
        full_dict = {}
        full_dict.update(self.task_epoch_summary.get_dict_summary())

    def on_test_batch_end(self, outputs: Any, batch: Any, batch_idx: int) -> None:
        self.test_step_outputs.append(outputs)

    def on_test_epoch_end(self) -> None:
        metrics_logs = self._general_epoch_end(outputs=self.test_step_outputs, step_name="test")
        self.test_step_outputs.clear()
        concatenated_metrics_logs = self.task_epoch_summary.concatenate_metrics_logs(metrics_logs)

        self.log_dict(concatenated_metrics_logs)

        # Save yaml file with the per-task metrics summaries
        full_dict = {}
        full_dict.update(self.task_epoch_summary.get_dict_summary())

    def on_train_start(self):
        hparams_log = deepcopy(self.hparams)
        hparams_log["n_params"] = self.n_params
        if self.logger is not None:
            self.logger.log_hyperparams(hparams_log)

    def get_progress_bar_dict(self) -> Dict[str, float]:
        prog_dict = {}
        prog_dict["loss"] = self.task_epoch_summary.weighted_loss.detach().cpu()
        results_on_progress_bar = self.task_epoch_summary.get_results_on_progress_bar("val")
        for task in self.tasks:
            prog_dict[self.task_epoch_summary.metric_log_name(task, "loss", "val")] = (
                self.task_epoch_summary.task_summaries[task].summaries["val"].loss
            )
            prog_dict.update(results_on_progress_bar)
        return prog_dict

    def __repr__(self) -> str:
        r"""
        Controls how the class is printed
        """
        model_str = self.model.__repr__()
        summary_str = self.summarize().__repr__()

        return model_str + "\n\n" + summary_str

    @staticmethod
    def list_pretrained_models():
        """List available pretrained models."""
        return GRAPHIUM_PRETRAINED_MODELS

    @staticmethod
    def load_pretrained_models(name: str):
        """Load a pretrained model from its name.

        Args:
            name: Name of the model to load. List available
                from `graphium.trainer.PredictorModule.list_pretrained_models()`.
        """

        if name not in GRAPHIUM_PRETRAINED_MODELS:
            raise ValueError(
                f"The model '{name}' is not available. Choose from {set(GRAPHIUM_PRETRAINED_MODELS.keys())}."
            )

        return PredictorModule.load_from_checkpoint(GRAPHIUM_PRETRAINED_MODELS[name])

    def set_max_nodes_edges_per_graph(self, datamodule: BaseDataModule, stages: Optional[List[str]] = None):
        datamodule.setup()

        max_nodes = datamodule.get_max_num_nodes_datamodule(stages)
        max_edges = datamodule.get_max_num_edges_datamodule(stages)

        self.model.set_max_num_nodes_edges_per_graph(max_nodes, max_edges)

    def get_num_graphs(self, data: Batch):
        """
        Method to compute number of graphs in a Batch.
        Essential to estimate throughput in graphs/s.
        """
        return torch.max(data.batch) + 1


# graphium/trainer/predictor_options.py

r"""Data classes to group together related arguments for the creation of a Predictor Module."""


"""
Replace the usage of **kwargs by adding checks to make sure that everything is type-safe.
    Stricter typing is important because:
        - It makes finding bugs easier since incorrect types will cause an obvious error.
        - Static analysis becomes easier, and the IDE can give hints about errors in the code before runtime.
Add the post-init function to do the checks immediately.
"""

from copy import deepcopy
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Type, Union
from inspect import signature, isclass

from torch import nn

from graphium.utils.spaces import LOSS_DICT
from graphium.utils.spaces import SCHEDULER_DICT


@dataclass
class ModelOptions:
    r"""
    This data class stores the arguments necessary to instantiate a model for the Predictor.

    Parameters:
        model_class:
            pytorch module used to create a model

        model_kwargs:
            Key-word arguments used to initialize the model from `model_class`.
    """
    model_class: Type[nn.Module]
    model_kwargs: Dict[str, Any]


@dataclass
class OptimOptions:
    r"""
    This data class stores the arguments necessary to configure the optimizer for the Predictor.

    Parameters:
        optim_kwargs:
            Dictionnary used to initialize the optimizer, with possible keys below.

            - lr `float`: Learning rate (Default=`1e-3`)
            - weight_decay `float`: Weight decay used to regularize the optimizer (Default=`0.`)

        torch_scheduler_kwargs:
            Dictionnary for the scheduling of learning rate, with possible keys below.

            - type `str`: Type of the learning rate to use from pytorch. Examples are
                `'ReduceLROnPlateau'` (default), `'CosineAnnealingWarmRestarts'`, `'StepLR'`, etc.
            - **kwargs: Any other argument for the learning rate scheduler

        scheduler_kwargs:
            Dictionnary for the scheduling of the learning rate modification used by pytorch-lightning

            - monitor `str`: metric to track (Default=`"loss/val"`)
            - interval `str`: Whether to look at iterations or epochs (Default=`"epoch"`)
            - strict `bool`: if set to True will enforce that value specified in monitor is available
                while trying to call scheduler.step(), and stop training if not found. If False will
                only give a warning and continue training (without calling the scheduler). (Default=`True`)
            - frequency `int`: **TODO: NOT REALLY SURE HOW IT WORKS!** (Default=`1`)

        scheduler_class: The class to use for the scheduler, or the str representing the scheduler.

    """
    optim_kwargs: Optional[Dict[str, Any]] = None
    torch_scheduler_kwargs: Optional[Dict[str, Any]] = None
    scheduler_kwargs: Optional[Dict[str, Any]] = None
    scheduler_class: Optional[Union[str, Type]] = None

    # Instead of passing a dictionary to be processed by the predictor,
    # this class will process the dictionary in advance and return the optimizer
    def set_kwargs(self):
        # Set the parameters and default value for the optimizer, and check values
        if self.optim_kwargs is None:
            self.optim_kwargs = {}
        self.optim_kwargs.setdefault("lr", 1e-3)
        self.optim_kwargs.setdefault("weight_decay", 0.0)
        assert self.optim_kwargs["lr"] > 0
        assert self.optim_kwargs["weight_decay"] >= 0

        # Set the lightning scheduler
        if self.scheduler_kwargs is None:
            self.scheduler_kwargs = {}
        self.scheduler_kwargs.setdefault("interval", "epoch")
        self.scheduler_kwargs.setdefault("monitor", "loss/val")
        self.scheduler_kwargs.setdefault("mode", "min")
        self.scheduler_kwargs.setdefault("frequency", 1)
        self.scheduler_kwargs.setdefault("strict", True)

        # Set the pytorch scheduler arguments
        if self.torch_scheduler_kwargs is None:
            self.torch_scheduler_kwargs = {}
        self.torch_scheduler_kwargs.setdefault("module_type", "ReduceLROnPlateau")

        # Get the class for the scheduler
        scheduler_class = self.torch_scheduler_kwargs.pop("module_type", None)
        if self.scheduler_class is None:
            if isinstance(scheduler_class, str):
                self.scheduler_class = SCHEDULER_DICT[scheduler_class]
            elif isclass(scheduler_class):
                self.scheduler_class = scheduler_class
            else:
                raise TypeError("`scheduler_class` should be a str or a class")

        # Add the `monitor` and `mode` variables
        sig = signature(self.scheduler_class.__init__)
        key_args = [p.name for p in sig.parameters.values()]
        if "monitor" in key_args:
            self.torch_scheduler_kwargs.setdefault("monitor", self.scheduler_kwargs["monitor"])
        if "mode" in key_args:
            self.torch_scheduler_kwargs.setdefault("mode", self.scheduler_kwargs["mode"])


@dataclass
class EvalOptions:
    r"""
    This data class stores the arguments necessary to instantiate a model for the Predictor.

    Parameters:
        loss_fun:
            Loss function used during training.
            Acceptable strings are graphium.utils.spaces.LOSS_DICT.keys().
            If a dict, must contain a 'name' key with one of the acceptable loss function strings
            as a value. The rest of the dict will be used as the arguments passed to the loss object.
            Otherwise, a callable object must be provided, with a method `loss_fun._get_name()`.

        metrics:
            A dictionnary of metrics to compute on the prediction, other than the loss function.
            These metrics will be logged into WandB or other.

        metrics_on_progress_bar:
            The metrics names from `metrics` to display also on the progress bar of the training

        metrics_on_training_set:
            The metrics names from `metrics` to be computed on the training set for each iteration.
            If `None`, all the metrics are computed. Using less metrics can significantly improve
            performance, depending on the number of readouts.
    """
    loss_fun: Union[str, Dict, Callable]
    metrics: Dict[str, Callable] = None
    metrics_on_progress_bar: List[str] = field(default_factory=List[str])
    metrics_on_training_set: Optional[List[str]] = None

    def check_metrics_validity(self):
        """
        Check that the metrics for the progress_par and training_set are valid
        """
        if self.metrics_on_progress_bar is not None:
            selected = set(self.metrics_on_progress_bar)
            assert selected.issubset(
                set(self.metrics.keys())
            ), f"Metrics {selected - set(self.metrics.keys())} not in `metrics` with choices {set(self.metrics.keys())}"

        if self.metrics_on_training_set is not None:
            selected = set(self.metrics_on_training_set)
            assert selected.issubset(
                set(self.metrics.keys())
            ), f"Metrics {selected - set(self.metrics.keys())} not in `metrics` with choices {set(self.metrics.keys())}"

    # Parse before or after?
    @staticmethod
    def parse_loss_fun(loss_fun: Union[str, Dict, Callable]) -> Callable:
        r"""
        Parse the loss function from a string or a dict

        Parameters:
            loss_fun:
                A callable corresponding to the loss function, a string specifying the loss
                function from `LOSS_DICT`, or a dict containing a key 'name' specifying the
                loss function, and the rest of the dict used as the arguments for the loss.
                Accepted strings are: graphium.utils.spaces.LOSS_DICT.keys().

        Returns:
            Callable:
                Function or callable to compute the loss, takes `preds` and `targets` as inputs.
        """

        if isinstance(loss_fun, str):
            if loss_fun not in LOSS_DICT.keys():
                raise ValueError(
                    f"`loss_fun` expected to be one of the strings in {LOSS_DICT.keys()}. "
                    f"Provided: {loss_fun}."
                )
            loss_fun = LOSS_DICT[loss_fun]()
        elif isinstance(loss_fun, dict):
            if loss_fun.get("name") is None:
                raise ValueError(f"`loss_fun` expected to have a key 'name'.")
            if loss_fun["name"] not in LOSS_DICT.keys():
                raise ValueError(
                    f"`loss_fun['name']` expected to be one of the strings in {LOSS_DICT.keys()}. "
                    f"Provided: {loss_fun}."
                )
            loss_fun = deepcopy(loss_fun)
            loss_name = loss_fun.pop("name")
            loss_fun = LOSS_DICT[loss_name](**loss_fun)
        elif not callable(loss_fun):
            raise ValueError(f"`loss_fun` must be `str`, `dict` or `callable`. Provided: {type(loss_fun)}")

        return loss_fun


@dataclass
class FlagOptions:
    r"""
    This data class stores the arguments necessary to instantiate a model for the Predictor.

    Parameters:
        flag_kwargs:
            Keyword arguments used for FLAG, and adversarial data augmentation for graph networks.
            See: https://arxiv.org/abs/2010.09891

            - n_steps: An integer that specifies the number of ascent steps when running FLAG during training.
                Default value of 0 trains GNNs without FLAG, and any value greater than 0 will use FLAG with that
                many iterations.

            - alpha: A float that specifies the ascent step size when running FLAG. Default=0.01
    """
    flag_kwargs: Dict[str, Any] = None

    # Set the parameters and default values for the FLAG adversarial augmentation, and check values
    def set_kwargs(self):
        if self.flag_kwargs is None:
            self.flag_kwargs = {}
        self.flag_kwargs.setdefault("alpha", 0.01)
        self.flag_kwargs.setdefault("n_steps", 0)
        assert isinstance(self.flag_kwargs["n_steps"], int) and (self.flag_kwargs["n_steps"] >= 0)
        assert self.flag_kwargs["alpha"] >= 0


# graphium/trainer/predictor_summaries.py

r"""Classes to store information about resulting evaluation metrics when using a Predictor Module."""

from typing import Any, Callable, Dict, List, Optional, Union
from loguru import logger

import numpy as np
import torch
from torch import Tensor

from graphium.utils.tensor import nan_mean, nan_std, nan_median


class SummaryInterface(object):
    r"""
    An interface to define the functions implemented by summary classes that implement SummaryInterface.
    """

    def set_results(self, **kwargs):
        raise NotImplementedError()

    def get_dict_summary(self):
        raise NotImplementedError()

    def update_predictor_state(self, **kwargs):
        raise NotImplementedError()

    def get_metrics_logs(self, **kwargs):
        raise NotImplementedError()


class Summary(SummaryInterface):
    # TODO (Gabriela): Default argument cannot be []
    def __init__(
        self,
        loss_fun: Union[str, Callable],
        metrics: Dict[str, Callable],
        metrics_on_training_set: List[str] = [],
        metrics_on_progress_bar: List[str] = [],
        monitor: str = "loss",
        mode: str = "min",
        task_name: Optional[str] = None,
    ):
        r"""
        A container to be used by the Predictor Module that stores the results for the given metrics on the predictions and targets provided.
        Parameters:
            loss_fun:
            Loss function used during training. Acceptable strings are 'mse', 'bce', 'mae', 'cosine'.
            Otherwise, a callable object must be provided, with a method `loss_fun._get_name()`.

            metrics:
            A dictionnary of metrics to compute on the prediction, other than the loss function.
            These metrics will be logged into WandB or similar.

            metrics_on_training_set:
            The metrics names from `metrics` to be computed on the training set for each iteration.
            If `None`, all the metrics are computed. Using less metrics can significantly improve
            performance, depending on the number of readouts.

            metrics_on_progress_bar:
            The metrics names from `metrics` to display also on the progress bar of the training

            monitor:
            `str` metric to track (Default=`"loss/val"`)

            task_name:
            name of the task (Default=`None`)

        """
        self.loss_fun = loss_fun
        self.metrics = metrics
        self.metrics_on_training_set = metrics_on_training_set
        self.metrics_on_progress_bar = metrics_on_progress_bar
        self.monitor = monitor
        self.mode = mode

        self.summaries = {}
        self.best_summaries = {}

        # Current predictor state
        # self.predictor_outputs = None
        self.step_name: str = None
        self.targets: Tensor = None
        self.predictions: Tensor = None
        self.loss = None  # What type?
        self.n_epochs: int = None

        self.task_name = task_name
        self.logged_metrics_exceptions = []  # Track which metric exceptions have been logged

    def update_predictor_state(
        self, step_name: str, targets: Tensor, predictions: Tensor, loss: Tensor, n_epochs: int
    ):
        r"""
        update the state of the predictor
        Parameters:
            step_name: which stage you are in, e.g. "train"
            targets: the targets tensor
            predictions: the predictions tensor
            loss: the loss tensor
            n_epochs: the number of epochs
        """
        self.step_name = step_name
        self.targets = targets
        self.predictions = predictions
        self.loss = loss
        self.n_epochs = n_epochs

    def set_results(
        self,
        metrics: Dict[str, Tensor],
    ):
        r"""
        set the reults from the metrics
        [!] This function requires that self.update_predictor_state() be called before it.
        Parameters:
            metrics: a dictionary of metrics
        """

        # Include the task_name in the loss for logging, and similarly for other metrics
        metrics[self.metric_log_name(self.task_name, "loss", self.step_name)] = self.loss
        self.summaries[self.step_name] = Summary.Results(
            targets=self.targets,
            predictions=self.predictions,
            loss=self.loss,
            metrics=metrics,  # Should include task name from get_metrics_logs()
            monitored_metric=f"{self.monitor}/{self.step_name}",  # Include task name?
            n_epochs=self.n_epochs,
        )
        if self.is_best_epoch(self.step_name, self.loss, metrics):
            self.best_summaries[self.step_name] = self.summaries[self.step_name]

    def is_best_epoch(self, step_name: str, loss: Tensor, metrics: Dict[str, Tensor]) -> bool:
        r"""
        check if the current epoch is the best epoch based on self.mode criteria
        Parameters:
            step_name: which stage you are in, e.g. "train"
            loss: the loss tensor
            metrics: a dictionary of metrics
        """

        # TODO (Gabriela): Check for bugs related to monitor_name
        if not (step_name in self.best_summaries.keys()):
            return True

        # Include the task_name in the loss for logging, and similarly for other metrics
        metrics[self.metric_log_name(self.task_name, "loss", self.step_name)] = loss
        monitor_name = f"{self.monitor}/{step_name}"  # Include task_name?
        if (
            not monitor_name in self.best_summaries.keys()
        ):  # Feels like there's a bug here. What is this trying to do???
            return True

        if self.mode == "max":
            return metrics[monitor_name] > self.best_summaries[step_name].monitored
        elif self.mode == "min":
            return metrics[monitor_name] < self.best_summaries[step_name].monitored
        else:
            ValueError(f"Mode must be 'min' or 'max', provided `{self.mode}`")

    def get_results(
        self,
        step_name: str,
    ):
        r"""
        retrieve the results for a given step
        Parameters:
            step_name: which stage you are in, e.g. "train"
        Returns:
            the results for the given step
        """
        return self.summaries[step_name]

    def get_best_results(
        self,
        step_name: str,
    ):
        r"""
        retrieve the best results for a given step
        Parameters:
            step_name: which stage you are in, e.g. "train"
        Returns:
            the best results for the given step
        """
        return self.best_summaries[step_name]

    def get_results_on_progress_bar(
        self,
        step_name: str,
    ) -> Dict[str, Tensor]:
        r"""
        retrieve the results to be displayed on the progress bar for a given step
        Parameters:
            step_name: which stage you are in, e.g. "train"
        Returns:
            the results to be displayed on the progress bar for the given step
        """
        results = self.summaries[step_name]
        results_prog = {
            # f"{kk}/{step_name}": results.metrics[f"{kk}/{step_name}"] for kk in self.metrics_on_progress_bar
            self.metric_log_name(self.task_name, kk, step_name): results.metrics[
                self.metric_log_name(self.task_name, kk, step_name)
            ]
            for kk in self.metrics_on_progress_bar
        }
        return results_prog

    def get_dict_summary(self) -> Dict[str, Any]:
        r"""
        retrieve the full summary in a dictionary
        Returns:
            the full summary in a dictionary
        """
        full_dict = {}
        # Get metric summaries
        full_dict["metric_summaries"] = {}
        for key, val in self.summaries.items():
            full_dict["metric_summaries"][key] = {k: v for k, v in val.metrics.items()}
            full_dict["metric_summaries"][key]["n_epochs"] = val.n_epochs

        # Get metric summaries at best epoch
        full_dict["best_epoch_metric_summaries"] = {}
        for key, val in self.best_summaries.items():
            full_dict["best_epoch_metric_summaries"][key] = val.metrics
            full_dict["best_epoch_metric_summaries"][key]["n_epochs"] = val.n_epochs

        return full_dict

    def get_metrics_logs(self) -> Dict[str, Any]:
        r"""
        Get the data about metrics to log.
        Note: This function requires that self.update_predictor_state() be called before it.
        Returns:
            A dictionary of metrics to log.
        """
        targets = self.targets.to(dtype=self.predictions.dtype, device=self.predictions.device)
        # Compute the metrics always used in regression tasks
        metric_logs = {}
        metric_logs[self.metric_log_name(self.task_name, "mean_pred", self.step_name)] = nan_mean(
            self.predictions
        )
        metric_logs[self.metric_log_name(self.task_name, "std_pred", self.step_name)] = nan_std(
            self.predictions
        )
        metric_logs[self.metric_log_name(self.task_name, "median_pred", self.step_name)] = nan_median(
            self.predictions
        )
        metric_logs[self.metric_log_name(self.task_name, "mean_target", self.step_name)] = nan_mean(targets)
        metric_logs[self.metric_log_name(self.task_name, "std_target", self.step_name)] = nan_std(targets)
        metric_logs[self.metric_log_name(self.task_name, "median_target", self.step_name)] = nan_median(
            targets
        )
        if torch.cuda.is_available():
            metric_logs[f"gpu_allocated_GB"] = torch.tensor(torch.cuda.memory_allocated() / (2**30))

        # Specify which metrics to use
        metrics_to_use = self.metrics
        if self.step_name == "train":
            metrics_to_use = {
                key: metric for key, metric in metrics_to_use.items() if key in self.metrics_on_training_set
            }
        # Compute the additional metrics
        for key, metric in metrics_to_use.items():
            metric_name = self.metric_log_name(
                self.task_name, key, self.step_name
            )  # f"{key}/{self.step_name}"
            try:
                metric_logs[metric_name] = metric(self.predictions, targets)
            except Exception as e:
                metric_logs[metric_name] = torch.as_tensor(float("nan"))
                # Warn only if it's the first warning for that metric
                if metric_name not in self.logged_metrics_exceptions:
                    self.logged_metrics_exceptions.append(metric_name)
                    logger.warning(f"Error for metric {metric_name}. NaN is returned. Exception: {e}")

        # Convert all metrics to CPU, except for the loss
        # metric_logs[f"{self.loss_fun._get_name()}/{self.step_name}"] = self.loss.detach().cpu()
        metric_logs[
            self.metric_log_name(self.task_name, self.loss_fun._get_name(), self.step_name)
        ] = self.loss.detach().cpu()
        # print("Metrics logs keys: ", metric_logs.keys())
        metric_logs = {key: metric.detach().cpu() for key, metric in metric_logs.items()}

        return metric_logs

    def metric_log_name(self, task_name, metric_name, step_name):
        if task_name is None:
            return f"{metric_name}/{step_name}"
        else:
            return f"{task_name}/{metric_name}/{step_name}"

    class Results:
        def __init__(
            self,
            targets: Tensor = None,
            predictions: Tensor = None,
            loss: float = None,  # Is this supposed to be a Tensor or float?
            metrics: dict = None,
            monitored_metric: str = None,
            n_epochs: int = None,
        ):
            r"""
            This inner class is used as a container for storing the results of the summary.
            Parameters:
                targets: the targets
                predictions: the prediction tensor
                loss: the loss, float or tensor
                metrics: the metrics
                monitored_metric: the monitored metric
                n_epochs: the number of epochs
            """
            self.targets = targets.detach().cpu()
            self.predictions = predictions.detach().cpu()
            self.loss = loss.item() if isinstance(loss, Tensor) else loss
            self.monitored_metric = monitored_metric
            if monitored_metric in metrics.keys():
                self.monitored = metrics[monitored_metric].detach().cpu()
            self.metrics = {
                key: value.tolist() if isinstance(value, (Tensor, np.ndarray)) else value
                for key, value in metrics.items()
            }
            self.n_epochs = n_epochs


class TaskSummaries(SummaryInterface):
    def __init__(
        self,
        task_loss_fun: Callable,
        task_metrics: Dict[str, Callable],
        task_metrics_on_training_set: List[str],
        task_metrics_on_progress_bar: List[str],
        monitor: str = "loss",
        mode: str = "min",
    ):
        r"""
        class to store the summaries of the tasks
        Parameters:
            task_loss_fun: the loss function for each task
            task_metrics: the metrics for each task
            task_metrics_on_training_set: the metrics to use on the training set
            task_metrics_on_progress_bar: the metrics to use on the progress bar
            monitor: the metric to monitor
            mode: the mode of the metric to monitor
        """
        self.task_loss_fun = task_loss_fun
        self.task_metrics = task_metrics
        self.task_metrics_on_progress_bar = task_metrics_on_progress_bar
        self.task_metrics_on_training_set = task_metrics_on_training_set
        self.monitor = monitor
        self.mode = mode

        self.task_summaries: Dict[str, Summary] = {}
        self.task_best_summaries: Dict[str, Summary] = {}
        self.tasks = list(task_loss_fun.keys())

        for task in self.tasks:
            self.task_summaries[task] = Summary(
                self.task_loss_fun[task],
                self.task_metrics[task],
                self.task_metrics_on_training_set[task],
                self.task_metrics_on_progress_bar[task],
                self.monitor,
                self.mode,
                task_name=task,
            )

        # Current predictor state
        self.weighted_loss = None
        self.step_name = None

    def update_predictor_state(
        self,
        step_name: str,
        targets: Dict[str, Tensor],
        predictions: Dict[str, Tensor],
        loss: Tensor,
        task_losses: Dict[str, Tensor],
        n_epochs: int,
    ):
        r"""
        update the state for all predictors
        Parameters:
            step_name: the name of the step
            targets: the target tensors
            predictions: the prediction tensors
            loss: the loss tensor
            task_losses: the task losses
            n_epochs: the number of epochs
        """
        self.weighted_loss = loss
        self.step_name = step_name
        for task in self.tasks:
            self.task_summaries[task].update_predictor_state(
                step_name,
                targets[task],
                predictions[task].detach(),
                task_losses[task].detach(),
                n_epochs,
            )

    def set_results(self, task_metrics: Dict[str, Dict[str, Tensor]]):
        """
        set the results for all tasks
        Parameters:
            task_metrics: the metrics for each task
        """
        for task in self.tasks:
            self.task_summaries[task].set_results(task_metrics[task])
            step_name = self.task_summaries[task].step_name
            loss = self.task_summaries[task].loss
            if self.task_summaries[task].is_best_epoch(step_name, loss, task_metrics[task]):
                self.task_summaries[task].best_summaries[step_name] = self.task_summaries[task].summaries[
                    step_name
                ]

    def get_results(
        self,
        step_name: str,
    ) -> Dict[str, Dict[str, Any]]:
        """
        retrieve the results
        Parameters:
            step_name: the name of the step, i.e. "train"
        Returns:
            the results
        """
        results = {}
        for task in self.tasks:
            results[task] = self.task_summaries[task].get_results(step_name)
        return results

    def get_best_results(
        self,
        step_name: str,
    ) -> Dict[str, Dict[str, Any]]:
        """
        retrieve the best results
        Parameters:
            step_name: the name of the step, i.e. "train"
        Returns:
            the best results
        """
        results = {}
        for task in self.tasks:
            results[task] = self.task_summaries[task].get_best_results(step_name)
        return results

    def get_results_on_progress_bar(
        self,
        step_name: str,
    ) -> Dict[str, Dict[str, Any]]:
        r"""
        return all results from all tasks for the progress bar
        Combine the dictionaries. Instead of having keys as task names, we merge all the task-specific dictionaries.
        Parameters:
            step_name: the name of the step
        Returns:
            the results for the progress bar
        """
        task_results_prog = {}
        for task in self.tasks:
            # task_results_prog[task] = self.task_summaries[task].get_results_on_progress_bar(step_name)
            task_results_prog.update(self.task_summaries[task].get_results_on_progress_bar(step_name))
        return task_results_prog

    def get_dict_summary(
        self,
    ) -> Dict[str, Dict[str, Any]]:
        r"""
        get task summaries in a dictionary
        Returns:
            the task summaries
        """
        task_full_dict = {}
        for task in self.tasks:
            task_full_dict[task] = self.task_summaries[task].get_dict_summary()
        return task_full_dict

    def get_metrics_logs(
        self,
    ) -> Dict[str, Dict[str, Tensor]]:
        r"""
        get the logs for the metrics
        Returns:
            the task logs for the metrics
        """
        task_metrics_logs = {}
        for task in self.tasks:
            task_metrics_logs[task] = self.task_summaries[task].get_metrics_logs()

        # Include global (weighted loss)
        task_metrics_logs["_global"] = {}
        task_metrics_logs["_global"][f"loss/{self.step_name}"] = self.weighted_loss.detach().cpu()
        return task_metrics_logs

    # TODO (Gabriela): This works to fix the logging on TB, but make it more efficient
    def concatenate_metrics_logs(
        self,
        metrics_logs: Dict[str, Dict[str, Tensor]],
    ) -> Dict[str, Tensor]:
        r"""
        concatenate the metrics logs
        Parameters:
            metrics_logs: the metrics logs
        Returns:
            the concatenated metrics logs
        """
        concatenated_metrics_logs = {}
        for task in list(self.tasks) + ["_global"]:
            concatenated_metrics_logs.update(metrics_logs[task])
        concatenated_metrics_logs[f"loss/{self.step_name}"] = self.weighted_loss.detach().cpu()
        return concatenated_metrics_logs

    def metric_log_name(
        self,
        task_name: str,
        metric_name: str,
        step_name: str,
    ) -> str:
        r"""
        print the metric name, task name and step name
        Returns:
            the metric name, task name and step name
        """
        if task_name is None:
            return f"{metric_name}/{step_name}"
        else:
            return f"{task_name}/{metric_name}/{step_name}"


# graphium/utils/__init__.py

from . import fs
from . import tensor
from . import dict_tensor


# graphium/utils/arg_checker.py

""" Argument checker module """
import collections
import numpy as np

# Global variable of accepted string types
KNOWN_TYPES = {
    "none": None,
    "str": str,
    "list": list,
    "tuple": tuple,
    "dict": dict,
    "int": int,
    "float": float,
    "complex": complex,
    "bool": bool,
    "callable": callable,
}


def _parse_type(type_to_validate, accepted_types):
    # Check if the provided type is accepted
    if (type_to_validate is not None) and (not isinstance(type_to_validate, accepted_types)):
        raise TypeError(
            "type_to_validate should be None, type or str. {} provided".format(type(type_to_validate))
        )
    if isinstance(type_to_validate, str):
        type_to_validate = type_to_validate.lower()
        if type_to_validate in KNOWN_TYPES.keys():
            type_to_validate = KNOWN_TYPES[type_to_validate]
        else:
            raise TypeError(
                "type_to_validate is not a known type. Known types are :"
                " \n{}\n Provided : \n{}".format(KNOWN_TYPES.keys(), type_to_validate)
            )
    return type_to_validate


def _enforce_iter_type(arg, enforce_type):
    # Cast the arg to be either a list or a tuple
    if enforce_type is not None:
        if (enforce_type == list) and (not isinstance(arg, list)):
            arg = list(arg)
        elif (enforce_type == tuple) and (not isinstance(arg, tuple)):
            arg = tuple(arg)
        elif enforce_type not in (list, tuple):
            raise TypeError('enforce_type should be None, "list" or "tuple", but is {}'.format(enforce_type))
    return arg


def check_arg_iterator(arg, enforce_type=None, enforce_subtype=None, cast_subtype: bool = True):
    r"""
    Verify if the type is an iterator. If it is `None`, convert to an empty list/tuple. If it is
    not a list/tuple/str, try to convert to an iterator. If it is a str or cannot be converted to
    an iterator, then put the `arg` inside an iterator.
    Possibly enforce the iterator type to `list` or `tuple`, if `enfoce_type` is not None.
    Possibly enforce the subtype to any given type if `enforce_subtype` is not None,
    and decide whether to cast the subtype or to throw an error.

    Parameters:
        arg (any type):
            The input to verify/convert to an iterator (list or tuple). If None, an empty iterator
            is returned.
        enforce_type (str or type):
            The type to enforce the iterator. The valid choices are :
            `None`, `list`, `tuple`, `'none'`, `'list'`, `'tuple'`.
            If `None`, then the iterator type is not enforced.

        enforce_subtype (type, np.dtype or str representing basic type):
            Verify if all the elements inside the iterator are the desired type.
            If `None`, then the sub-type is not enforced.
            Accepted strings are ['none', 'str', 'list', 'tuple', 'dict', 'int',
            'float', 'complex', 'bool', 'callable']

        cast_subtype:
            If True, then the type specified by `enforce_subtype` is used to cast the
            elements inside the iterator. If False, then an error is thrown if the
            types do not match.

    Returns:
        output (iterator):
            An iterator based on the input of the desired type (list or tuple) and
            the desired subtypes.

    """

    # If not list or tuple, put into a list
    if arg is None:
        arg = []
    elif isinstance(arg, str):
        arg = [arg]
    elif isinstance(arg, tuple):
        if enforce_type is None:
            enforce_type = tuple
        arg = list(arg)
    elif not isinstance(arg, (tuple, list)):
        try:
            arg = list(arg)
        except Exception:
            arg = [arg]

    output = arg

    # Make sure that enforce_type and enforce_subtype are a good inputs
    enforce_type = _parse_type(enforce_type, (type, str))
    enforce_subtype = _parse_type(enforce_subtype, (type, str, np.dtype))

    # Cast all the subtypes of the list/tuple into the desired subtype
    if enforce_subtype is not None:
        if enforce_type is None:
            arg2 = output
        elif not isinstance(output, enforce_type):
            arg2 = list(output)
        else:
            arg2 = output
        try:
            for idx, a in enumerate(output):
                if not isinstance(a, enforce_subtype):
                    if cast_subtype:
                        arg2[idx] = enforce_subtype(a)
                    else:
                        raise TypeError(
                            "iter subtype is {}, desired subtype is {}, "
                            "but cast_subtype is set to False".format(type(arg2[idx]), enforce_subtype)
                        )
        except Exception as e:
            raise TypeError(
                "iterator subtype is {} and cannot be casted to {}\n{}".format(type(a), enforce_subtype, e)
            )

        output = _enforce_iter_type(arg2, enforce_type)

    output = _enforce_iter_type(output, enforce_type)

    return output


def check_list1_in_list2(list1, list2, throw_error=True):
    r"""
    Verify if the list1 (iterator) is included in list2 (iterator). If not, raise an error.

    Parameters:
        list1, list2: list, tuple or object
            A list or tuple containing the elements to verify the inclusion.
            If an object is provided other than a list or tuple,
            then it is considered as a list of a single element.
        throw_error: bool
            Whether to throw an error if list1 is not in list2

    Returns:
        list1_in_list2: bool
            A boolean representing the inclusion of list1 in list2. It is returned if
            throw_error is set to false


    """

    list1 = check_arg_iterator(list1)
    list2 = check_arg_iterator(list2)

    # If all elements of list1 are not in list2, throw an error
    list1_in_list2 = all(elem in list2 for elem in list1)
    if not list1_in_list2 and throw_error:
        raise ValueError(
            ("Elements in list1 should be contained in list2." + "\n\nlist1 = {} \n\n list2 = {}").format(
                list1, list2
            )
        )

    return list1_in_list2


def check_columns_choice(dataframe, columns_choice, extra_accepted_cols=None, enforce_type="list"):
    r"""
    Verify if the choice of column `columns_choice` is inside the dataframe or
    the extra_accepted_cols. Otherwise, errors are thrown by the sub-functions.

    Parameters:
        dataframe: (pd.DataFrame)
            The dataframe on which to verify if the column choice is valid.
            columns_choice: str, iterator(str)
            The columns chosen from the dataframe
        extra_accepted_cols: str, iterator(str)
            A list

        enforce_type: str or type
            The type to enforce the iterator. The valid choices are :
            `None`, `list`, `tuple`, `'none'`, `'list'`, `'tuple'`.
            If `None`, then the iterator type is not enforced.


    Returns:
        output: iterator
            A str iterator based on the input of the desired type (list or tuple)

    """
    extra_accepted_cols = [] if extra_accepted_cols is None else extra_accepted_cols
    valid_columns = list(dataframe.columns)
    kwargs_iterator = {
        "enforce_type": enforce_type,
        "enforce_subtype": None,
        "cast_subtype": False,
    }
    columns_choice = check_arg_iterator(columns_choice, **kwargs_iterator)
    extra_accepted_cols = check_arg_iterator(extra_accepted_cols, **kwargs_iterator)
    valid_columns = check_arg_iterator(valid_columns, **kwargs_iterator)
    valid_columns_full = valid_columns + extra_accepted_cols
    check_list1_in_list2(columns_choice, valid_columns_full)

    return columns_choice


# graphium/utils/command_line_utils.py

import re
from collections import defaultdict
from typing import List, Dict


def get_anchors_and_aliases(filepath):
    """Utility function to extract anchors and aliases from YAML config file

    In a YAML file we can specify an anchor as  `some_name: &anchor_name value`
    This is then picked up with alises in the rest of the config as
    `other_name: *anchor_name`
    Using this format in the YAML file all anchors and aliases will be extracted

    Args:
        filepath (str): path to the config file

    Returns:
        anchors (Dict): A dictionary containing the YAML paths of the anchors
                        as keys, and a list of YAML paths to alises.
    """
    anchors = defaultdict(list)
    current_level = {}
    anchor_to_path = {}

    with open(filepath, "r") as file:
        for line in file:
            indent = len(line) - len(line.lstrip(" "))
            key_match = re.search(r"(\w+):", line)
            anchor_match = re.search(r"&(\w+)", line)
            alias_match = re.search(r"\*(\w+)", line)
            if key_match:
                key = key_match.group(1)
                # Compute the full path of the current key.
                full_path = ".".join(
                    [current_level[i] for i in sorted(current_level.keys()) if i < indent] + [key]
                )
                current_level[indent] = key
                # Remove any keys that are indented more than the current line.
                keys_to_remove = [i for i in current_level if i > indent]
                for i in keys_to_remove:
                    del current_level[i]
            else:
                full_path = ".".join([current_level[i] for i in sorted(current_level.keys())])
            if anchor_match:
                anchor = anchor_match.group(1)
                anchor_to_path[anchor] = full_path
            if alias_match:
                alias = alias_match.group(1)
                if alias in anchor_to_path:
                    anchors[anchor_to_path[alias]].append(full_path)
    return anchors


def update_config(cfg: Dict, unknown: List, anchors: List):
    """
    Update the configuration dictionary with command line arguments.
    """
    for arg in unknown:
        if arg.startswith("--"):
            key, value = arg[2:].split("=")
            if key in anchors.keys():
                all_refs = anchors[key]
            else:
                all_refs = []
            all_refs.append(key)
            for key in all_refs:
                keys = key.split(".")
                temp_cfg = cfg
                for k in keys[:-1]:
                    temp_cfg = temp_cfg[k]
                temp_cfg[keys[-1]] = type(temp_cfg[keys[-1]])(value) if keys[-1] in temp_cfg else value

    return cfg


# graphium/utils/custom_lr.py

import warnings
from torch.optim.lr_scheduler import _LRScheduler


class WarmUpLinearLR(_LRScheduler):
    """Custom linear learning rate with warmup.

    Args:
        optimizer (Optimizer): Wrapped optimizer.
        max_num_epochs (int): Maximum number of epochs.
        warmup_epochs (int): The number of epochs for learning rate warmup. Default: 0.
        min_lr (float): the minimum learning rate value. Default: 0.
        last_epoch (int): The index of last epoch. Default: -1.
        verbose (bool): If ``True``, prints a message to stdout for
            each update. Default: ``False``.
    """

    def __init__(self, optimizer, max_num_epochs, warmup_epochs=0, min_lr=0, last_epoch=-1, verbose=False):
        self.max_num_epochs = max_num_epochs
        self.warmup_epochs = warmup_epochs
        self.min_lr = min_lr
        self.last_epoch = last_epoch
        super(WarmUpLinearLR, self).__init__(optimizer, last_epoch, verbose)

    def get_lr(self):
        if not self._get_lr_called_within_step:
            warnings.warn(
                "To get the last learning rate computed by the scheduler, " "please use `get_last_lr()`.",
                UserWarning,
            )
        if self.warmup_epochs > 0 and (self.last_epoch + 1) < self.warmup_epochs:
            return [(self.last_epoch + 1) * base_lr / self.warmup_epochs for base_lr in self.base_lrs]
        else:
            # check epoch_diff in case there is a division by zero error
            epoch_diff = self.max_num_epochs - self.warmup_epochs
            if epoch_diff <= 0:
                factor = 0
            else:
                factor = ((self.last_epoch + 1) - self.warmup_epochs) / epoch_diff
            return [factor * self.min_lr + (1 - factor) * base_lr for base_lr in self.base_lrs]

    def _get_closed_form_lr(self):
        lr_list = []
        for base_lr in self.base_lrs:
            if self.warmup_epochs > 0 and (self.last_epoch + 1) < self.warmup_epochs:
                lr = (self.last_epoch + 1) * base_lr / self.warmup_epochs
            else:
                factor = ((self.last_epoch + 1) - self.warmup_epochs) / (
                    self.max_num_epochs - self.warmup_epochs
                )
                lr = factor * self.min_lr + (1 - factor) * base_lr
            lr_list.append(lr)
        return lr_list


# graphium/utils/decorators.py

class classproperty(property):
    r"""
    Decorator used to declare a class property, defined for the class
    without needing to instanciate an object.

    !!! Example

        ``` python linenums="1"
            @classproperty
            def my_class_property(cls):
                return 5
        ```
    """

    def __get__(self, cls, owner):
        return classmethod(self.fget).__get__(None, owner)()


# graphium/utils/dict_tensor.py

from __future__ import annotations
from typing import Union, Any, Callable, Dict, Hashable
from torch import Tensor


class DictTensor(dict):
    """
    A class that combines the functionality of `dict` and `torch.Tensor`.
    Specifically, it is a dict of Tensor, but it has all the methods and attributes of a Tensor.
    When a given method or attribute is called, it will be called on each element
    of the dict, and a new dict is returned.

    All methods from `torch.Tensor` and other functions are expected to work on `DictTensor`
    with the following rules:
        - The function receives one `DictTensor`: The function will be applied on
          all values of the dictionary. Examples are `torch.sum` or `torch.abs`.
        - The function receives two `DictTensor`: The function will be applied on
          each pair of Tensor with matching keys. If the keys don't match, an error
          will be thrown. Example are operators such as `dict_tensor1 + dict_tensor2`
          or `dict_tensor1 == dict_tensor2`.

    The output of the functions that act on the `torch.Tensor` level is a `dict[Any]`.
    If the output is a `dict[torch.Tensor]`, it is converted automatically to `DictTensor`.

    If a given method is available in both `dict` and `torch.Tensor`, then the one
    from `Tensor` is not available. Exceptions to the above rules are the `__dict__`
    and all comparison methods:

    - `__dict__`: Not available for this class.
    - `__lt__` or `<`: Supports two-way comparison between `DictTensor` and number, `torch.Tensor` or `DictTensor`.
    - `__le__` or `<=`: Supports two-way comparison between `DictTensor` and number, `torch.Tensor` or `DictTensor`.
    - `__eq__` or `==`: Supports two-way comparison between `DictTensor` and number, `torch.Tensor` or `DictTensor`.
    - `__ne__` or `!=`: Supports two-way comparison between `DictTensor` and number, `torch.Tensor` or `DictTensor`.
    - `__gt__` or `>`: Supports two-way comparison between `DictTensor` and number, `torch.Tensor` or `DictTensor`.
    - `__ge__` or `>=`: Supports two-way comparison between `DictTensor` and number, `torch.Tensor` or `DictTensor`.

    The only major function from `torch.Tensor` that doesn't work (to my knowledge) is the indexing.
    Indexing has to be done by manually looping the dictionary.

    Example:
        ```
        # Summing a float to a DictTensor
        dict_ten = DictTensor({
                "a": torch.zeros(5),
                "b": torch.zeros(2, 3),
                "c": torch.zeros(1, 2),
            })
        dict_ten + 2
        >>
        {'a': tensor([2., 2., 2., 2., 2.]),
        'b': tensor([[2., 2., 2.],
                [2., 2., 2.]]),
        'c': tensor([[2., 2.]])}

        # Summing a DictTensor to a DictTensor
        dict_ten = DictTensor({
                "a": torch.zeros(5),
                "b": torch.zeros(2, 3) + 0.1,
                "c": torch.zeros(1, 2) + 0.5,
            })
        dict_ten2 = DictTensor({
                "a": torch.zeros(5) + 0.2,
                "b": torch.zeros(2, 3) + 0.3,
                "c": torch.zeros(1, 2) + 0.4,
            })
        dict_ten + dict_ten2
        >>
        {'a': tensor([0.2000, 0.2000, 0.2000, 0.2000, 0.2000]),
        'b': tensor([[0.4000, 0.4000, 0.4000],
                [0.4000, 0.4000, 0.4000]]),
        'c': tensor([[0.9000, 0.9000]])}


        # Summing a accross the first axis
        dict_ten = DictTensor({
                "a": torch.zeros(5) + 0.1,
                "b": torch.zeros(2, 3) + 0.2,
                "c": torch.zeros(1, 2) + 0.5,
            })
        dict_ten.sum(axis=0)
        >>
        {'a': tensor(0.5000),
        'b': tensor([0.4000, 0.4000, 0.4000]),
        'c': tensor([0.5000, 0.5000])}

        # Getting the shape of each tensor
        dict_ten = DictTensor({
                "a": torch.zeros(5) + 0.1,
                "b": torch.zeros(2, 3) + 0.2,
                "c": torch.zeros(1, 2) + 0.5,
            })
        dict_ten.shape
        >>
        {'a': torch.Size([5]), 'b': torch.Size([2, 3]), 'c': torch.Size([1, 2])}
        ```
    """

    not_accepted_methods = {"__dict__"}
    from_tensor_methods = {
        "__lt__",
        "__le__",
        "__eq__",
        "__ne__",
        "__gt__",
        "__ge__",
    }

    def _dict_func_wrapper(self, func: Callable):
        """
        A function that wraps another function to be called on each value of the `self` dictionary,
        then returns a new dictionary.
        If the first arg is also a `DictTensor`, then also loop it's values. This will allow to sum
        or multiply different dict tensors.
        If the output is a dict of tensors, convert it to `DictTensor`.

        Parameters:
            func: The function to be wrapped.
        """

        def wrap(first_input: Union[DictTensor, Any], *args, **kwargs):
            # The following lines are for when a function is 'bounded' and
            # is called form `torch.FUNCTION()` instead of `Tensor.FUNCTION()`.
            this_func, this_args, this_kwargs = func, args, kwargs
            if hasattr(func, "__self__"):
                if len(args) == 3:
                    this_func, this_class, this_args = args
                elif len(args) == 4:
                    this_func, this_class, this_args, this_kwargs = args
                first_input, this_args = this_args[0], this_args[1:]

            arg0 = args[0] if len(args) > 0 else None
            if isinstance(first_input, DictTensor):
                # Loop all the elements in the dict to apply the function
                out = {}
                for k, v in first_input.items():
                    if isinstance(arg0, DictTensor):
                        # If the first arg is also a `DictTensor`, loop it as well to apply
                        # the function to each pair of values from the `DictTensor`
                        assert set(first_input.keys()) == set(
                            arg0.keys()
                        ), f"Keys do not match. \nkeys1={first_input.keys()}\nkeys2={arg0.keys()}"
                        out[k] = this_func(v, arg0[k], *this_args[1:], **this_kwargs)
                    else:
                        # In the regular case, simply apply the function to each value of the `DictTensor`
                        out[k] = this_func(v, *this_args, **this_kwargs)
            elif isinstance(arg0, DictTensor):
                out = {}
                for k, v in arg0.items():
                    if isinstance(arg0, DictTensor):
                        # If the first arg is also a `DictTensor`, loop it as well to apply
                        # the function to each pair of values from the `DictTensor`
                        out[k] = this_func(arg0[k], v, *this_args[1:], **this_kwargs)
            else:
                raise TypeError("type `DictTensor` not found")

            out = self._to_dict_tensor(out, raise_if_type_error=False)
            return out

        return wrap

    def _create_property(self, prop_name: str):
        """
        A function that creates a property from a name.
        When the property is called, it will be called on each value of the `self` dictionary,
        then returns a new dictionary.
        """
        setattr(
            self.__class__,
            prop_name,
            property(fget=lambda self: {k: getattr(v, prop_name) for k, v in self.items()}),
        )

    def __init__(self, dic: Dict[Hashable, Tensor]):
        """
        Take a dictionary of `torch.Tensors`, and transform it into a `DictTensor`.
        Register all the required methods from `torch.Tensors`, but modify them
        to work on dictionary instead.
        """
        super().__init__(dic)

        # Assert that the dictionary is a dict of Tensor
        assert isinstance(dic, dict), f"Must be a dict, got {type(dic)}"
        assert all(
            [isinstance(val, Tensor) for val in dic.values()]
        ), f"Must only contain `torch.Tensor`, found {[type(v) for v in dic.values()]}"

        # From `torch.Tensor`, find the functions/methods/attributes to register
        tensor_func_names = {f for f in dir(Tensor)}  # if not ((f.endswith("_")) or (f.startswith("_")))}
        dict_func_names = {f for f in dir(dict)}  # if not ((f.endswith("_")) or (f.startswith("_")))}
        func_names_to_register = (
            tensor_func_names - dict_func_names - self.not_accepted_methods
        ) | self.from_tensor_methods

        # Loop all selected functions/methods/attributes to register them in the current class
        for func_name in func_names_to_register:
            func = getattr(Tensor, func_name)
            if isinstance(func, Callable):
                # Register the methods and functions
                setattr(DictTensor, func_name, self._dict_func_wrapper(func))
            else:
                # Register the attributes as properties
                self._create_property(func_name)

    @staticmethod
    def _to_dict_tensor(dic: Dict[Hashable, Tensor], raise_if_type_error: bool = True) -> DictTensor:
        """
        Convert a dictionary of tensors into a `DictTensor`
        """
        if all([isinstance(v, Tensor) for v in dic.values()]):
            dic = DictTensor(dic)
        else:
            if raise_if_type_error:
                raise TypeError("All values must be Tensors")
        return dic

    def apply(self, func: Callable, *args, **kwargs) -> Union[DictTensor, Dict[Any]]:
        """
        Apply a function on every Tensor "value" of the current dictionary

        Parameters:
            func: function to be called on each Tensor
            *args, **kwargs: Additional parameters to the function. Note that
                the Tensor should be the first input to the function.

        Returns:
            `DictTensor` if the output of the function is a `torch.Tensor`,
            or `dict[X]` if the output of the function is of type`X`.
        """
        out = {k: func(v, *args, **kwargs) for k, v in self.items()}
        out = self._to_dict_tensor(out, raise_if_type_error=False)
        return out


# graphium/utils/fs.py

from typing import Union
from typing import Optional

import os
import io
import platformdirs
import pathlib

from tqdm.auto import tqdm
import fsspec


def get_cache_dir(suffix: str = None, create: bool = True) -> pathlib.Path:
    """Get a local cache directory. You can append a suffix folder to it and optionnaly create
    the folder if it doesn't exist.
    """

    cache_dir = pathlib.Path(platformdirs.user_cache_dir(appname="graphium"))

    if suffix is not None:
        cache_dir /= suffix

    if create:
        cache_dir.mkdir(exist_ok=True, parents=True)

    return cache_dir


def get_mapper(path: Union[str, os.PathLike]):
    """Get the fsspec mapper.
    Args:
        path: a path supported by `fsspec` such as local, s3, gcs, etc.
    """
    return fsspec.get_mapper(str(path))


def get_basename(path: Union[str, os.PathLike]):
    """Get the basename of a file or a folder.
    Args:
        path: a path supported by `fsspec` such as local, s3, gcs, etc.
    """
    path = str(path)
    mapper = get_mapper(path)
    clean_path = path.rstrip(mapper.fs.sep)
    return str(clean_path).split(mapper.fs.sep)[-1]


def get_extension(path: Union[str, os.PathLike]):
    """Get the extension of a file.
    Args:
        path: a path supported by `fsspec` such as local, s3, gcs, etc.
    """
    basename = get_basename(path)
    return basename.split(".")[-1]


def exists(path: Union[str, os.PathLike, fsspec.core.OpenFile, io.IOBase]):
    """Check whether a file exists.
    Args:
        path: a path supported by `fsspec` such as local, s3, gcs, etc.
    """

    if isinstance(path, fsspec.core.OpenFile):
        return path.fs.exists(path.path)

    elif isinstance(path, (str, pathlib.Path)):
        mapper = get_mapper(str(path))
        return mapper.fs.exists(path)

    else:
        # NOTE(hadim): file-like objects always exist right?
        return True


def exists_and_not_empty(path: Union[str, os.PathLike]):
    """Check whether a directory exists and is not empty."""

    if not exists(path):
        return False

    fs = get_mapper(path).fs

    return len(fs.ls(path)) > 0


def mkdir(path: Union[str, os.PathLike], exist_ok: bool = True):
    """Create directory including potential parents."""
    fs = get_mapper(path).fs
    fs.mkdirs(path, exist_ok=exist_ok)


def rm(path: Union[str, os.PathLike], recursive=False, maxdepth=None):
    """Delete a file or a directory with all nested files."""
    fs = get_mapper(path).fs
    fs.rm(path, recursive=recursive, maxdepth=maxdepth)


def join(*paths):
    """Join paths together. The first element determine the
    filesystem to use (and so the separator.
    Args:
        paths: a list of paths supported by `fsspec` such as local, s3, gcs, etc.
    """
    paths = [str(path) for path in paths]
    source_path = paths[0]
    fs = get_mapper(source_path).fs
    full_path = fs.sep.join(paths)
    return full_path


def get_size(file: Union[str, os.PathLike, io.IOBase, fsspec.core.OpenFile]) -> Optional[int]:
    """Get the size of a file given its path. Return None if the
    size can't be retrieved.
    """

    if isinstance(file, io.IOBase) and hasattr(file, "name"):
        fs_local = fsspec.filesystem("file")
        file_size = fs_local.size(getattr(file, "name"))

    elif isinstance(file, (str, pathlib.Path)):
        fs = get_mapper(str(file)).fs
        file_size = fs.size(str(file))

    elif isinstance(file, fsspec.core.OpenFile):
        file_size = file.fs.size(file.path)

    else:
        file_size = None

    return file_size


def copy(
    source: Union[str, os.PathLike, io.IOBase, fsspec.core.OpenFile],
    destination: Union[str, os.PathLike, io.IOBase, fsspec.core.OpenFile],
    chunk_size: int = None,
    force: bool = False,
    progress: bool = False,
    leave_progress: bool = True,
):
    """Copy one file to another location across different filesystem (local, S3, GCS, etc).

    Args:
        source: path or file-like object to copy from.
        destination: path or file-like object to copy to.
        chunk_size: the chunk size to use. If progress is enabled the chunk
            size is `None`, it is set to 2048.
        force: whether to overwrite the destination file it it exists.
        progress: whether to display a progress bar.
        leave_progress: whether to hide the progress bar once the copy is done.
    """

    if progress and chunk_size is None:
        chunk_size = 2048

    if isinstance(source, (str, os.PathLike)):
        source_file = fsspec.open(str(source), "rb")
    else:
        source_file = source

    if isinstance(destination, (str, os.PathLike)):
        # adapt the file mode of the destination depending on the source file.
        destination_mode = "wb"
        if hasattr(source_file, "mode"):
            destination_mode = "wb" if "b" in getattr(source_file, "mode") else "w"
        elif isinstance(source_file, io.BytesIO):
            destination_mode = "wb"
        elif isinstance(source_file, io.StringIO):
            destination_mode = "w"

        destination_file = fsspec.open(str(destination), destination_mode)
    else:
        destination_file = destination

    if not exists(source_file):
        raise ValueError(f"The file being copied does not exist: {source}")

    if not force and exists(destination_file):
        raise ValueError(f"The destination file to copy already exists: {destination}")

    with source_file as source_stream:
        with destination_file as destination_stream:
            if chunk_size is None:
                # copy without chunks
                destination_stream.write(source_stream.read())

            else:
                # copy with chunks

                # determine the size of the source file
                source_size = None
                if progress:
                    source_size = get_size(source)

                # init progress bar
                pbar = tqdm(
                    total=source_size,
                    leave=leave_progress,
                    disable=not progress,
                    unit="B",
                    unit_divisor=1024,
                    unit_scale=True,
                )

                # start the loop
                while True:
                    data = source_stream.read(chunk_size)
                    if not data:
                        break
                    destination_stream.write(data)
                    pbar.update(chunk_size)

                pbar.close()


# graphium/utils/hashing.py

from typing import Any
import hashlib
import yaml


def get_md5_hash(object: Any) -> str:
    """
    MD5 hash of any object.
    The object is converted to a YAML string before being hashed.
    This allows for nested dictionaries/lists and for hashing of classes and their attributes.
    """
    dhash = hashlib.md5()
    encoded = yaml.dump(object, sort_keys=True).encode()
    dhash.update(encoded)
    return dhash.hexdigest()


# graphium/utils/moving_average_tracker.py

from dataclasses import dataclass


@dataclass
class MovingAverageTracker:
    num_samples: int = 0
    mean_value: float = 0.0

    def update(self, value: float):
        self.mean_value = self.mean_value * (self.num_samples / (self.num_samples + 1)) + value / (
            self.num_samples + 1
        )
        self.num_samples += 1

    def reset(self):
        # no need to update mean_value, it will be reset when update is called
        self.num_samples = 0


# graphium/utils/mup.py

##### Code adapted from the `mup` package from Microsoft https://github.com/microsoft/mup

from torch.nn import Linear
from torch.nn.modules.conv import _ConvNd
from mup import get_shapes, assert_hidden_size_inf, MuReadout, rescale_linear_bias, save_base_shapes
from mup.shape import _zip_infshape_dict, _extract_shapes

from graphium.nn.base_layers import MuReadoutGraphium


def apply_infshapes(model, infshapes):
    """
    Modified from the regular `mup.apply_infshapes` by explicitly adding `base_dim` to the `MuReadoutGraphium`.
    This allows the code to work on IPUs.
    """
    for name, p in model.named_parameters():
        p.infshape = infshapes[name]
    for _, module in model.named_modules():
        if isinstance(module, MuReadoutGraphium):
            module.base_width = module.weight.infshape[-1].base_dim


def set_base_shapes(model, base, rescale_params=True, delta=None, savefile=None, do_assert=True):
    """Sets the `p.infshape` attribute for each parameter `p` of `model`.

    Code taken from the `mup` package from Microsoft https://github.com/microsoft/mup.
    No change except in the `apply_inf_shapes`, using the one from Graphium instead of `mup`

    Inputs:
        model: nn.Module instance
        base: The base model.
            Can be nn.Module, a dict of shapes, a str, or None.
            If None, then defaults to `model`
            If str, then treated as filename for yaml encoding of a dict of base shapes.
        rescale_params:
            assuming the model is initialized using the default pytorch init (or
            He initialization etc that scale the same way with fanin): If True
            (default), rescales parameters to have the correct (μP) variances.
        do_assert:
    Output:
        same object as `model`, after setting the `infshape` attribute of each parameter.
    """
    if base is None:
        base = model
    base_shapes = _extract_shapes(base)
    if delta is not None:
        delta_shapes = _extract_shapes(delta)
        base_shapes = _zip_infshape_dict(base_shapes, delta_shapes)
    shapes = get_shapes(model)
    infshapes = _zip_infshape_dict(base_shapes, shapes)
    if savefile is not None:
        save_base_shapes(infshapes, savefile)
    apply_infshapes(model, infshapes)
    if do_assert:
        assert_hidden_size_inf(model)
    if rescale_params:
        for name, module in model.named_modules():
            if isinstance(module, MuReadout):
                module._rescale_parameters()
            elif isinstance(module, (Linear, _ConvNd)):
                rescale_linear_bias(module)
    return model


# graphium/utils/packing.py

# Copyright (c) 2022 Graphcore Ltd. All rights reserved.

from typing import List, Tuple, Iterable, Optional
import numpy as np
import torch


class MolPack:
    """
    Class that keeps track of the number of atoms and indices that are added
    to each pack. Useful when doing packing, or other forms of smart batching.
    A pack is a batch, but with optimized memory consumption.
    """

    def __init__(self):
        self.num_nodes = 0
        self.num_graphs = 0
        self.average_atom = 0
        self.indices = []

    def add_mol(self, num_nodes: int, idx: int) -> "MolPack":
        """
        Add a molecule and it's index to the batch

        Parameters:
            num_nodes: Number of atoms of the new molecule

            idx: Index associated to the molecule
        """
        self.num_nodes += num_nodes
        self.num_graphs += 1
        self.average_atom = self.num_nodes / self.num_graphs
        self.indices.append(idx)
        return self

    def expected_atoms(self, remaining_mean_num_nodes: float, batch_size: int) -> float:
        """
        Given a desired batch size, and given the remaining mean number of
        atoms, find the expected number of atoms of the current batch when it is full

        Parameters:
            remaining_mean_num_nodes: Average number of atoms per molecule
                left to be sampled and distributed across tasks.

            batch_size: Desired batch size

        Returns:
            expected_atoms: The expected number of atoms in this batch if we
                sample randomly the remaining molecules.
        """
        return self.num_nodes + ((batch_size - self.num_graphs) * remaining_mean_num_nodes)

    def __repr__(self) -> str:
        """
        Print the main attributes of the current class
        """
        return f"{self.__class__.__name__}(m: {self.num_graphs},\ta: {self.num_nodes},\tav: {self.average_atom:.1f})"


def smart_packing(num_nodes: List[int], batch_size: int) -> List[List[int]]:
    """
    Simple and fast algorithm for packing graphs such that each batch has roughly the
    same number of atoms.
    Has for-loop scalability issues `O(num_graphs * ipu_batch_size)` = `O(num_graphs^2 / batch_size)`

    Parameters:
        num_nodes: List of the number of atoms per molecule for the entire global batch.
            Must be of length `batch_size * ipu_batch_size`.

        batch_size: The batch size per iteration, considering a single device and single
            forward pass.
            The global batch size is `batch_size * device_iterations * replication_factor * gradient_accumulation`

    Returns:
        packed_indices: A list of packs, each containing a list of indices, such that
            if we collect `num_nodes` from the indices, then each pack has roughly the
            same total number of atoms.
    """

    # Sort the list
    num_nodes = np.asarray(num_nodes)
    argsort_num_nodes = np.argsort(num_nodes)
    sorted_num_nodes = num_nodes[argsort_num_nodes]
    ipu_batch_size = int(len(num_nodes) / batch_size)
    sorted_num_nodes, initial_num_nodes = (
        sorted_num_nodes[:-ipu_batch_size],
        sorted_num_nodes[-ipu_batch_size:],
    )
    reverse_cumsum = np.sum(sorted_num_nodes) - np.cumsum(sorted_num_nodes) + sorted_num_nodes[-1]

    # Start with the largest element in separate packs
    mol_batches = [
        MolPack().add_mol(initial_num_nodes[-ii - 1], argsort_num_nodes[-ii - 1])
        for ii in range(ipu_batch_size)
    ]

    # Loop from smallest to largest molecule, and add each molecule to the pack with smallest expected sum
    for ii, num_atom in enumerate(sorted_num_nodes):
        remaining_mean = reverse_cumsum[ii] / (len(sorted_num_nodes) - ii)
        max_expected, idx_max_expected = 0, 0
        for jj, m in enumerate(mol_batches):
            if m.num_graphs >= batch_size:
                continue
            expected = m.num_nodes + (
                (batch_size - m.num_graphs) * remaining_mean
            )  # Faster than calling m.expected_atoms
            if expected > max_expected:
                max_expected = expected
                idx_max_expected = jj
        mol_batches[idx_max_expected].add_mol(num_atom, argsort_num_nodes[ii])

    packed_indices = [batch.indices for batch in mol_batches]

    return packed_indices


def fast_packing(num_nodes: List[int], batch_size: int) -> List[List[int]]:
    """
    Super fast algorithm for packing graphs such that each batch has roughly the
    same number of atoms. Not as good as `smart_packing` but
    faster and more scalable for-loop complexity of `O(batch_size)`.

    Parameters:
        num_nodes: List of the number of atoms per molecule for the entire global batch.
            Must be of length `batch_size * ipu_batch_size`.

        batch_size: The batch size per iteration, considering a single device and single
            forward pass.
            The global batch size is `batch_size * device_iterations * replication_factor * gradient_accumulation`

    Returns:
        packed_indices: A list of packs, each containing a list of indices, such that
            if we collect `num_nodes` from the indices, then each pack has roughly the
            same total number of atoms.
    """
    num_nodes = np.asarray(num_nodes)
    argsort_num_nodes = np.argsort(num_nodes)
    ipu_batch_size = int(len(num_nodes) / batch_size)

    packed_indices = np.stack(
        [
            np.random.permutation(argsort_num_nodes[ii * ipu_batch_size : (ii + 1) * ipu_batch_size])
            for ii in range(batch_size)
        ],
        axis=0,
    ).T.tolist()
    return packed_indices


def hybrid_packing(num_nodes: List[int], batch_size: int) -> List[List[int]]:
    """
    Uses a combination of the `smart_packing` `O(n^2)` on the most important data points,
    and the `fast_packing` `O(n)` on the average-sized data points.

    Depending on the expected complexity

    Parameters:
        num_nodes: List of the number of atoms per molecule for the entire global batch.
            Must be of length `batch_size * ipu_batch_size`.

        batch_size: The batch size per iteration, considering a single device and single
            forward pass.
            The global batch size is `batch_size * device_iterations * replication_factor * gradient_accumulation`

    Returns:
        packed_indices: A list of packs, each containing a list of indices, such that
            if we collect `num_nodes` from the indices, then each pack has roughly the
            same total number of atoms.
    """

    # Determine the parameters based on the complexity of the smart-packing.
    # The bigger the complexity, the more the `fast_packing` algorithm becomes
    # statistically powerful, and the more speed benefits it provides.
    smart_packing_complexity = len(num_nodes) ** 2 / batch_size
    if smart_packing_complexity < 1e4:
        return smart_packing(num_nodes=num_nodes, batch_size=batch_size)
    elif smart_packing_complexity < 1e5:
        big, small = 3, 6
    else:
        return fast_packing(num_nodes=num_nodes, batch_size=batch_size)

    # Small datasets benefit from smart-packing, without compute burden
    ipu_batch_size = int(len(num_nodes) / batch_size)
    if len(num_nodes) < (big + small) * ipu_batch_size:
        return smart_packing(num_nodes=num_nodes, batch_size=batch_size)

    # Sort the list
    num_nodes = np.asarray(num_nodes)
    argsort_num_nodes = np.argsort(num_nodes)

    # Smallest and biggest graphs are often outliers and will benefit from the `smart_packing`
    biggest_graphs = argsort_num_nodes[-big * ipu_batch_size :]
    smallest_graphs = argsort_num_nodes[: small * ipu_batch_size]
    big_n_small_graphs = np.concatenate([biggest_graphs, smallest_graphs])
    big_n_small_packs = smart_packing(num_nodes[big_n_small_graphs], batch_size=big + small)
    big_n_small_indices = [big_n_small_graphs[pack] for pack in big_n_small_packs]
    big_n_small_nodes = [num_nodes[pack] for pack in big_n_small_indices]

    # Medium graphs will be packed faster
    medium_graphs = argsort_num_nodes[small * ipu_batch_size : -big * ipu_batch_size]
    medium_packs = fast_packing(num_nodes[medium_graphs], batch_size=batch_size - big - small)
    medium_indices = [medium_graphs[pack] for pack in medium_packs]
    medium_nodes = [num_nodes[pack] for pack in medium_indices]

    # Pack the big/small with the medium in a smart way
    big_n_small_sort = np.argsort(np.sum(np.stack(big_n_small_nodes, axis=1), axis=0))
    medium_sort = np.argsort(np.sum(np.stack(medium_nodes, axis=1), axis=0))
    packed_indices = [
        np.concatenate([medium_indices[medium_sort[ii]], big_n_small_indices[big_n_small_sort[-ii]]])
        for ii in range(len(medium_indices))
    ]

    return packed_indices


def get_pack_sizes(packed_indices, num_nodes):
    """
    Get the number of atoms of each pack
    """
    pack_sums = []
    for pack in packed_indices:
        pack_sum = 0
        for idx in pack:
            pack_sum += num_nodes[idx]
        pack_sums.append(pack_sum)
    return pack_sums


def estimate_max_pack_node_size(num_nodes: Iterable[int], batch_size: int, combined_batch_size: int):
    """
    Estimate the value of `max_num_nodes`, which represents the maximum number of nodes
    needed in a batch to fit the data.

    Parameters:
        num_nodes: Number of nodes for all the graphs in the dataset
        batch_size: The regular batch size per IPU
        combined_batch_size: batch_size * device_iterations
                             * replication_factor * gradient_accumulation

    """

    # Estimate the packing size needed
    rand_indices = np.arange(len(num_nodes))
    np.random.shuffle(rand_indices)
    max_pack_size = 0
    for ii in range(0, len(num_nodes), combined_batch_size):
        this_indices = rand_indices[ii : ii + combined_batch_size]
        choice = num_nodes[this_indices]
        if len(choice) == combined_batch_size:
            packed_indices = hybrid_packing(choice, batch_size)
            max_pack_size = max(max_pack_size, max(get_pack_sizes(packed_indices, num_nodes[this_indices])))
    max_pack_size_per_graph = max_pack_size / batch_size

    return max_pack_size, max_pack_size_per_graph


def node_to_pack_indices_mask(
    packed_indices: Iterable[Iterable[int]], all_num_nodes: Iterable[int], max_pack_size: Optional[int] = None
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Given a list of packed indices, and the number of nodes in each graph,
    return a tensor of shape (sum(all_num_nodes), 2) where the first column
    is the pack index, and the second column is the node index within the pack.

    Can be used to generate a dense packing of the nodes as follows:
    ```
    # node_features: A tensor of shape (num_nodes, num_node_features)
    # num_packs: The number of packs desired
    # max_nodes_per_pack: The maximum number of nodes per pack
    # dense_pack: A tensor of shape (num_packs, max_nodes_per_pack, num_node_features)

    dense_pack = torch.zeros([num_packs, max_nodes_per_pack, num_node_features])
    dense_pack[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]] = node_features
    ```

    This is useful when using a Transformer, to avoid wasteful padding when the
    the longest sequence is much longer than the average sequence length.

    Parameters:
        packed_indices: A list of lists of graph indices, where each sub-list
                        represents a pack of graphs
        all_num_nodes: The number of nodes in each graph
        max_pack_size: The maximum number of nodes per pack. If None, will be
                          infered from the provided packs.
                          Useful to determine the shape of the `pack_attn_mask`.

    Returns:
        pack_from_node_idx: A tensor of shape (num_nodes, 2) where the first column
                        is the pack index, and the second column is the node index within the pack.

        pack_attn_mask: A tensor of shape (num_packs, max_pack_size, max_pack_size),
                            that represents the attention masking for each pack,
                            such that the graphs in the pack are masked out from each other.
    """

    all_num_nodes = torch.as_tensor(all_num_nodes, dtype=torch.long)
    cumsum_num_nodes = torch.cumsum(all_num_nodes, dim=0)
    if max_pack_size is None:
        pack_sizes = get_pack_sizes(packed_indices, all_num_nodes)
        max_pack_size = max(pack_sizes)

    # Get the node indices associated to the packs, with 0 padding
    pack_from_node_idx = torch.zeros(sum(all_num_nodes), 2, dtype=torch.long)
    pack_attn_mask = []  # masks for the attention
    for ii, pack in enumerate(packed_indices):
        jj = 0  # Counter for the number of nodes in the pack
        this_pack_attn_mask = torch.ones((max_pack_size, max_pack_size), dtype=torch.bool)
        for graph_idx in pack:
            num_nodes = all_num_nodes[graph_idx]
            node_idx = torch.arange(cumsum_num_nodes[graph_idx] - num_nodes, cumsum_num_nodes[graph_idx])
            this_pack_attn_mask[jj : jj + num_nodes, jj : jj + num_nodes] = False
            pack_from_node_idx[node_idx, 0] = ii
            pack_from_node_idx[node_idx, 1] = jj + torch.arange(num_nodes)
            jj += num_nodes
        pack_attn_mask.append(this_pack_attn_mask)
    pack_attn_mask = torch.stack(pack_attn_mask, dim=0)

    return pack_from_node_idx, pack_attn_mask


# graphium/utils/read_file.py

""" Utiles for data parsing"""
import os
import warnings
import numpy as np
import pandas as pd
import datamol as dm
from functools import partial
from copy import copy
import fsspec

from loguru import logger
from rdkit import Chem
from rdkit.Chem.Descriptors import ExactMolWt

from graphium.utils.tensor import parse_valid_args, arg_in_func


def read_file(filepath, as_ext=None, **kwargs):
    r"""
    Allow to read different file format and parse them into a MolecularDataFrame.
    Supported formats are:
    * csv (.csv, .smile, .smiles, .tsv)
    * txt (.txt)
    * xls (.xls, .xlsx, .xlsm, .xls*)
    * sdf (.sdf)
    * pkl (.pkl)

    Arguments
    -----------

        filepath: str
            The full path and name of the file to read.
            It also supports the s3 url path.
        as_ext: str, Optional
            The file extension used to read the file. If None, the extension is deduced
            from the extension of the file. Otherwise, no matter the file extension,
            the file will be read according to the specified ``as_ext``.
            (Default=None)
        **kwargs: All the optional parameters required for the desired file reader.

    TODO: unit test to make sure it works well with all extensions

    Returns
    ---------
        df: pandas.DataFrame
            The ``pandas.DataFrame`` containing the parsed data

    """

    # Get the file extension
    if as_ext is None:
        file_ext = os.path.splitext(filepath)[-1].lower()[1:]
    else:
        file_ext = as_ext
        if not isinstance(file_ext, str):
            raise "`file_type` must be a `str`. Provided: {}".format(file_ext)

    open_mode = "r"

    # Read the file according to the right extension
    if file_ext in ["csv", "smile", "smiles", "smi", "tsv"]:
        file_reader = pd.read_csv
    elif file_ext == "txt":
        file_reader = pd.read_table
    elif file_ext[0:3] == "xls":
        open_mode = "rb"
        file_reader = partial(pd.read_excel, engine="openpyxl")
    elif file_ext == "sdf":
        file_reader = parse_sdf_to_dataframe
    elif file_ext == "pkl":
        open_mode = "rb"
        file_reader = pd.read_pickle
    else:
        raise 'File extension "{}" not supported'.format(file_ext)

    kwargs = parse_valid_args(fn=file_reader, param_dict=kwargs)

    if file_ext[0:3] not in ["sdf", "xls"]:
        with file_opener(filepath, open_mode) as file_in:
            data = file_reader(file_in, **kwargs)
    else:
        data = file_reader(filepath, **kwargs)
    return data


def parse_sdf_to_dataframe(sdf_path, as_cxsmiles=True, skiprows=None):
    r"""
    Allows to read an SDF file containing molecular informations, convert
    it to a pandas DataFrame and convert the molecules to SMILES. It also
    lists a warning of all the molecules that couldn't be read.

    Arguments
    -----------

        sdf_path: str
            The full path and name of the sdf file to read
        as_cxsmiles: bool, optional
            Whether to use the CXSMILES notation, which preserves atomic coordinates,
            stereocenters, and much more.
            See `https://dl.chemaxon.com/marvin-archive/latest/help/formats/cxsmiles-doc.html`
            (Default = True)
        skiprows: int, list
            The rows to skip from dataset. The enumerate index starts from 1 insted of 0.
            (Default = None)

    """

    # read the SDF file
    # locally or from s3
    data = dm.read_sdf(sdf_path)

    # For each molecule in the SDF file, read all the properties and add it to a list of dict.
    # Also count the number of molecules that cannot be read.
    data_list = []
    count_none = 0
    if skiprows is not None:
        if isinstance(skiprows, int):
            skiprows = range(0, skiprows - 1)
        skiprows = np.array(skiprows) - 1

    for idx, mol in enumerate(data):
        if (skiprows is not None) and (idx in skiprows):
            continue

        if (mol is not None) and (ExactMolWt(mol) > 0):
            mol_dict = mol.GetPropsAsDict()
            data_list.append(mol_dict)
            if as_cxsmiles:
                smiles = Chem.rdmolfiles.MolToCXSmiles(mol, canonical=True)
            else:
                smiles = dm.to_smiles(mol, canonical=True)
            data_list[-1]["SMILES"] = smiles
        else:
            count_none += 1
            logger.info(f"Could not read molecule # {idx}")

    # Display a message or warning after the SDF is done parsing
    if count_none == 0:
        logger.info("Successfully read the SDF file without error: {}".format(sdf_path))
    else:
        warnings.warn(
            (
                'Error reading {} molecules from the "{}" file.\
         {} molecules read successfully.'
            ).format(count_none, sdf_path, len(data_list))
        )
    return pd.DataFrame(data_list)


def file_opener(filename, mode="r"):
    """File reader stream"""
    filename = str(filename)
    if "w" in mode:
        filename = "simplecache::" + filename
    if filename.endswith(".gz"):
        instream = fsspec.open(filename, mode=mode, compression="gzip")
    else:
        instream = fsspec.open(filename, mode=mode)
    return instream


# graphium/utils/safe_run.py

from loguru import logger
import traceback as tb


class SafeRun:
    def __init__(self, name: str, raise_error: bool = True, verbose: int = 2) -> None:
        """
        Run some code with error handling and some printing, using the with statment.

        Example:
            In the example below, the `2+None`, an error will be caught and printed.
            ```
            with SafeRun(name="Addition that fails", raise_error=False):
                2 + None
            ```

        Parameters:
            name: Name of the code, used for printing
            raise_error: Whether to raise an error, or to catch it and print instead
            verbose: The level of verbosity
                0: Do not print anything
                1: Print only the traced stack when an error is caught and `raise_error` is False
                2: Print headers and footers at the start and exit of the with statement.
        """
        self.name = name
        self.raise_error = raise_error
        self.verbose = verbose

    def __enter__(self):
        """
        Print that the with-statement started, if `self.verbose >= 2`
        """
        if self.verbose >= 2:
            logger.info(f"\n------------ {self.name} STARTED ------------")

    def __exit__(self, type, value, traceback):
        """
        Handle the error. Raise it if `self.raise_error==True`, otherwise ignore it
        and print it if `self.verbose >= 1`. Also print that the with-statement is
        completed if `self.verbose >= 2`.
        """
        if traceback is not None:
            if self.raise_error:
                if self.verbose >= 1:
                    logger.error(f"------------ {self.name} ERROR: ------------")
                return False
            else:
                if self.verbose >= 1:
                    logger.error(f"------------ {self.name} ERROR: ------------\nERROR skipped. Traceback:\n")
                    logger.trace(print("".join(tb.format_exception(None, value, traceback))))
                return True
        else:
            if self.verbose >= 2:
                logger.info("\n------------ {self.name} COMPLETED ------------\n\n")
            return True


# graphium/utils/spaces.py

from copy import deepcopy
import torch
import torch.optim.lr_scheduler as sc
import torchmetrics.functional as TorchMetrics

import graphium.nn.base_layers as BaseLayers
import graphium.utils.custom_lr as CustomLR
import graphium.data.datamodule as Datamodules
import graphium.ipu.ipu_losses as IPULosses
import graphium.ipu.ipu_metrics as Metrics
import graphium.nn.pyg_layers as PygLayers
import graphium.nn.residual_connections as Residuals
import graphium.nn.encoders as Encoders
import graphium.trainer.losses as Losses

PE_ENCODERS_DICT = {
    "laplacian_pe": Encoders.laplace_pos_encoder.LapPENodeEncoder,
    "mlp": Encoders.mlp_encoder.MLPEncoder,
    "signnet": Encoders.signnet_pos_encoder.SignNetNodeEncoder,
    "gaussian_kernel": Encoders.gaussian_kernel_pos_encoder.GaussianKernelPosEncoder,
    "bessel_kernel": Encoders.bessel_pos_encoder.BesselSphericalPosEncoder,
}


FC_LAYERS_DICT = {
    "fc": BaseLayers.FCLayer,
}

PYG_LAYERS_DICT = {
    "pyg:gcn": PygLayers.GCNConvPyg,
    "pyg:gin": PygLayers.GINConvPyg,
    "pyg:gine": PygLayers.GINEConvPyg,
    "pyg:gated-gcn": PygLayers.GatedGCNPyg,
    "pyg:pna-msgpass": PygLayers.PNAMessagePassingPyg,
    "pyg:gps": PygLayers.GPSLayerPyg,
    "pyg:dimenet": PygLayers.DimeNetPyg,
}

LAYERS_DICT = deepcopy(FC_LAYERS_DICT)
LAYERS_DICT.update(deepcopy(PYG_LAYERS_DICT))


RESIDUALS_DICT = {
    "none": Residuals.ResidualConnectionNone,
    "simple": Residuals.ResidualConnectionSimple,
    "weighted": Residuals.ResidualConnectionWeighted,
    "concat": Residuals.ResidualConnectionConcat,
    "densenet": Residuals.ResidualConnectionDenseNet,
    "random": Residuals.ResidualConnectionRandom,
}

LOSS_DICT = {
    "mse": torch.nn.MSELoss,
    "bce": torch.nn.BCELoss,
    "l1": torch.nn.L1Loss,
    "mae": torch.nn.L1Loss,
    "hybrid_ce": Losses.HybridCELoss,
    "bce_ipu": IPULosses.BCELossIPU,
    "bce_logits_ipu": IPULosses.BCEWithLogitsLossIPU,
    "mse_ipu": IPULosses.MSELossIPU,
    "mae_ipu": IPULosses.L1LossIPU,
    "l1_ipu": IPULosses.L1LossIPU,
    "hybrid_ce_ipu": IPULosses.HybridCELossIPU,
}


SCHEDULER_DICT = {
    "CosineAnnealingLR": sc.CosineAnnealingLR,
    "CosineAnnealingWarmRestarts": sc.CosineAnnealingWarmRestarts,
    "CyclicLR": sc.CyclicLR,
    "ExponentialLR": sc.ExponentialLR,
    "LambdaLR": sc.LambdaLR,
    "MultiStepLR": sc.MultiStepLR,
    "ReduceLROnPlateau": sc.ReduceLROnPlateau,
    "StepLR": sc.StepLR,
    "ConstantLR": sc.ConstantLR,
    "WarmUpLinearLR": CustomLR.WarmUpLinearLR,
}

METRICS_CLASSIFICATION = {
    "accuracy": TorchMetrics.accuracy,
    "averageprecision": TorchMetrics.average_precision,
    "auroc": TorchMetrics.auroc,
    "confusionmatrix": TorchMetrics.confusion_matrix,
    "f1": TorchMetrics.f1_score,
    "fbeta": TorchMetrics.fbeta_score,
    "precisionrecallcurve": TorchMetrics.precision_recall_curve,
    "precision": TorchMetrics.precision,
    "recall": TorchMetrics.recall,
    "mcc": TorchMetrics.matthews_corrcoef,
    "auroc_ipu": Metrics.auroc_ipu,
    "accuracy_ipu": Metrics.accuracy_ipu,
    "average_precision_ipu": Metrics.average_precision_ipu,
    "f1_ipu": Metrics.f1_score_ipu,
    "fbeta_ipu": Metrics.fbeta_score_ipu,
    "precision_ipu": Metrics.precision_ipu,
    "recall_ipu": Metrics.recall_ipu,
}

METRICS_REGRESSION = {
    "mae": TorchMetrics.mean_absolute_error,
    "mape": TorchMetrics.mean_absolute_percentage_error,
    "mse": TorchMetrics.mean_squared_error,
    "msle": TorchMetrics.mean_squared_log_error,
    "pearsonr": TorchMetrics.pearson_corrcoef,
    "spearmanr": TorchMetrics.spearman_corrcoef,
    "r2": TorchMetrics.r2_score,
    "cosine": TorchMetrics.cosine_similarity,
    "pearsonr_ipu": Metrics.pearson_ipu,
    "spearmanr_ipu": Metrics.spearman_ipu,
    "r2_score_ipu": Metrics.r2_score_ipu,
    "mae_ipu": Metrics.mean_absolute_error_ipu,
    "mse_ipu": Metrics.mean_squared_error_ipu,
}

METRICS_DICT = deepcopy(METRICS_CLASSIFICATION)
METRICS_DICT.update(METRICS_REGRESSION)


DATAMODULE_DICT = {
    "GraphOGBDataModule": Datamodules.GraphOGBDataModule,
    "MultitaskFromSmilesDataModule": Datamodules.MultitaskFromSmilesDataModule,
    "ADMETBenchmarkDataModule": Datamodules.ADMETBenchmarkDataModule,
    "FakeDataModule": Datamodules.FakeDataModule,
}


# graphium/utils/tensor.py

import os
import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from typing import Iterable, List, Union, Any, Callable
from inspect import getfullargspec
from copy import copy, deepcopy
from loguru import logger

from rdkit.Chem import AllChem

import torch
from torch import Tensor


def save_im(im_dir, im_name: str, ext: List[str] = ["svg", "png"], dpi: int = 600) -> None:
    if not os.path.exists(im_dir):
        if im_dir[-1] not in ["/", "\\"]:
            im_dir += "/"
        os.makedirs(im_dir)

    if isinstance(ext, str):
        ext = [ext]

    full_name = os.path.join(im_dir, im_name)
    for this_ext in ext:
        plt.savefig(f"{full_name}.{this_ext}", dpi=dpi, bbox_inches="tight", pad_inches=0)


def is_dtype_torch_tensor(dtype: Union[np.dtype, torch.dtype]) -> bool:
    r"""
    Verify if the dtype is a torch dtype

    Parameters:
        dtype: dtype
            The dtype of a value. E.g. np.int32, str, torch.float

    Returns:
        A boolean saying if the dtype is a torch dtype
    """
    return isinstance(dtype, torch.dtype) or (dtype == Tensor)


def is_dtype_numpy_array(dtype: Union[np.dtype, torch.dtype]) -> bool:
    r"""
    Verify if the dtype is a numpy dtype

    Parameters:
        dtype: dtype
            The dtype of a value. E.g. np.int32, str, torch.float

    Returns:
        A boolean saying if the dtype is a numpy dtype
    """
    is_torch = is_dtype_torch_tensor(dtype)
    is_num = dtype in (int, float, complex)
    if hasattr(dtype, "__module__"):
        is_numpy = dtype.__module__ == "numpy"
    else:
        is_numpy = False

    return (is_num or is_numpy) and not is_torch


def one_of_k_encoding(val: Any, classes: Iterable[Any]) -> List[int]:
    r"""Converts a single value to a one-hot vector.

    Parameters:
        val: int
            class to be converted into a one hot vector
            (integers from 0 to num_classes).
        num_classes: iterator
            a list or 1D array of allowed
            choices for val to take
    Returns:
        A list of length len(num_classes) + 1
    """
    encoding = [False] * (len(classes) + 1)
    found = False
    for i, v in enumerate(classes):
        if v == val:
            encoding[i] = True
            found = True
            break
    if not found:
        encoding[-1] = True
    return encoding


def is_device_cuda(device: torch.device, ignore_errors: bool = False) -> bool:
    r"""Check wheter the given device is a cuda device.

    Parameters:
        device: str, torch.device
            object to check for cuda
        ignore_errors: bool
            Whether to ignore the error if the device is not recognized.
            Otherwise, ``False`` is returned in case of errors.
    Returns:
        is_cuda: bool
    """

    if ignore_errors:
        is_cuda = False
        try:
            is_cuda = torch.device(device).type == "cuda"
        except:
            pass
    else:
        is_cuda = torch.device(device).type == "cuda"
    return is_cuda


def nan_mean(input: Tensor, *args, **kwargs) -> Tensor:
    r"""
    Return the mean of all elements, while ignoring the NaNs.

    Parameters:

        input: The input tensor.

        dim (int or tuple(int)): The dimension or dimensions to reduce.

        keepdim (bool): whether the output tensor has dim retained or not.

        dtype (torch.dtype, optional):
            The desired data type of returned tensor.
            If specified, the input tensor is casted to dtype before the operation is performed.
            This is useful for preventing data type overflows. Default: None.

    Returns:
        output: The resulting mean of the tensor
    """

    sum = torch.nansum(input, *args, **kwargs)
    num = torch.sum(~torch.isnan(input), *args, **kwargs)
    mean = sum / num
    return mean


def nan_median(input: Tensor, **kwargs) -> Tensor:
    r"""
    Return the median of all elements, while ignoring the NaNs.
    Contrarily to `torch.nanmedian`, this function supports a list
    of dimensions, or `dim=None`, and does not return the index of the median

    Parameters:

        input: The input tensor.

        dim (int or tuple(int)): The dimension or dimensions to reduce.

        keepdim (bool): whether the output tensor has dim retained or not.

        dtype (torch.dtype, optional):
            The desired data type of returned tensor.
            If specified, the input tensor is casted to dtype before the operation is performed.
            This is useful for preventing data type overflows. Default: None.

    Returns:
        output: The resulting median of the tensor.
            Contrarily to `torch.median`, it does not return the index of the median
    """

    dim = kwargs.pop("dim", None)
    keepdim = kwargs.pop("keepdim", False)

    if isinstance(dim, Iterable) and not isinstance(dim, str):
        dim = list(dim)
        dim.sort()
        # Implement the median for a list of dimensions
        for d in dim:
            input = input.unsqueeze(-1)
            input = input.transpose(d, -1)
        input = input.flatten(-len(dim))
        median, _ = torch.nanmedian(input, dim=-1, keepdim=False)
        if not keepdim:
            for d in dim[::-1]:
                median = median.squeeze(d)
    else:
        if dim is None:
            median = torch.nanmedian(input.flatten().float()).to(input.dtype)
        else:
            median, _ = torch.nanmedian(input, dim=dim, keepdim=keepdim)

    return median


def nan_var(input: Tensor, unbiased: bool = True, **kwargs) -> Tensor:
    r"""
    Return the variace of all elements, while ignoring the NaNs.
    If unbiased is True, Bessel’s correction will be used.
    Otherwise, the sample deviation is calculated, without any correction.

    Parameters:

        input: The input tensor.

        unbiased: whether to use Bessel’s correction (δN=1\delta N = 1δN=1).

        dim (int or tuple(int)): The dimension or dimensions to reduce.

        keepdim (bool): whether the output tensor has dim retained or not.

        dtype (torch.dtype, optional):
            The desired data type of returned tensor.
            If specified, the input tensor is casted to dtype before the operation is performed.
            This is useful for preventing data type overflows. Default: None.

    Returns:
        output: The resulting variance of the tensor
    """

    mean_kwargs = deepcopy(kwargs)
    mean_kwargs.pop("keepdim", None)
    dim = mean_kwargs.pop("dim", [ii for ii in range(input.ndim)])
    mean = nan_mean(input, dim=dim, keepdim=True, **mean_kwargs)
    dist = input - mean
    dist2 = dist * dist
    var = nan_mean(dist2, **kwargs)

    if unbiased:
        num = torch.sum(~torch.isnan(input), **kwargs)
        var = var * num / (num - 1)

    return var


def nan_std(input: Tensor, unbiased: bool = True, **kwargs) -> Tensor:
    r"""
    Return the standard deviation of all elements, while ignoring the NaNs.
    If unbiased is True, Bessel’s correction will be used.
    Otherwise, the sample deviation is calculated, without any correction.

    Parameters:

        input: The input tensor.

        unbiased: whether to use Bessel’s correction (δN=1\delta N = 1δN=1).

        dim (int or tuple(int)): The dimension or dimensions to reduce.

        keepdim (bool): whether the output tensor has dim retained or not.

        dtype (torch.dtype, optional):
            The desired data type of returned tensor.
            If specified, the input tensor is casted to dtype before the operation is performed.
            This is useful for preventing data type overflows. Default: None.

    Returns:
        output: The resulting standard deviation of the tensor
    """
    variances = nan_var(input=input, unbiased=unbiased, **kwargs)
    return torch.sqrt(variances.float()).to(variances.dtype)


def nan_mad(input: Tensor, normal: bool = True, **kwargs) -> Tensor:
    r"""
    Return the median absolute deviation of all elements, while ignoring the NaNs.

    Parameters:

        input: The input tensor.

        normal: whether to multiply the result by 1.4826 to mimic the
            standard deviation for normal distributions.

        dim (int or tuple(int)): The dimension or dimensions to reduce.

        keepdim (bool): whether the output tensor has dim retained or not.

        dtype (torch.dtype, optional):
            The desired data type of returned tensor.
            If specified, the input tensor is casted to dtype before the operation is performed.
            This is useful for preventing data type overflows. Default: None.

    Returns:
        output: The resulting median absolute deviation of the tensor
    """
    median_kwargs = deepcopy(kwargs)
    median_kwargs.pop("keepdim", None)
    dim = median_kwargs.pop("dim", [ii for ii in range(input.ndim)])
    median = nan_median(input, dim=dim, keepdim=True, **median_kwargs)
    dist = (input - median).abs()
    mad = nan_median(dist, **kwargs)
    if normal:
        mad = mad * 1.4826
    return mad


class ModuleWrap(torch.nn.Module):
    r"""
    Wrap a function into a `torch.nn.Module`, with possible `*args` and `**kwargs`

    Parameters:
        func: function to wrap into a module
    """

    def __init__(self, func, *args, **kwargs) -> None:
        super().__init__()
        self.func = func
        self.__name__ = f"ModuleWrap({self.func.__name__})"
        self.args = args
        self.kwargs = kwargs

    def forward(self, *args, **kwargs):
        """
        Calls the function `self.func` with the arguments `self.func(*self.args, *args, **self.kwargs, **kwargs)`
        """
        return self.func(*self.args, *args, **self.kwargs, **kwargs)

    def __repr__(self):
        return self.__name__


class ModuleListConcat(torch.nn.ModuleList):
    r"""
    A list of neural modules similar to `torch.nn.ModuleList`,
    but where the modules are applied on the same input and
    concatenated together, instead of being applied sequentially.

    Parameters:
        dim: The dimension for the concatenation
    """

    def __init__(self, dim: int = -1):
        super().__init__()
        self.dim = dim

    def forward(self, *args, **kwargs) -> Tensor:
        r"""
        Apply all layers on the `args` and `kwargs`, and concatenate
        their output alongside the dimension `self.dim`.
        """
        h = []
        for module in self:
            h.append(module.forward(*args, **kwargs))

        return torch.cat(h, dim=self.dim)


def parse_valid_args(param_dict, fn):
    r"""
    Check if a function takes the given argument.

    Parameters
    ----------
    fn: func
        The function to check the argument.
    param_dict: dict
        Dictionary of the argument.

    Returns
    -------
        param_dict: dict
            Valid paramter dictionary for the given fucntions.
    """
    param_dict_cp = copy(param_dict)
    for key, param in param_dict.items():
        if not arg_in_func(fn=fn, arg=key):
            logger.warning(
                f"{key} is not an available argument for {fn.__name__}, and is ignored by default."
            )
            param_dict_cp.pop(key)

    return param_dict_cp


def arg_in_func(fn, arg):
    r"""
    Check if a function takes the given argument.

    Parameters
    ----------
    fn: func
        The function to check the argument.
    arg: str
        The name of the argument.

    Returns
    -------
        res: bool
            True if the function contains the argument, otherwise False.
    """
    fn_args = getfullargspec(fn)
    return (fn_args.varkw is not None) or (arg in fn_args[0])


# graphium/visualization/vis_utils.py

import inspect
from copy import deepcopy

from matplotlib.offsetbox import AnchoredText


def annotate_metric(ax, metrics, x, y, fontsize=10, loc="upper left", **kwargs):
    # Compute each metric from `metrics` on the `x` and `y` data.
    # Then Annotate the plot with with the the results of each metric
    # Compute the metrics and generate strings for each metric

    stat_text = ""

    for metric_name, metric in metrics.items():
        kwargs_copy = deepcopy(kwargs)
        stat = metric(x, y, **kwargs_copy)
        stat_text += "\n" + metric_name + " = {:0.3f}".format(stat)

    stat_text = stat_text[1:]

    # Display the metrics on the plot
    _annotate(ax=ax.ax_joint, text=stat_text, loc=loc)


def _annotate(ax, text, loc="upper left", bbox_to_anchor=(1.2, 1)):
    text = text.strip()

    text_loc_outside = dict()
    if loc == "outside":
        text_loc_outside["bbox_to_anchor"] = bbox_to_anchor
        text_loc_outside["bbox_transform"] = ax.transAxes
        loc = "upper left"

    anchored_text = AnchoredText(text, loc=loc, **text_loc_outside)
    anchored_text.patch._alpha = 0.25
    ax.add_artist(anchored_text)


# graphium/__init__.py

from ._version import __version__

from .config import load_config

from . import utils
from . import features
from . import data
from . import nn
from . import trainer


# graphium/_version.py

from importlib.metadata import version
from importlib.metadata import PackageNotFoundError

try:
    __version__ = version("graphium")
except PackageNotFoundError:
    # package is not installed
    __version__ = "dev"


# graphium/finetuning/__init__.py

from .finetuning import MolecularFinetuning

from .utils import modify_cfg_for_finetuning

# graphium/finetuning/finetuning.py

from typing import Iterable, List, Dict, Tuple, Union, Callable, Any, Optional, Type

import pytorch_lightning as pl

from torch.optim.optimizer import Optimizer
from pytorch_lightning.callbacks import BaseFinetuning


class MolecularFinetuning(BaseFinetuning):
    def __init__(
        self,
        cfg,
        train_bn: bool = False

    ):
        """
        Finetuning logic
        """
        super().__init__()

        cfg_finetune = cfg['finetuning']
        self.depth = cfg_finetune['added_depth']
        try:
            self.depth += cfg_finetune['add_finetune_depth']
        except:
            pass
        self.task = cfg_finetune['task']
        self.train_bn = train_bn
        self.epoch_unfreeze_all = cfg_finetune['epoch_unfreeze_all']

    def freeze_before_training(
        self, pl_module
    ):
        # Initially freeze everything
        self.freeze(modules=pl_module, train_bn=self.train_bn)

    def finetune_function(
        self,
        pl_module: pl.LightningModule,
        epoch: int,
        optimizer: Optimizer
    ):
        # Unfreeze finetuning layers
        if epoch == 0:       
            self.unfreeze_and_add_param_group(
                modules=pl_module.model.task_heads.task_heads[self.task].layers[-self.depth:],
                optimizer=optimizer,
                train_bn=self.train_bn
            )

        if epoch == self.epoch_unfreeze_all:       
            self.unfreeze_and_add_param_group(
                modules=pl_module,
                optimizer=optimizer,
                train_bn=self.train_bn
            )

# graphium/finetuning/utils.py

from copy import deepcopy


def modify_cfg_for_finetuning(cfg):

    cfg_finetune = cfg['finetuning']
    module = cfg_finetune['module_from_pretrained']
    task = cfg_finetune['task']
    level = cfg_finetune['level']
    task_head_from_pretrained = cfg_finetune['task_head_from_pretrained']

    new_module_kwargs = deepcopy(cfg['architecture'][module][task_head_from_pretrained])

    upd_kwargs = {
        'task_level': level,
        'out_dim': cfg_finetune['new_out_dim'],
        'depth': new_module_kwargs['depth'] + cfg_finetune['added_depth'] - cfg_finetune['drop_depth']
    }

    new_module_kwargs.update(upd_kwargs)

    cfg['architecture'][module] = {
        task: new_module_kwargs
    }

    return cfg

# graphium/finetuning/finetuning_architecture.py
from typing import Iterable, List, Dict, Tuple, Union, Callable, Any, Optional, Type

from copy import deepcopy

import torch
import torch.nn as nn

from torch import Tensor
from torch_geometric.data import Batch

from graphium.data.utils import get_keys
from graphium.nn.architectures.encoder_manager import EncoderManager
from graphium.nn.architectures import FullGraphMultiTaskNetwork, FeedForwardNN, FeedForwardPyg, TaskHeads
from graphium.nn.architectures.global_architectures import FeedForwardGraph
from graphium.trainer.predictor_options import ModelOptions
from graphium.nn.utils import MupMixin

FINETUNING_HEAD_DICT = {
    "mlp": FeedForwardNN,
    "gnn": FeedForwardPyg
}


class FullGraphFinetuningNetwork(FullGraphMultiTaskNetwork):
    def __init__(
        self,
        gnn_kwargs: Dict[str, Any],
        pre_nn_kwargs: Optional[Dict[str, Any]] = None,
        pre_nn_edges_kwargs: Optional[Dict[str, Any]] = None,
        pe_encoders_kwargs: Optional[Dict[str, Any]] = None,
        task_heads_kwargs: Optional[Dict[str, Any]] = None,
        graph_output_nn_kwargs: Optional[Dict[str, Any]] = None,
        finetuning_head_kwargs: Optional[Dict[str, Any]] = None,
        accelerator_kwargs: Optional[Dict[str, Any]] = None,
        num_inference_to_average: int = 1,
        last_layer_is_readout: bool = False,
        name: str = "FullFinetuningGNN",
    ):
        r"""
        Class that allows to implement a full graph neural network architecture,
        including the pre-processing MLP and the post processing MLP.

        Parameters:

            gnn_kwargs:
                key-word arguments to use for the initialization of the pre-processing
                GNN network using the class `FeedForwardGraph`.
                It must respect the following criteria:

                - gnn_kwargs["in_dim"] must be equal to pre_nn_kwargs["out_dim"]
                - gnn_kwargs["out_dim"] must be equal to graph_output_nn_kwargs["in_dim"]

            pe_encoders_kwargs:
                key-word arguments to use for the initialization of all positional encoding encoders.
                See the class `EncoderManager` for more details.

            pre_nn_kwargs:
                key-word arguments to use for the initialization of the pre-processing
                MLP network of the node features before the GNN, using the class `FeedForwardNN`.
                If `None`, there won't be a pre-processing MLP.

            pre_nn_edges_kwargs:
                key-word arguments to use for the initialization of the pre-processing
                MLP network of the edge features before the GNN, using the class `FeedForwardNN`.
                If `None`, there won't be a pre-processing MLP.

            task_heads_kwargs:
                This argument is a list of dictionaries containing the arguments for task heads. Each argument is used to
                initialize a task-specific MLP.

            graph_output_nn_kwargs:
                This argument is a list of dictionaries corresponding to the arguments for a FeedForwardNN.
                Each dict of arguments is used to initialize a shared MLP.

            finetuning_head_kwargs:
                key-word arguments to use for the finetuning head.
                It must respect the following criteria:

                - [last_used_module]_kwarg["out_level"] must be equal to finetuning_head_kwargs["in_level"]
                - [last_used_module]_kwarg["out_dim"] must be equal to finetuning_head_kwargs["in_dim"]

                Here, [last_used_module] represents the module that is finetuned from,
                e.g., gnn, graph_output or (one of the) task_heads

            accelerator_kwargs:
                key-word arguments specific to the accelerator being used,
                e.g. pipeline split points

            num_inference_to_average:
                Number of inferences to average at val/test time. This is used to avoid the noise introduced
                by positional encodings with sign-flips. In case no such encoding is given,
                this parameter is ignored.
                NOTE: The inference time will be slowed-down proportionaly to this parameter.

            last_layer_is_readout: Whether the last layer should be treated as a readout layer.
                Allows to use the `mup.MuReadout` from the muTransfer method https://github.com/microsoft/mup

            name:
                Name attributed to the current network, for display and printing
                purposes.
        """

        super().__init__(
            gnn_kwargs,
            pre_nn_kwargs,
            pre_nn_edges_kwargs,
            pe_encoders_kwargs,
            task_heads_kwargs,
            graph_output_nn_kwargs,
            accelerator_kwargs,
            num_inference_to_average,
            last_layer_is_readout,
            name
        )

        self.name = name
        self.num_inference_to_average = num_inference_to_average
        self.last_layer_is_readout = last_layer_is_readout
        self._concat_last_layers = None
        self.pre_nn, self.pre_nn_edges, self.task_heads, self.finetuning_head = None, None, None, None
        self.pe_encoders_kwargs = deepcopy(pe_encoders_kwargs)
        self.graph_output_nn_kwargs = graph_output_nn_kwargs
        self.finetuning_head_kwargs = finetuning_head_kwargs
        self.encoder_manager = EncoderManager(pe_encoders_kwargs)
        self.max_num_nodes_per_graph = None
        self.max_num_edges_per_graph = None

        # Initialize the pre-processing neural net for nodes (applied directly on node features)
        if pre_nn_kwargs is not None:
            name = pre_nn_kwargs.pop("name", "pre-NN")
            self.pre_nn = FeedForwardNN(**pre_nn_kwargs, name=name)
            next_in_dim = self.pre_nn.out_dim
            gnn_kwargs.setdefault("in_dim", next_in_dim)
            assert (
                next_in_dim == gnn_kwargs["in_dim"]
            ), f"Inconsistent dimensions between pre-NN output ({next_in_dim}) and GNN input ({gnn_kwargs['in_dim']})"

        # Initialize the pre-processing neural net for edges (applied directly on edge features)
        if pre_nn_edges_kwargs is not None:
            name = pre_nn_edges_kwargs.pop("name", "pre-NN-edges")
            self.pre_nn_edges = FeedForwardNN(**pre_nn_edges_kwargs, name=name)
            next_in_dim = self.pre_nn_edges.out_dim
            gnn_kwargs.setdefault("in_dim_edges", next_in_dim)
            assert (
                next_in_dim == gnn_kwargs["in_dim_edges"]
            ), f"Inconsistent dimensions between pre-NN-edges output ({next_in_dim}) and GNN input ({gnn_kwargs['in_dim_edges']})"

        # Initialize the graph neural net (applied after the pre_nn)
        name = gnn_kwargs.pop("name", "GNN")
        gnn_class = FeedForwardGraph
        gnn_kwargs.setdefault(
            "last_layer_is_readout", self.last_layer_is_readout and (task_heads_kwargs is None)
        )
        self.gnn = gnn_class(**gnn_kwargs, name=name)
        next_in_dim = self.gnn.out_dim

        if task_heads_kwargs is not None:
            self.task_heads = TaskHeads(
                in_dim=self.out_dim,
                in_dim_edges=self.out_dim_edges,
                task_heads_kwargs=task_heads_kwargs,
                graph_output_nn_kwargs=graph_output_nn_kwargs,
            )
            self._task_heads_kwargs = task_heads_kwargs

        if finetuning_head_kwargs is not None:
            self.finetuning_head = FinetuningHead(finetuning_head_kwargs)
            # model_type = finetuning_head_kwargs.pop("model_type", "mlp")
            # finetuning_model = FINETUNING_HEAD_DICT[model_type]
            # self.finetuning_head = finetuning_model(**finetuning_head_kwargs)

        if accelerator_kwargs is not None:
            accelerator = accelerator_kwargs["_accelerator"]
            if accelerator == "ipu":
                self._apply_ipu_options(accelerator_kwargs)

    def forward(self, g: Batch) -> Tensor:
        r"""
        Apply the pre-processing neural network, the graph neural network,
        and the post-processing neural network on the graph features.

        Parameters:

            g:
                pyg Batch graph on which the convolution is done.
                Must contain the following elements:

                - Node key `"feat"`: `torch.Tensor[..., N, Din]`.
                  Input node feature tensor, before the network.
                  `N` is the number of nodes, `Din` is the input features dimension ``self.pre_nn.in_dim``

                - Edge key `"edge_feat"`: `torch.Tensor[..., N, Ein]` **Optional**.
                  The edge features to use. It will be ignored if the
                  model doesn't supporte edge features or if
                  `self.in_dim_edges==0`.

                - Other keys related to positional encodings `"pos_enc_feats_sign_flip"`,
                  `"pos_enc_feats_no_flip"`.

        Returns:

            `torch.Tensor[..., M, Dout]` or `torch.Tensor[..., N, Dout]`:
                Node or graph feature tensor, after the network.
                `N` is the number of nodes, `M` is the number of graphs,
                `Dout` is the output dimension ``self.graph_output_nn.out_dim``
                If the `self.gnn.pooling` is [`None`], then it returns node features and the output dimension is `N`,
                otherwise it returns graph features and the output dimension is `M`

        """

        # Apply the positional encoders
        g = self.encoder_manager(g)

        g["feat"] = g["feat"]
        e = None

        if "edge_feat" in get_keys(g):
            g["edge_feat"] = g["edge_feat"]

        # Run the pre-processing network on node features
        if self.pre_nn is not None:
            g["feat"] = self.pre_nn.forward(g["feat"])

        # Run the pre-processing network on edge features
        # If there are no edges, skip the forward and change the dimension of e
        if self.pre_nn_edges is not None:
            e = g["edge_feat"]
            if torch.prod(torch.as_tensor(e.shape[:-1])) == 0:
                e = torch.zeros(
                    list(e.shape[:-1]) + [self.pre_nn_edges.out_dim], device=e.device, dtype=e.dtype
                )
            else:
                e = self.pre_nn_edges.forward(e)
            g["edge_feat"] = e

        # Run the graph neural network
        g = self.gnn.forward(g)

        if self.task_heads is not None:
            g = self.task_heads.forward(g)

        # if self.task_heads is not None and self.finetuning_head is not None:
        #     task, g = list(g.items())[0]

        if self.finetuning_head is not None:
            # g = {task: self.finetuning_head.forward(g)}
            g = self.finetuning_head.forward(g)

        return g

    def make_mup_base_kwargs(self, divide_factor: float = 2.0) -> Dict[str, Any]:
        """
        Create a 'base' model to be used by the `mup` or `muTransfer` scaling of the model.
        The base model is usually identical to the regular model, but with the
        layers width divided by a given factor (2 by default)

        Parameter:
            divide_factor: Factor by which to divide the width.

        Returns:
            Dictionary with the kwargs to create the base model.
        """
        kwargs = dict(
            gnn_kwargs=None,
            pre_nn_kwargs=None,
            pre_nn_edges_kwargs=None,
            pe_encoders_kwargs=None,
            finetuning_head_kwargs=None,
            num_inference_to_average=self.num_inference_to_average,
            last_layer_is_readout=self.last_layer_is_readout,
            name=self.name,
        )

        # For the pre-nn network, get the smaller dimensions.
        # For the input dim, only divide the features coming from the pe-encoders
        if self.pre_nn is not None:
            kwargs["pre_nn_kwargs"] = self.pre_nn.make_mup_base_kwargs(
                divide_factor=divide_factor, factor_in_dim=False
            )
            pe_enc_outdim = 0 if self.encoder_manager is None else self.pe_encoders_kwargs.get("out_dim", 0)
            pre_nn_indim = kwargs["pre_nn_kwargs"]["in_dim"] - pe_enc_outdim
            kwargs["pre_nn_kwargs"]["in_dim"] = round(pre_nn_indim + (pe_enc_outdim / divide_factor))

        # For the pre-nn on the edges, factor all dimensions, except the in_dim
        if self.pre_nn_edges is not None:
            kwargs["pre_nn_edges_kwargs"] = self.pre_nn_edges.make_mup_base_kwargs(
                divide_factor=divide_factor, factor_in_dim=False
            )
            pe_enc_edge_outdim = (
                0 if self.encoder_manager is None else self.pe_encoders_kwargs.get("edge_out_dim", 0)
            )
            pre_nn_edge_indim = kwargs["pre_nn_edges_kwargs"]["in_dim"] - pe_enc_edge_outdim
            kwargs["pre_nn_edges_kwargs"]["in_dim"] = round(
                pre_nn_edge_indim + (pe_enc_edge_outdim / divide_factor)
            )

        # For the pe-encoders, don't factor the in_dim and in_dim_edges
        if self.encoder_manager is not None:
            kwargs["pe_encoders_kwargs"] = self.encoder_manager.make_mup_base_kwargs(
                divide_factor=divide_factor
            )

        if self.task_heads is not None:
            task_heads_kwargs = self.task_heads.make_mup_base_kwargs(
                divide_factor=divide_factor, factor_in_dim=True
            )
            kwargs["task_heads_kwargs"] = task_heads_kwargs["task_heads_kwargs"]
            kwargs["graph_output_nn_kwargs"] = task_heads_kwargs["graph_output_nn_kwargs"]

        # For the gnn network, all the dimension are divided, except the input dims if pre-nn are missing
        if self.gnn is not None:
            factor_in_dim = self.pre_nn is not None
            kwargs["gnn_kwargs"] = self.gnn.make_mup_base_kwargs(
                divide_factor=divide_factor,
                factor_in_dim=factor_in_dim,
            )

        if self.finetuning_head is not None:
            kwargs["finetuning_head_kwargs"] = self.finetuning_head.make_mup_base_kwargs(
                divide_factor=divide_factor,
                factor_in_dim=True
            )

        return kwargs
    

class FinetuningHead(nn.Module, MupMixin):
    def __init__(
        self,
        finetuning_head_kwargs: Dict[str, Any]
    ):
        r"""
        A flexible neural network architecture, with variable hidden dimensions,
        support for multiple layer types, and support for different residual
        connections.

        This class is meant to work with different graph neural networks
        layers. Any layer must inherit from `graphium.nn.base_graph_layer.BaseGraphStructure`
        or `graphium.nn.base_graph_layer.BaseGraphLayer`.

        Parameters:

            ...

        """

        super().__init__()
        self.task = finetuning_head_kwargs.pop("task", None)
        self.previous_module = finetuning_head_kwargs.pop("previous_module", "task_heads")
        self.incoming_level = finetuning_head_kwargs.pop("incoming_level", "graph")

        model_type = finetuning_head_kwargs.pop("model_type", "mlp")
        finetuning_model = FINETUNING_HEAD_DICT[model_type]
        self.finetuning_head = finetuning_model(**finetuning_head_kwargs)

    def forward(self, g: Union[torch.Tensor, Batch]):

        if self.previous_module == "task_heads":
            g = list(g.values())[0]

        g = self.finetuning_head.forward(g)

        return {self.task: g}

    def make_mup_base_kwargs(self, divide_factor: float = 2.0, factor_in_dim: bool = False) -> Dict[str, Any]:
        """
        Create a 'base' model to be used by the `mup` or `muTransfer` scaling of the model.
        The base model is usually identical to the regular model, but with the
        layers width divided by a given factor (2 by default)

        Parameter:
            divide_factor: Factor by which to divide the width.
            factor_in_dim: Whether to factor the input dimension

        Returns:
            Dictionary with the kwargs to create the base model.
        """
        # For the post-nn network, all the dimension are divided

        return self.finetuning_head.make_mup_base_kwargs(divide_factor=divide_factor, factor_in_dim=factor_in_dim)