# tests/mtl/test_mtl_architecture.py

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

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

import torch
import unittest as ut
import dgl
from copy import deepcopy

import goli
from goli.config._loader import load_architecture
from goli.nn.architectures import TaskHeads, FullDGLMultiTaskNetwork, TaskHead, TaskHeadParams
from goli.nn.base_layers import FCLayer
from goli.nn.residual_connections import (
    ResidualConnectionConcat,
    ResidualConnectionDenseNet,
    ResidualConnectionNone,
    ResidualConnectionSimple,
    ResidualConnectionWeighted,
)

kwargs = {
    "activation": "relu",
    "last_activation": "none",
    "normalization": "none",
    "dropout": 0.2,
    "name": "LNN",
    "layer_type": FCLayer,
    "residual_type": "none",
    "residual_skip_steps": 1,
}

task_1_kwargs = {
    "task_name": "task_1",
    "out_dim": 5,
    "hidden_dims": [5, 6, 7],
}
task_2_kwargs = {
    "task_name": "task_2",
    "out_dim": 3,
    "hidden_dims": [8, 9, 10],
}
task_3_kwargs = {
    "task_name": "task_3",
    "out_dim": 4,
    "hidden_dims": [2, 2, 2],
}

# The params to create the task head MLPs. 
task_1_params = TaskHeadParams(
    **task_1_kwargs,
    **kwargs,
)
task_2_params = TaskHeadParams(
    **task_2_kwargs,
    **kwargs,
)
task_3_params = TaskHeadParams(
    **task_3_kwargs,
    **kwargs,
)

class test_TaskHeads(ut.TestCase):
    def test_task_head_forward(self):
        """Task heads are FeedForwardNNs, with the addition of a 'task_name' attribute.
        We want to make sure that they work as expected even though they are basically the same as FFNNs.
        
        This test matches `test_forward_no_residual` for the FeedForwardNN."""
        
        # Assume same setup as when testing the FFNN
        kwargs = {
            "activation": "relu",
            "last_activation": "none",
            "normalization": "none",
            "dropout": 0.2,
            "name": "LNN",
            "layer_type": FCLayer,
        }
        
        task_name = "task"
        in_dim = 8
        out_dim = 16
        hidden_dims = [5, 6, 7]
        batch = 2

        task_head = TaskHead(
            task_name=task_name,
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="none",
            residual_skip_steps=1,
            **kwargs,
        )

        # Check that the task name is correct
        self.assertEqual(task_head.task_name, task_name)

        # Check that the layers were created as intended
        self.assertEqual(len(task_head.layers), len(hidden_dims) + 1)
        self.assertEqual(task_head.layers[0].in_dim, in_dim)
        self.assertEqual(task_head.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(task_head.layers[2].in_dim, hidden_dims[1])
        self.assertEqual(task_head.layers[3].in_dim, hidden_dims[2])

        h = torch.FloatTensor(batch, in_dim)
        h_out = task_head.forward(h)                # For now just inherits the forward() of FFNN

        # Check that the output has the correct dimensions
        self.assertListEqual(list(h_out.shape), [batch, out_dim])

    def test_task_heads_forward(self):

        in_dim = 4      # Dimension of the incoming data
        batch = 2

        task_heads_params = [task_1_params, task_2_params, task_3_params]

        # 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, task_heads_params=task_heads_params)

        # Test the sizes of the MLPs for each head
        # Head for task_1
        task_1_head = multi_head_nn.task_heads["task_1"]
        # Check that the task name is correct
        self.assertEqual(task_1_head.task_name, task_1_kwargs["task_name"])
        # 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 that the task name is correct
        self.assertEqual(task_2_head.task_name, task_2_kwargs["task_name"])
        # 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 that the task name is correct
        self.assertEqual(task_3_head.task_name, task_3_kwargs["task_name"])
        # 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])

        h = torch.FloatTensor(batch, in_dim)
        h_out = multi_head_nn.forward(h)

        # Check the output: It's a per-task prediction!
        self.assertListEqual(list(h_out[task_1_kwargs["task_name"]].shape), [batch, task_1_kwargs["out_dim"]])
        self.assertListEqual(list(h_out[task_2_kwargs["task_name"]].shape), [batch, task_2_kwargs["out_dim"]])
        self.assertListEqual(list(h_out[task_3_kwargs["task_name"]].shape), [batch, task_3_kwargs["out_dim"]])


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

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

    g1 = dgl.graph((torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])))
    g2 = dgl.graph((torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])))
    g1.ndata["feat"] = torch.zeros(g1.num_nodes(), in_dim, dtype=torch.float32)
    g1.edata["feat"] = torch.ones(g1.num_edges(), in_dim_edges, dtype=torch.float32)
    g2.ndata["feat"] = torch.ones(g2.num_nodes(), in_dim, dtype=torch.float32)
    g2.edata["feat"] = torch.zeros(g2.num_edges(), in_dim_edges, dtype=torch.float32)
    batch = [g1, g2, deepcopy(g1), deepcopy(g2)]
    batch = [dgl.add_self_loop(g) for g in batch]
    bg = dgl.batch(batch)

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

    gnn_layers_kwargs = {
        "gcn": {},
        "gin": {},
        "gat": {"layer_kwargs": {"num_heads": 3}},
        "gated-gcn": {"in_dim_edges": in_dim_edges, "hidden_dims_edges": hidden_dims},
        "pna-conv": {"layer_kwargs": pna_kwargs},
        "pna-msgpass#1": {"layer_kwargs": pna_kwargs, "in_dim_edges": 0},
        "pna-msgpass#2": {"layer_kwargs": pna_kwargs, "in_dim_edges": in_dim_edges},
    }

    def test_fulldgl_multitask_forward(self):
        # Task heads
        task_heads_params = [task_1_params, task_2_params, task_3_params]

        # Params for the FullDGLNetwork
        temp_dim_1 = 5
        temp_dim_2 = 17

        pre_nn_kwargs = dict(in_dim=self.in_dim, out_dim=temp_dim_1, hidden_dims=[4, 4, 4, 4, 4])

        post_nn_kwargs = dict(in_dim=temp_dim_2, out_dim=self.fulldgl_out_dim, hidden_dims=[3, 3, 3, 3])

        for pooling in [["none"], ["sum"], ["mean", "s2s", "max"]]:
            for residual_skip_steps in [1, 2, 3]:
                for virtual_node in self.virtual_nodes:
                    for normalization in self.norms:
                        for layer_name, this_kwargs in self.gnn_layers_kwargs.items():
                            err_msg = f"pooling={pooling}, virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                            layer_type = layer_name.split("#")[0]

                            gnn_kwargs = dict(
                                in_dim=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,
                                pooling=pooling,
                                **this_kwargs,
                                **self.fulldgl_kwargs,
                            )

                            multitask_fulldgl_nn = FullDGLMultiTaskNetwork(
                                task_heads_kwargs=task_heads_params,
                                gnn_kwargs=gnn_kwargs,
                                pre_nn_kwargs=pre_nn_kwargs,
                                post_nn_kwargs=post_nn_kwargs,
                            )

                            bg = deepcopy(self.bg)
                            h_out = multitask_fulldgl_nn.forward(bg)

                            dim_1 = bg.num_nodes() if pooling == ["none"] else bg.batch_size
                            #self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)

                            self.assertListEqual(list(h_out[task_1_kwargs["task_name"]].shape), [dim_1, task_1_kwargs["out_dim"]], msg=err_msg)
                            self.assertListEqual(list(h_out[task_2_kwargs["task_name"]].shape), [dim_1, task_2_kwargs["out_dim"]], msg=err_msg)
                            self.assertListEqual(list(h_out[task_3_kwargs["task_name"]].shape), [dim_1, task_3_kwargs["out_dim"]], msg=err_msg)


class test_FullDGLMultiTaskNetwork(ut.TestCase):

    in_dim_nodes = 7
    in_dim_edges = 13

    g1 = dgl.graph((torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])))
    g2 = dgl.graph((torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])))
    g1.ndata["feat"] = torch.zeros(g1.num_nodes(), in_dim_nodes, dtype=torch.float32)
    g1.edata["feat"] = torch.ones(g1.num_edges(), in_dim_edges, dtype=torch.float32)
    g2.ndata["feat"] = torch.ones(g2.num_nodes(), in_dim_nodes, dtype=torch.float32)
    g2.edata["feat"] = torch.zeros(g2.num_edges(), in_dim_edges, dtype=torch.float32)
    batch = [g1, g2, deepcopy(g1), deepcopy(g2)]
    batch = [dgl.add_self_loop(g) for g in batch]
    bg = dgl.batch(batch)

    def test_fulldglmultitasknetwork_from_config(self):
        cfg = goli.load_config(name="zinc_default_multitask_fulldgl")

        # Initialize the network
        model_class, model_kwargs = load_architecture(
            cfg,
            in_dim_nodes=self.in_dim_nodes,
            in_dim_edges=self.in_dim_edges,
        )

        multitask_fulldgl_nn = model_class(**model_kwargs)

        # Test 
        bg = deepcopy(self.bg)
        h_out = multitask_fulldgl_nn.forward(bg)

        dim_1 = self.bg.batch_size

        self.assertListEqual(list(h_out[task_1_kwargs["task_name"]].shape), [dim_1, task_1_kwargs["out_dim"]])
        self.assertListEqual(list(h_out[task_2_kwargs["task_name"]].shape), [dim_1, task_2_kwargs["out_dim"]])
        self.assertListEqual(list(h_out[task_3_kwargs["task_name"]].shape), [dim_1, task_3_kwargs["out_dim"]])


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


# goli/nn/dgl_layers/__init__.py

from goli.nn.dgl_layers.base_dgl_layer import BaseDGLLayer
from goli.nn.dgl_layers.gat_layer import GATLayer
from goli.nn.dgl_layers.gcn_layer import GCNLayer
from goli.nn.dgl_layers.gin_layer import GINLayer
from goli.nn.dgl_layers.gated_gcn_layer import GatedGCNLayer
from goli.nn.dgl_layers.pna_layer import PNAConvolutionalLayer, PNAMessagePassingLayer
from goli.nn.dgl_layers.dgn_layer import DGNConvolutionalLayer, DGNMessagePassingLayer


# goli/nn/dgl_layers/base_dgl_layer.py

import torch
import torch.nn as nn
import abc
from typing import List, Dict, Tuple, Union, Callable

from goli.nn.base_layers import get_activation
from goli.utils.decorators import classproperty


class BaseDGLLayer(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",
    ):
        r"""
        Abstract class used to standardize the implementation of DGL 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
        """

        super().__init__()

        # Basic attributes
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.normalization = normalization
        self.dropout = dropout
        self.activation = activation

        # Build the layers
        self.activation_layer = get_activation(activation)

        self.dropout_layer = None
        if dropout > 0:
            self.dropout_layer = nn.Dropout(p=dropout)

        self.norm_layer = self._parse_norm(normalization)

    def _parse_norm(self, normalization):

        if normalization is None or normalization == "none":
            parsed_norm = None
        elif callable(normalization):
            parsed_norm = normalization
        elif normalization == "batch_norm":
            parsed_norm = nn.BatchNorm1d(self.out_dim * self.out_dim_factor)
        elif normalization == "layer_norm":
            parsed_norm = nn.LayerNorm(self.out_dim * self.out_dim_factor)
        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,
        h: torch.Tensor,
        normalization: Union[str, Callable] = True,
        activation: bool = True,
        dropout: bool = True,
    ):
        r"""
        Apply the different normalization and the dropout to the
        output layer.

        Parameters:

            h:
                Feature tensor, to be normalized

            normalization:
                Normalization to use. Choices:

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

            activation:
                Whether to apply the activation layer

            dropout:
                Whether to apply the dropout layer

        Returns:

            h:
                Normalized and dropped-out features

        """

        if self.norm_layer is not None:
            h = self.norm_layer(h)

        if activation and (self.activation_layer is not None):
            h = self.activation_layer(h)

        if dropout and (self.dropout_layer is not None):
            h = self.dropout_layer(h)

        return h

    @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
        """
        ...

    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})"


# goli/nn/dgl_layers/dgn_layer.py

import torch
from typing import Dict, List, Tuple, Union, Callable
from functools import partial

from goli.nn.dgl_layers.pna_layer import PNAConvolutionalLayer, PNAMessagePassingLayer
from goli.nn.pna_operations import PNA_AGGREGATORS
from goli.nn.dgn_operations import DGN_AGGREGATORS


class BaseDGNLayer:
    def parse_aggregators(self, aggregators_name: List[str]) -> List[Callable]:
        r"""
        Parse the aggregators from a list of strings into a list of callables.

        The possibilities are:

        - `"mean"`
        - `"sum"`
        - `"min"`
        - `"max"`
        - `"std"`
        - `"dir{dir_idx:int}/smooth/{Optional[temperature:float]}"`
        - `"dir{dir_idx:int}/dx_abs/{Optional[temperature:float]}"`
        - `"dir{dir_idx:int}/dx_no_abs/{Optional[temperature:float]}"`
        - `"dir{dir_idx:int}/dx_abs_balanced/{Optional[temperature:float]}"`
        - `"dir{dir_idx:int}/forward/{Optional[temperature:float]}"`
        - `"dir{dir_idx:int}/backward/{Optional[temperature:float]}"`

        `dir_idx` is an integer specifying the index of the positional encoding
        to use for direction. In the case of eigenvector-based directions, `dir_idx=1`
        is chosen for the first non-trivial eigenvector and `dir_idx=2` for the second.

        `temperature` is used to harden the direction using a softmax on the directional
        matrices. If it is not provided, then no softmax is applied. The larger the temperature,
        the more weight is attributed to the dominant direction.

        Example:
            ```
            In:     self.parse_aggregators(["dir1/dx_abs", "dir2/smooth/0.2"])
            Out:    [partial(aggregate_dir_dx_abs, dir_idx=1, temperature=None),
                     partial(aggregate_dir_smooth, dir_idx=2, temperature=0.2)]
            ```

        Parameters:
            aggregators_name: The list of all aggregators names to use, selected
                from the list of possible strings.

        Returns:
            aggregators: The list of all callable aggregators.

        """
        aggregators = []

        for agg_name in aggregators_name:
            agg_name = agg_name.lower()
            this_agg = None

            # Get the aggregator from PNA if not a directional aggregation
            if agg_name in PNA_AGGREGATORS.keys():
                this_agg = PNA_AGGREGATORS[agg_name]

            # If the directional, get the right aggregator
            elif "dir" == agg_name[:3]:
                agg_split = agg_name.split("/")
                agg_dir, agg_fn_name = agg_split[0], agg_split[1]
                dir_idx = int(agg_dir[3:])
                temperature = None
                if len(agg_split) > 2:
                    temperature = float(agg_split[2])
                this_agg = partial(DGN_AGGREGATORS[agg_fn_name], dir_idx=dir_idx, temperature=temperature)

            if this_agg is None:
                raise ValueError(f"aggregator `{agg_name}` not a valid choice.")

            aggregators.append(this_agg)

        return aggregators

    def message_func(self, edges) -> Dict[str, torch.Tensor]:
        r"""
        The message function to generate messages along the edges.
        """
        return {
            "e": edges.data["e"],
            "source_pos": edges.data["source_pos"],
            "dest_pos": edges.data["dest_pos"],
        }

    def reduce_func(self, nodes) -> Dict[str, torch.Tensor]:
        r"""
        The reduce function to aggregate the messages.
        Apply the aggregators and scalers, and concatenate the results.
        """
        h_in = nodes.data["h"]
        h = nodes.mailbox["e"]
        source_pos = nodes.mailbox["source_pos"]
        dest_pos = nodes.mailbox["dest_pos"]
        D = h.shape[-2]

        # aggregators and scalers
        h_to_cat = [
            aggr(h=h, h_in=h_in, source_pos=source_pos, dest_pos=dest_pos) for aggr in self.aggregators
        ]
        h = torch.cat(h_to_cat, dim=-1)
        if len(self.scalers) > 1:
            h = torch.cat([scale(h, D=D, avg_d=self.avg_d) for scale in self.scalers], dim=1)

        return {"h": h}


class DGNConvolutionalLayer(BaseDGNLayer, PNAConvolutionalLayer):
    r"""
    Implementation of the convolutional architecture of the DGN layer,
    previously known as `DGNSimpleLayer`. This layer aggregates the
    neighbouring messages using multiple aggregators and scalers,
    concatenates their results, then applies an MLP on the concatenated
    features.

    DGN: Directional Graph Networks
    Dominique Beaini, Saro Passaro, Vincent Létourneau, William L. Hamilton, Gabriele Corso, Pietro Liò
    https://arxiv.org/pdf/2010.02863.pdf
    """

    def parse_aggregators(self, aggregators: List[str]) -> List[Callable]:
        return BaseDGNLayer.parse_aggregators(self, aggregators)

    def message_func(self, edges) -> Dict[str, torch.Tensor]:
        return BaseDGNLayer.message_func(self, edges)

    def reduce_func(self, nodes) -> Dict[str, torch.Tensor]:
        return BaseDGNLayer.reduce_func(self, nodes)

    def pretrans_edges(self, edges):
        pretrans = PNAConvolutionalLayer.pretrans_edges(self, edges)
        pretrans.update({"source_pos": edges.src["pos_dir"], "dest_pos": edges.dst["pos_dir"]})
        return pretrans


class DGNMessagePassingLayer(BaseDGNLayer, PNAMessagePassingLayer):
    r"""
    Implementation of the message passing architecture of the DGN message passing layer,
    previously known as `DGNLayerComplex`. 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 using
    multiple aggregators and scalers,
    concatenates their results, then applies an MLP on the concatenated
    features.

    DGN: Directional Graph Networks
    Dominique Beaini, Saro Passaro, Vincent Létourneau, William L. Hamilton, Gabriele Corso, Pietro Liò
    https://arxiv.org/pdf/2010.02863.pdf
    """

    def parse_aggregators(self, aggregators: List[str]) -> List[Callable]:
        return BaseDGNLayer.parse_aggregators(self, aggregators)

    def message_func(self, edges) -> Dict[str, torch.Tensor]:
        return BaseDGNLayer.message_func(self, edges)

    def reduce_func(self, nodes) -> Dict[str, torch.Tensor]:
        return BaseDGNLayer.reduce_func(self, nodes)

    def pretrans_edges(self, edges):
        pretrans = PNAMessagePassingLayer.pretrans_edges(self, edges)
        pretrans.update({"source_pos": edges.src["pos_dir"], "dest_pos": edges.dst["pos_dir"]})
        return pretrans


# goli/nn/dgl_layers/gat_layer.py

import torch
from typing import Union, Callable

from dgl.nn.pytorch import GATConv
from dgl import DGLGraph

from goli.nn.dgl_layers.base_dgl_layer import BaseDGLLayer
from goli.utils.decorators import classproperty

"""
    GAT: Graph Attention Network
    Graph Attention Networks (Veličković et al., ICLR 2018)
    https://arxiv.org/abs/1710.10903
"""


class GATLayer(BaseDGLLayer):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        num_heads: int,
        activation="elu",
        dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
    ):
        r"""
        GAT: Graph Attention Network
        Graph Attention Networks (Veličković et al., ICLR 2018)
        https://arxiv.org/abs/1710.10903

        The implementation is built on top of the DGL ``GATCONV`` layer

        Parameters:

            in_dim: int
                Input feature dimensions of the layer

            out_dim: int
                Output feature dimensions of the layer

            num_heads: int
                Number of heads in Multi-Head Attention

            activation: str, Callable
                activation function to use in the layer

            dropout: float
                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
        """

        self.num_heads = num_heads

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

        self.gatconv = GATConv(
            in_feats=self.in_dim,
            out_feats=self.out_dim,
            num_heads=self.num_heads,
            feat_drop=self.dropout,
            attn_drop=self.dropout,
            activation=None,  # Activation is applied after
        )

    def forward(self, g: DGLGraph, h: torch.Tensor) -> torch.Tensor:
        r"""
        Apply the graph convolutional layer, with the specified activations,
        normalizations and dropout.

        Parameters:

            g: dgl.DGLGraph
                graph on which the convolution is done

            h: `torch.Tensor[..., N, Din]`
                Node feature tensor, before convolution.
                N is the number of nodes, Din is the input dimension ``self.in_dim``

        Returns:

            `torch.Tensor[..., N, Dout]`:
                Node feature tensor, after convolution.
                N is the number of nodes, Dout is the output dimension ``self.out_dim``

        """

        h = self.gatconv(g, h).flatten(1)
        self.apply_norm_activation_dropout(h, normalization="batch_norm", activation=True, dropout=False)

        return h

    @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:

            uses_edges: 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:

            uses_edges: 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:

            dim_factor: int
                Always ``self.num_heads`` for the current class
        """
        return self.num_heads


# goli/nn/dgl_layers/gated_gcn_layer.py

import torch
import torch.nn as nn
from typing import Tuple, Union, Callable
from dgl import DGLGraph

from goli.nn.dgl_layers.base_dgl_layer import BaseDGLLayer
from goli.utils.decorators import classproperty

"""
    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
"""


class GatedGCNLayer(BaseDGLLayer):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        in_dim_edges: int,
        out_dim_edges: int,
        activation: Union[Callable, str] = "relu",
        dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
    ):
        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

        """

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

        self.A = nn.Linear(in_dim, out_dim, bias=True)
        self.B = nn.Linear(in_dim, out_dim, bias=True)
        self.C = nn.Linear(in_dim_edges, out_dim, bias=True)
        self.D = nn.Linear(in_dim, out_dim, bias=True)
        self.E = nn.Linear(in_dim, out_dim, bias=True)

    def message_func(self, edges):
        Bh_j = edges.src["Bh"]
        e_ij = edges.data["Ce"] + edges.src["Dh"] + edges.dst["Eh"]  # e_ij = Ce_ij + Dhi + Ehj
        edges.data["e"] = e_ij
        return {"Bh_j": Bh_j, "e_ij": e_ij}

    def reduce_func(self, nodes):
        Ah_i = nodes.data["Ah"]
        Bh_j = nodes.mailbox["Bh_j"]
        e = nodes.mailbox["e_ij"]
        sigma_ij = torch.sigmoid(e)  # sigma_ij = sigmoid(e_ij)
        # h = Ah_i + torch.mean( sigma_ij * Bh_j, dim=1 ) # hi = Ahi + mean_j alpha_ij * Bhj
        h = Ah_i + torch.sum(sigma_ij * Bh_j, dim=1) / (
            torch.sum(sigma_ij, dim=1) + 1e-6
        )  # hi = Ahi + sum_j eta_ij/sum_j' eta_ij' * Bhj <= dense attention
        return {"h": h}

    def forward(self, g: DGLGraph, h: torch.Tensor, e: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        r"""
        Apply the graph convolutional layer, with the specified activations,
        normalizations and dropout.

        Parameters:

            g:
                graph on which the convolution is done

            h: `torch.Tensor[..., N, Din]`
                Node feature tensor, before convolution.
                N is the number of nodes, Din is the input dimension ``self.in_dim``

            e: `torch.Tensor[..., N, Din_edges]`
                Edge feature tensor, before convolution.
                N is the number of nodes, Din is the input edge dimension  ``self.in_dim_edges``

        Returns:
            `torch.Tensor[..., N, Dout]`:
                Node feature tensor, after convolution.
                N is the number of nodes, Dout is the output dimension ``self.out_dim``

            `torch.Tensor[..., N, Dout]`:
                Edge feature tensor, after convolution.
                N is the number of nodes, Dout_edges is the output edge dimension ``self.out_dim``

        """

        g.ndata["h"] = h
        g.ndata["Ah"] = self.A(h)
        g.ndata["Bh"] = self.B(h)
        g.ndata["Dh"] = self.D(h)
        g.ndata["Eh"] = self.E(h)
        g.edata["e"] = e
        g.edata["Ce"] = self.C(e)
        g.update_all(self.message_func, self.reduce_func)
        h = g.ndata["h"]  # result of graph convolution
        e = g.edata["e"]  # result of graph convolution

        h = self.apply_norm_activation_dropout(h)
        e = self.apply_norm_activation_dropout(e)

        return h, e

    @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


# goli/nn/dgl_layers/gcn_layer.py

import torch
from typing import Union, Callable

from dgl.nn.pytorch import GraphConv
from dgl import DGLGraph

from goli.nn.dgl_layers.base_dgl_layer import BaseDGLLayer
from goli.utils.decorators import classproperty

"""
    GCN: Graph Convolutional Networks
    Thomas N. Kipf, Max Welling, Semi-Supervised Classification with Graph Convolutional Networks (ICLR 2017)
    http://arxiv.org/abs/1609.02907
"""


class GCNLayer(BaseDGLLayer):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        activation: Union[str, Callable] = "relu",
        dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
    ):
        r"""
        Graph convolutional network (GCN) layer from
        Thomas N. Kipf, Max Welling, Semi-Supervised Classification with Graph Convolutional Networks (ICLR 2017)
        http://arxiv.org/abs/1609.02907

        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
        """

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

        self.conv = GraphConv(
            in_feats=in_dim,
            out_feats=out_dim,
            norm="both",
            weight=True,
            bias=True,
            activation=None,
            allow_zero_in_degree=False,
        )

    def forward(self, g: DGLGraph, h: torch.Tensor) -> torch.Tensor:
        r"""
        Apply the graph convolutional layer, with the specified activations,
        normalizations and dropout.

        Parameters:

            g:
                graph on which the convolution is done

            h: `torch.Tensor[..., N, Din]`
                Node feature tensor, before convolution.
                N is the number of nodes, Din is the input dimension ``self.in_dim``

        Returns:

            `torch.Tensor[..., N, Dout]`:
                Node feature tensor, after convolution.
                N is the number of nodes, Dout is the output dimension ``self.out_dim``

        """

        h = self.conv(g, h)
        h = self.apply_norm_activation_dropout(h)

        return h

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

        Returns:

            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


# goli/nn/dgl_layers/gin_layer.py

import torch
import dgl.function as fn
from dgl import DGLGraph
from typing import Callable, Union

from goli.nn.dgl_layers.base_dgl_layer import BaseDGLLayer
from goli.nn.base_layers import MLP
from goli.utils.decorators import classproperty

"""
    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
"""


class GINLayer(BaseDGLLayer):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        activation: Union[Callable, str] = "relu",
        dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
        init_eps: float = 0.0,
        learn_eps: bool = True,
    ):
        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 adapted from dgl 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,
        )

        # Specify to consider the edges weight in the aggregation

        # to specify whether eps is trainable or not.
        if learn_eps:
            self.eps = torch.nn.Parameter(torch.FloatTensor([init_eps]))
        else:
            self.register_buffer("eps", torch.FloatTensor([init_eps]))

        # The weights of the model, applied after the aggregation
        self.mlp = MLP(
            in_dim=self.in_dim,
            hidden_dim=self.in_dim,
            out_dim=self.out_dim,
            layers=2,
            activation=self.activation_layer,
            last_activation="none",
            normalization=self.normalization,
            last_normalization="none",
        )

    def message_func(self, g):
        r"""
        If edge weights are provided, use them to weight the messages
        """

        if "w" in g.edata.keys():
            func = fn.u_mul_e("h", "w", "m")
        else:
            func = fn.copy_u("h", "m")
        return func

    def forward(self, g: DGLGraph, h: torch.Tensor) -> torch.Tensor:
        r"""
        Apply the GIN convolutional layer, with the specified activations,
        normalizations and dropout.

        Parameters:

            g:
                graph on which the convolution is done

            h: `torch.Tensor[..., N, Din]`
                Node feature tensor, before convolution.
                N is the number of nodes, Din is the input dimension ``self.in_dim``

        Returns:

            `torch.Tensor[..., N, Dout]`:
                Node feature tensor, after convolution.
                N is the number of nodes, Dout is the output dimension ``self.out_dim``

        """

        # Aggregate the message
        g = g.local_var()
        g.ndata["h"] = h
        func = fn.copy_u("h", "m")
        g.update_all(self.message_func(g), fn.sum("m", "neigh"))
        h = (1 + self.eps) * h + g.ndata["neigh"]

        # Apply the MLP
        h = self.mlp(h)
        h = self.apply_norm_activation_dropout(h)

        return h

    @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


# goli/nn/dgl_layers/pna_layer.py

import torch
import dgl
from dgl import DGLGraph
from typing import Dict, List, Tuple, Union, Callable
from copy import deepcopy

from goli.nn.pna_operations import PNA_AGGREGATORS, PNA_SCALERS
from goli.nn.base_layers import MLP, get_activation
from goli.nn.dgl_layers.base_dgl_layer import BaseDGLLayer
from goli.utils.decorators import classproperty

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


class BasePNALayer(BaseDGLLayer):
    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: float = 1.0,
        last_activation: Union[Callable, str] = "none",
        in_dim_edges: int = 0,
    ):
        r"""
        Abstract class used to standardize the implementation of PNA layers
        in the current library.

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

        Method ``layer_inputs_edges()`` needs to be implemented in children classes

        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

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

        """

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

        # Edge dimensions
        self.in_dim_edges = in_dim_edges
        self.edge_features = self.in_dim_edges > 0

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

        # Initializing aggregators and scalers
        self.aggregators = self.parse_aggregators(aggregators)
        self.scalers = self._parse_scalers(scalers)

    def parse_aggregators(self, aggregators: List[str]) -> List[Callable]:
        r"""
        Parse the aggregators from a list of strings into a list of callables
        """
        return [PNA_AGGREGATORS[aggr] for aggr in aggregators]

    def _parse_scalers(self, scalers: List[str]) -> List[Callable]:
        r"""
        Parse the scalers from a list of strings into a list of callables
        """
        return [PNA_SCALERS[scale] for scale in scalers]

    def message_func(self, edges) -> Dict[str, torch.Tensor]:
        r"""
        The message function to generate messages along the edges.
        """
        return {"e": edges.data["e"]}

    def reduce_func(self, nodes) -> Dict[str, torch.Tensor]:
        r"""
        The reduce function to aggregate the messages.
        Apply the aggregators and scalers, and concatenate the results.
        """
        h_in = nodes.data["h"]
        h = nodes.mailbox["e"]
        D = h.shape[-2]
        h_to_cat = [aggr(h=h, h_in=h_in) for aggr in self.aggregators]
        h = torch.cat(h_to_cat, dim=-1)

        if len(self.scalers) > 1:
            h = torch.cat([scale(h, D=D, avg_d=self.avg_d) for scale in self.scalers], dim=-1)

        return {"h": h}

    def add_virtual_graph_if_no_edges(self, g: DGLGraph) -> Tuple[DGLGraph, bool]:
        r"""
        When all elements of a given batch don't have any edges
        (e.g. molecule with a single atom), the message function will
        be skipped, and the number of features will be inconsistent due
        to the variable number of aggregators.

        To fix this issue, this method creates a new graph with self-loop
        and appends it to the batch, only if all the elements of the batch
        have degree 0.

        Parameters:
            g: The batched graphs

        Returns:
            g: The batched graphs, with a possible new graph appended at the end
            no_edges:
                Whether a graph was appended to the end of the batch
                if all the elements of the batch had no edges.
        """
        no_edges = g.num_edges() == 0
        if no_edges:
            new_g = deepcopy(dgl.unbatch(g)[0])
            new_g.add_edges(0, 0)
            g = dgl.batch([g, new_g])

        return g, no_edges

    def remove_virtual_graph_if_no_edges(self, g: DGLGraph, no_edges: bool) -> DGLGraph:
        r"""
        This removes the added graph from the method
        `add_virtual_graph_if_no_edges`.

        Parameters:
            g: The batched graphs
            no_edges:
                Whether to remove the last graph of the batch
                if all the elements of the batch had no edges.

        Returns:
            g: The batched graphs, with a possible new graph appended at the end
        """
        if no_edges:
            g = dgl.batch(dgl.unbatch(g)[:-1])

        return g

    @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

    @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:
                Returns ``self.edge_features``
        """
        return self.edge_features


class PNAConvolutionalLayer(BasePNALayer):
    r"""
    Implementation of the convolutional architecture of the PNA layer,
    previously known as `PNASimpleLayer`. This layer aggregates the
    neighbouring messages using 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
    """

    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},
        last_activation: Union[Callable, str] = "none",
        posttrans_layers: int = 1,
        in_dim_edges: int = 0,
    ):
        r"""

        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

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

        """

        super().__init__(
            in_dim=in_dim,
            out_dim=out_dim,
            aggregators=aggregators,
            scalers=scalers,
            avg_d=avg_d,
            activation=activation,
            dropout=0,
            normalization="none",
            last_activation=last_activation,
            in_dim_edges=in_dim_edges,
        )

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

    def pretrans_edges(self, edges) -> Dict[str, torch.Tensor]:
        r"""
        Return a mapping to the features of the source nodes, concatenated to the
        edge data.
        """
        if self.edge_features:
            edata = torch.cat([edges.src["h"], edges.data["ef"]], dim=-1)
        else:
            edata = edges.src["h"]
        return {"e": edata}

    def forward(self, g: DGLGraph, h: torch.Tensor, e: torch.Tensor = None) -> torch.Tensor:
        r"""
        Apply the PNA convolutional layer, with the specified post transformation

        Parameters:

            g:
                graph on which the convolution is done

            h: `torch.Tensor[..., N, Din]`
                Node feature tensor, before convolution.
                N is the number of nodes, Din is the input dimension ``self.in_dim``

            e: `torch.Tensor[..., N, Din_edges]` or `None`
                Edge feature tensor, before convolution.
                N is the number of nodes, Din is the input edge dimension

                Can be set to None if the layer does not use edge features
                i.e. ``self.layer_inputs_edges -> False``

        Returns:

            `torch.Tensor[..., N, Dout]`:
                Node feature tensor, after convolution.
                N is the number of nodes, Dout is the output dimension ``self.out_dim``

        """

        g.ndata["h"] = h
        if self.edge_features:  # add the edges information only if edge_features = True
            g.edata["ef"] = e

        g, no_edges = self.add_virtual_graph_if_no_edges(g)

        g.apply_edges(self.pretrans_edges)

        # aggregation
        g.update_all(self.message_func, self.reduce_func)
        g = self.remove_virtual_graph_if_no_edges(g, no_edges=no_edges)
        h = g.ndata["h"]

        # post-transformation
        h = self.posttrans(h)

        return h


class PNAMessagePassingLayer(BasePNALayer):
    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
    """

    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},
        last_activation: Union[Callable, str] = "none",
        posttrans_layers: int = 1,
        pretrans_layers: int = 1,
        in_dim_edges: int = 0,
    ):
        r"""

        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

        """

        super().__init__(
            in_dim=in_dim,
            out_dim=out_dim,
            aggregators=aggregators,
            scalers=scalers,
            avg_d=avg_d,
            activation=activation,
            dropout=0,
            normalization="none",
            last_activation=last_activation,
            in_dim_edges=in_dim_edges,
        )

        # 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_dim=in_dim,
            out_dim=in_dim,
            layers=pretrans_layers,
            activation=self.activation,
            last_activation=self.last_activation,
            dropout=dropout,
            normalization=normalization,
            last_normalization=normalization,
        )

        # MLP used on the aggregated messages MLP(h'_u)
        self.posttrans = MLP(
            in_dim=(len(self.aggregators) * len(self.scalers) + 1) * in_dim,
            hidden_dim=out_dim,
            out_dim=out_dim,
            layers=posttrans_layers,
            activation=self.activation,
            last_activation=self.last_activation,
            dropout=dropout,
            normalization=normalization,
            last_normalization=normalization,
        )

    def pretrans_edges(self, edges) -> Dict[str, torch.Tensor]:
        r"""
        Return a mapping to the concatenation of the features from
        the source node, the destination node, and the edge between them (if applicable).
        """
        if self.edge_features:
            z2 = torch.cat([edges.src["h"], edges.dst["h"], edges.data["ef"]], dim=-1)
        else:
            z2 = torch.cat([edges.src["h"], edges.dst["h"]], dim=-1)
        return {"e": self.pretrans(z2)}

    def forward(self, g: DGLGraph, h: torch.Tensor, e: torch.Tensor = None) -> torch.Tensor:
        r"""
        Apply the PNA Message passing layer, with the specified pre/post transformations

        Parameters:

            g:
                graph on which the convolution is done

            h: `torch.Tensor[..., N, Din]`
                Node feature tensor, before convolution.
                N is the number of nodes, Din is the input dimension ``self.in_dim``

            e: `torch.Tensor[..., N, Din_edges]` or `None`
                Edge feature tensor, before convolution.
                N is the number of nodes, Din is the input edge dimension

                Can be set to None if the layer does not use edge features
                i.e. ``self.layer_inputs_edges -> False``

        Returns:

            `torch.Tensor[..., N, Dout]`:
                Node feature tensor, after convolution.
                N is the number of nodes, Dout is the output dimension ``self.out_dim``

        """

        g.ndata["h"] = h
        if self.edge_features:  # add the edges information only if edge_features = True
            g.edata["ef"] = e
        g, no_edges = self.add_virtual_graph_if_no_edges(g)

        # pretransformation
        g.apply_edges(self.pretrans_edges)

        # aggregation
        g.update_all(self.message_func, self.reduce_func)
        g = self.remove_virtual_graph_if_no_edges(g, no_edges=no_edges)
        h = torch.cat([h, g.ndata["h"]], dim=-1)

        # post-transformation
        h = self.posttrans(h)

        return h

    @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


# goli/nn/dgl_layers/pooling.py

import torch
import torch.nn as nn
from typing import List, Union, Callable, Tuple

import dgl
from dgl.nn.pytorch.glob import SumPooling, AvgPooling, MaxPooling, Set2Set, GlobalAttentionPooling
from dgl import mean_nodes, sum_nodes, max_nodes

from goli.nn.base_layers import MLP, FCLayer
from goli.utils.tensor import ModuleListConcat


EPS = 1e-6


class S2SReadout(nn.Module):
    r"""
    Performs a Set2Set aggregation of all the graph nodes' features followed by a series of fully connected layers
    """

    def __init__(
        self,
        in_dim,
        hidden_dim,
        out_dim,
        fc_layers=3,
        device="cpu",
        final_activation: Union[str, Callable] = "relu",
    ):
        super().__init__()

        # set2set aggregation
        self.set2set = Set2Set(in_dim, device=device)

        # fully connected layers
        self.mlp = MLP(
            in_dim=2 * in_dim,
            hidden_dim=hidden_dim,
            out_dim=out_dim,
            layers=fc_layers,
            activation="relu",
            last_activation=final_activation,
            normalization="batch_norm",
            last_normalization="none",
            device=device,
        )

    def forward(self, x):
        x = self.set2set(x)
        return self.mlp(x)


class StdPooling(nn.Module):
    r"""Apply standard deviation pooling over the nodes in the graph.

    $$r^{(i)} = \sigma_{k=1}^{N_i}\left( x^{(i)}_k \right)$$
    """

    def __init__(self):
        super().__init__()
        self.sum_pooler = SumPooling()
        self.relu = nn.ReLU()

    def forward(self, graph, feat):
        r"""Compute standard deviation pooling.

        Parameters:
            graph : DGLGraph
                The graph.
            feat : torch.Tensor
                The input feature with shape :math:`(N, *)` where
                :math:`N` is the number of nodes in the graph.

        Returns:
            torch.Tensor
                The output feature with shape :math:`(B, *)`, where
                :math:`B` refers to the batch size.
        """

        readout = torch.sqrt(
            self.relu((self.sum_pooler(graph, feat**2)) - (self.sum_pooler(graph, feat) ** 2)) + EPS
        )
        return readout


class MinPooling(MaxPooling):
    r"""Apply min pooling over the nodes in the graph.

    $$r^{(i)} = \min_{k=1}^{N_i}\left( x^{(i)}_k \right)$$
    """

    def forward(self, graph, feat):
        r"""Compute max pooling.

        Parameters:
            graph : DGLGraph
                The graph.
            feat : torch.Tensor
                The input feature with shape :math:`(N, *)` where
                :math:`N` is the number of nodes in the graph.

        Returns:
            readout: torch.Tensor
                The output feature with shape :math:`(B, *)`, where
                :math:`B` refers to the batch size.
        """

        return -super().forward(graph, -feat)


class DirPooling(nn.Module):
    r"""
    Apply pooling over the nodes in the graph using a directional potential
    with an inner product.

    In most cases, this is a pooling using the Fiedler vector.
    This is basically equivalent to computing a Fourier transform for the
    Fiedler vector. Then, we use the absolute value due to the sign ambiguity

    """

    def __init__(self, dir_idx):
        super().__init__()
        self.sum_pooler = SumPooling()
        self.dir_idx = dir_idx

    def forward(self, graph, feat):
        r"""Compute directional inner-product pooling, and return absolute value.

        Parameters:
            graph : DGLGraph
                The graph. Must have the key `graph.ndata["pos_dir"]`
            feat : torch.Tensor
                The input feature with shape :math:`(N, *)` where
                :math:`N` is the number of nodes in the graph.

        Returns:
            readout: torch.Tensor
                The output feature with shape :math:`(B, *)`, where
                :math:`B` refers to the batch size.
        """

        dir = graph.ndata["pos_dir"][:, self.dir_idx].unsqueeze(-1)
        pooled = torch.abs(self.sum_pooler(graph, feat * dir))

        return pooled


class LogSumPooling(AvgPooling):
    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$$
    """

    def forward(self, graph, feat):
        r"""Compute log-sum pooling.

        Parameters:
            graph : DGLGraph
                The graph.
            feat : torch.Tensor
                The input feature with shape :math:`(N, *)` where
                :math:`N` is the number of nodes in the graph.

        Returns:
            readout: torch.Tensor
                The output feature with shape :math:`(B, *)`, where
                :math:`B` refers to the batch size.
        """
        mean_pool = super().forward(graph=graph, feat=feat)
        lognum = torch.log(torch.as_tensor(graph.batch_num_nodes(), dtype=feat.dtype, device=feat.device))
        pool = mean_pool * lognum.unsqueeze(-1)
        return pool


def parse_pooling_layer(in_dim: int, pooling: Union[str, List[str]], n_iters: int = 2, n_layers: int = 2):
    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:

            - "sum": `SumPooling`
            - "mean": `MeanPooling`
            - "max": `MaxPooling`
            - "min": `MinPooling`
            - "std": `StdPooling`
            - "s2s": `Set2Set`
            - "dir{int}": `DirPooling`

        n_iters:
            IGNORED FOR ALL POOLING LAYERS, EXCEPT "s2s".
            The number of iterations.

        n_layers:
            IGNORED FOR ALL POOLING LAYERS, EXCEPT "s2s".
            The number of recurrent layers.
    """

    # TODO: Add configuration for the pooling layer kwargs

    # Create the pooling layer
    pool_layer = ModuleListConcat()
    out_pool_dim = 0
    if isinstance(pooling, str):
        pooling = [pooling]

    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(SumPooling())
        elif this_pool == "mean":
            pool_layer.append(AvgPooling())
        elif this_pool == "logsum":
            pool_layer.append(LogSumPooling())
        elif this_pool == "max":
            pool_layer.append(MaxPooling())
        elif this_pool == "min":
            pool_layer.append(MinPooling())
        elif this_pool == "std":
            pool_layer.append(StdPooling())
        elif this_pool == "s2s":
            pool_layer.append(Set2Set(input_dim=in_dim, n_iters=n_iters, n_layers=n_layers))
            out_pool_dim += in_dim
        elif isinstance(this_pool, str) and (this_pool[:3] == "dir"):
            dir_idx = int(this_pool[3:])
            pool_layer.append(DirPooling(dir_idx=dir_idx))
        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 VirtualNode(nn.Module):
    def __init__(
        self,
        dim: 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,
    ):
        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

            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

        """
        super().__init__()
        if (vn_type is None) or (vn_type.lower() == "none"):
            self.vn_type = None
            self.fc_layer = None
            self.residual = None
            return

        self.vn_type = vn_type.lower()
        self.residual = residual
        self.fc_layer = FCLayer(
            in_dim=dim,
            out_dim=dim,
            activation=activation,
            dropout=dropout,
            normalization=normalization,
            bias=bias,
        )

    def forward(
        self, g: dgl.DGLGraph, h: torch.Tensor, vn_h: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        r"""
        Apply the virtual node layer.

        Parameters:

            g:
                graph on which the convolution is done

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

            vn_h (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

        Returns:

            `h = 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_h = 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

        """

        g.ndata["h"] = h

        # Pool the features
        if self.vn_type is None:
            return h, vn_h
        elif self.vn_type == "mean":
            pool = mean_nodes(g, "h")
        elif self.vn_type == "max":
            pool = max_nodes(g, "h")
        elif self.vn_type == "sum":
            pool = sum_nodes(g, "h")
        elif self.vn_type == "logsum":
            pool = mean_nodes(g, "h")
            lognum = torch.log(torch.tensor(g.batch_num_nodes(), dtype=h.dtype, device=h.device))
            pool = pool * lognum.unsqueeze(-1)
        else:
            raise ValueError(
                f'Undefined input "{self.pooling}". Accepted values are "none", "sum", "mean", "logsum"'
            )

        # Compute the new virtual node features
        vn_h_temp = self.fc_layer.forward(vn_h + pool)
        if self.residual:
            vn_h = vn_h + vn_h_temp
        else:
            vn_h = vn_h_temp

        # Add the virtual node value to the graph features
        temp_h = torch.cat(
            [vn_h[ii : ii + 1].repeat(num_nodes, 1) for ii, num_nodes in enumerate(g.batch_num_nodes())],
            dim=0,
        )
        h = h + temp_h

        return h, vn_h


# expts/main_run.py

# General imports
import os
from os.path import dirname, abspath
import yaml
from copy import deepcopy
from omegaconf import DictConfig


# Current project imports
import goli
from goli.config._loader import load_datamodule, load_metrics, load_architecture, load_predictor, load_trainer


# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(goli.__file__)))
# CONFIG_FILE = "expts/config_molPCBA.yaml"
# CONFIG_FILE = "expts/config_molHIV.yaml"
# CONFIG_FILE = "expts/config_molPCQM4M.yaml"
# CONFIG_FILE = "expts/config_molPCQM4Mv2.yaml"
CONFIG_FILE = "expts/config_micro_ZINC.yaml"
# CONFIG_FILE = "expts/config_micro-PCBA.yaml"
# CONFIG_FILE = "expts/config_ZINC_bench_gnn.yaml"
os.chdir(MAIN_DIR)


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

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

    # 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)

    print(predictor.model)
    print(predictor.summarize(max_depth=4))

    trainer = load_trainer(cfg)

    # Run the model training
    print("\n------------ TRAINING STARTED ------------")
    try:
        trainer.fit(model=predictor, datamodule=datamodule)
        print("\n------------ TRAINING COMPLETED ------------\n\n")

    except Exception as e:
        if not cfg["constants"]["raise_train_error"]:
            print("\n------------ TRAINING ERROR: ------------\n\n", e)
        else:
            raise e

    print("\n------------ TESTING STARTED ------------")
    try:
        ckpt_path = trainer.checkpoint_callbacks[0].best_model_path
        trainer.test(model=predictor, datamodule=datamodule, ckpt_path=ckpt_path)
        print("\n------------ TESTING COMPLETED ------------\n\n")

    except Exception as e:
        if not cfg["constants"]["raise_train_error"]:
            print("\n------------ TESTING ERROR: ------------\n\n", e)
        else:
            raise e


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_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

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


# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(goli.__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://goli-public/pretrained-models"
    MODEL_NAME = "goli-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(predictor.summarize(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_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

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


# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(goli.__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(predictor.summarize(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

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

from goli.trainer.predictor import PredictorModule


# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(goli.__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(predictor.summarize(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)


# profiling/profile_forward.py

from time import time
import torch
import yaml
import fsspec
import dgl
from torch.profiler import profile, record_function, ProfilerActivity

from goli.config._loader import load_datamodule, load_metrics, load_predictor, load_architecture


def timed_fulldgl_forward(self, g: dgl.DGLGraph, flip_pos_enc: str) -> torch.Tensor:
    # Get the node features and positional embedding

    t0 = time()
    h = g.ndata["feat"]
    if "pos_enc_feats_sign_flip" in g.ndata.keys():
        pos_enc = g.ndata["pos_enc_feats_sign_flip"]
        if flip_pos_enc == "random":
            rand_sign_shape = ([1] * (pos_enc.ndim - 1)) + [pos_enc.shape[-1]]
            rand_sign = torch.sign(torch.randn(rand_sign_shape, dtype=h.dtype, device=h.device))
            pos_enc = pos_enc * rand_sign
        elif flip_pos_enc == "no-flip":
            pass
        elif flip_pos_enc == "sign-flip":
            pos_enc = -pos_enc
        h = torch.cat((h, pos_enc), dim=-1)
    if "pos_enc_feats_no_flip" in g.ndata.keys():
        pos_enc = g.ndata["pos_enc_feats_no_flip"]
        h = torch.cat((h, pos_enc), dim=-1)

    g.ndata["h"] = h
    print("pos encoding: ", time() - t0)

    if "feat" in g.edata.keys():
        g.edata["e"] = g.edata["feat"]

    t0 = time()
    # Run the pre-processing network on node features
    if self.pre_nn is not None:
        h = g.ndata["h"]
        h = self.pre_nn.forward(h)
        g.ndata["h"] = h
    print("node features: ", time() - t0)

    t0 = time()
    # 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.edata["e"]
        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.edata["e"] = e
    print("edge features: ", time() - t0)

    t0 = time()
    # Run the graph neural network
    h = self.gnn.forward(g)
    print("GNN: ", time() - t0)

    t0 = time()
    # Run the output network
    if self.post_nn is not None:
        if self.concat_last_layers is None:
            h = self.post_nn.forward(h)
        else:
            # Concatenate the output of the last layers according to `self._concat_last_layers``.
            # Useful for generating fingerprints
            h = [h]
            for ii in range(len(self.post_nn.layers)):
                h.insert(0, self.post_nn.layers[ii].forward(h[0]))  # Append in reverse order
            h = torch.cat([h[ii] for ii in self._concat_last_layers], dim=-1)
    print("output nn: ", time() - t0)

    return h


def timed_dgllayer_forward(self, layer, g, h, e, step_idx) -> torch.Tensor:
    # Apply the GNN layer with the right inputs/outputs
    t0 = time()
    if layer.layer_inputs_edges and layer.layer_outputs_edges:
        h, e = layer(g=g, h=h, e=e)
    elif layer.layer_inputs_edges:
        h = layer(g=g, h=h, e=e)
    elif layer.layer_outputs_edges:
        h, e = layer(g=g, h=h)
    else:
        h = layer(g=g, h=h)
    print("gnn layer: ", time() - t0)

    t0 = time()
    # Apply the residual layers on the features and edges (if applicable)
    if step_idx < len(self.layers) - 1:
        h, h_prev = self.residual_layer.forward(h, h_prev, step_idx=step_idx)
        if (self.residual_edges_layer is not None) and (layer.layer_outputs_edges):
            e, e_prev = self.residual_edges_layer.forward(e, e_prev, step_idx=step_idx)
    print("residual: ", time() - t0)

    return h, e, h_prev, e_prev


def main():
    CONFIG_PATH = "profiling/configs_profiling.yaml"
    DATA_PATH = "https://storage.googleapis.com/goli-public/datasets/goli-zinc-bench-gnn/smiles_score.csv.gz"
    BATCH_SIZE = 1000
    NUM_FORWARD_LOOPS = 6
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

    with fsspec.open(CONFIG_PATH, "r") as f:
        cfg = yaml.safe_load(f)
    cfg["datamodule"]["args"]
    cfg["datamodule"]["args"]["cache_data_path"] = "goli/data/cache/profiling/forward_data.cache"
    cfg["datamodule"]["args"]["df_path"] = DATA_PATH
    cfg["datamodule"]["args"]["sample_size"] = NUM_FORWARD_LOOPS * BATCH_SIZE
    cfg["datamodule"]["args"]["prepare_dict_or_graph"] = "dglgraph"

    datamodule = load_datamodule(cfg)
    datamodule.prepare_data()

    dglgraphs = [datamodule.dataset[ii]["features"] for ii in range(len(datamodule.dataset))]
    dglgraphs = [
        dgl.batch(dglgraphs[ii * (BATCH_SIZE + 1) : (ii + 1) * (BATCH_SIZE + 1)])
        for ii in range(NUM_FORWARD_LOOPS)
    ]
    for ii in range(NUM_FORWARD_LOOPS):
        dglgraphs[ii].ndata["feat"] = dglgraphs[ii].ndata["feat"].to(dtype=torch.float32)
        dglgraphs[ii].edata["feat"] = dglgraphs[ii].edata["feat"].to(dtype=torch.float32)
        dglgraphs[ii] = dglgraphs[ii].to(device=DEVICE)

    # 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)

    predictor = load_predictor(cfg, model_class, model_kwargs, metrics)
    model = predictor.model
    model = model.to(device=DEVICE)

    with torch.no_grad():
        for ii in range(NUM_FORWARD_LOOPS):
            print("---------------------------------")
            timed_fulldgl_forward(model, dglgraphs[ii], flip_pos_enc="random")
            print("---------------------------------")

    print("Done :)")


if __name__ == "__main__":
    main()


# profiling/profile_mol_to_graph.py

from tqdm import tqdm
import datamol as dm
import pickle

from goli.data.utils import load_micro_zinc
from goli.features.featurizer import mol_to_dglgraph, mol_to_adj_and_features


def main():
    df = load_micro_zinc()
    smiles = df["SMILES"].values.tolist()
    smiles = smiles * 1
    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,
        "pos_encoding_as_directions": pos_enc,
    }

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

    print(graphs[0])


if __name__ == "__main__":
    main()


# profiling/profile_one_of_k_encoding.py

from tqdm import tqdm
from goli.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 goli.config._loader import load_datamodule, load_metrics, load_trainer, load_predictor, load_architecture
from pytorch_lightning import Trainer


def main():
    CONFIG_PATH = "expts/config_micro-PCBA.yaml"
    # DATA_PATH = "https://storage.googleapis.com/goli-public/datasets/goli-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"] = "goli/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()


# 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 goli/nn/architectures...

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

import torch
import unittest as ut
import dgl
from copy import deepcopy

from goli.nn.architectures import FeedForwardNN, FeedForwardDGL, FullDGLNetwork, FullDGLSiameseNetwork
from goli.nn.base_layers import FCLayer
from goli.nn.residual_connections import (
    ResidualConnectionConcat,
    ResidualConnectionDenseNet,
    ResidualConnectionNone,
    ResidualConnectionSimple,
    ResidualConnectionWeighted,
)


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])

        h = torch.FloatTensor(batch, in_dim)
        h_out = lnn.forward(h)

        self.assertListEqual(list(h_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])

        h = torch.FloatTensor(batch, in_dim)
        h_out = lnn.forward(h)

        self.assertListEqual(list(h_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)

            h = torch.FloatTensor(batch, in_dim)
            h_out = lnn.forward(h)

            self.assertListEqual(list(h_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])

        h = torch.FloatTensor(batch, in_dim)
        h_out = lnn.forward(h)

        self.assertListEqual(list(h_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])

        h = torch.FloatTensor(batch, in_dim)
        h_out = lnn.forward(h)

        self.assertListEqual(list(h_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])

        h = torch.FloatTensor(batch, in_dim)
        h_out = lnn.forward(h)

        self.assertListEqual(list(h_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])

        h = torch.FloatTensor(batch, in_dim)
        h_out = lnn.forward(h)

        self.assertListEqual(list(h_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])

        h = torch.FloatTensor(batch, in_dim)
        h_out = lnn.forward(h)

        self.assertListEqual(list(h_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))

        h = torch.FloatTensor(batch, in_dim)
        h_out = lnn.forward(h)

        self.assertListEqual(list(h_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))

        h = torch.FloatTensor(batch, in_dim)
        h_out = lnn.forward(h)

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


class test_FeedForwardDGL(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]

    g1 = dgl.graph((torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])))
    g2 = dgl.graph((torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])))
    g1.ndata["h"] = torch.zeros(g1.num_nodes(), in_dim, dtype=torch.float32)
    g1.edata["e"] = torch.ones(g1.num_edges(), in_dim_edges, dtype=torch.float32)
    g2.ndata["h"] = torch.ones(g2.num_nodes(), in_dim, dtype=torch.float32)
    g2.edata["e"] = torch.zeros(g2.num_edges(), in_dim_edges, dtype=torch.float32)
    batch = [g1, g2, deepcopy(g1), deepcopy(g2)]
    batch = [dgl.add_self_loop(g) for g in batch]
    bg = dgl.batch(batch)

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

    layers_kwargs = {
        "gcn": {},
        "gin": {},
        "gat": {"layer_kwargs": {"num_heads": 3}},
        "gated-gcn": {"in_dim_edges": in_dim_edges, "hidden_dims_edges": hidden_dims},
        "pna-conv": {"layer_kwargs": pna_kwargs},
        "pna-msgpass#1": {"layer_kwargs": pna_kwargs, "in_dim_edges": 0},
        "pna-msgpass#2": {"layer_kwargs": pna_kwargs, "in_dim_edges": in_dim_edges},
    }

    def test_dgl_forward_no_residual(self):

        for pooling in [["none"], ["sum"], ["mean", "s2s", "max"]]:
            for residual_skip_steps in [1, 2, 3]:
                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"pooling={pooling}, virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                            layer_type = layer_name.split("#")[0]

                            gnn = FeedForwardDGL(
                                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,
                                pooling=pooling,
                                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)

                            bg = deepcopy(self.bg)
                            h_out = gnn.forward(bg)

                            dim_1 = bg.num_nodes() if pooling == ["none"] else bg.batch_size
                            self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)

    def test_dgl_forward_simple_residual(self):

        for pooling in [["none"], ["sum"], ["mean", "s2s", "max"]]:
            for residual_skip_steps in [1, 2, 3]:
                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"pooling={pooling}, virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                            layer_type = layer_name.split("#")[0]

                            gnn = FeedForwardDGL(
                                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,
                                pooling=pooling,
                                **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)

                            bg = deepcopy(self.bg)
                            h_out = gnn.forward(bg)

                            dim_1 = bg.num_nodes() if pooling == ["none"] else bg.batch_size
                            self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)

    def test_dgl_forward_weighted_residual(self):

        for pooling in [["none"], ["sum"], ["mean", "s2s", "max"]]:
            for residual_skip_steps in [1, 2, 3]:
                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"pooling={pooling}, virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                            layer_type = layer_name.split("#")[0]

                            gnn = FeedForwardDGL(
                                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,
                                pooling=pooling,
                                **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)

                            bg = deepcopy(self.bg)
                            h_out = gnn.forward(bg)

                            dim_1 = bg.num_nodes() if pooling == ["none"] else bg.batch_size
                            self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)

    def test_dgl_forward_concat_residual(self):

        for pooling in [["none"], ["sum"], ["mean", "s2s", "max"]]:
            for residual_skip_steps in [1, 2, 3]:
                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"pooling={pooling}, virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                            layer_type = layer_name.split("#")[0]

                            gnn = FeedForwardDGL(
                                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,
                                pooling=pooling,
                                **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)

                            bg = deepcopy(self.bg)
                            h_out = gnn.forward(bg)

                            dim_1 = bg.num_nodes() if pooling == ["none"] else bg.batch_size
                            self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)

    def test_dgl_forward_densenet_residual(self):

        for pooling in [["none"], ["sum"], ["mean", "s2s", "max"]]:
            for residual_skip_steps in [1, 2, 3]:
                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"pooling={pooling}, virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                            layer_type = layer_name.split("#")[0]

                            gnn = FeedForwardDGL(
                                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,
                                pooling=pooling,
                                **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)

                            bg = deepcopy(self.bg)
                            h_out = gnn.forward(bg)

                            dim_1 = bg.num_nodes() if pooling == ["none"] else bg.batch_size
                            self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)


class test_FullDGLNetwork(ut.TestCase):

    kwargs = {
        "activation": "relu",
        "last_activation": "none",
        "normalization": "none",
        "dropout": 0.2,
        "name": "LNN",
    }

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

    g1 = dgl.graph((torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])))
    g2 = dgl.graph((torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])))
    g1.ndata["feat"] = torch.zeros(g1.num_nodes(), in_dim, dtype=torch.float32)
    g1.edata["feat"] = torch.ones(g1.num_edges(), in_dim_edges, dtype=torch.float32)
    g2.ndata["feat"] = torch.ones(g2.num_nodes(), in_dim, dtype=torch.float32)
    g2.edata["feat"] = torch.zeros(g2.num_edges(), in_dim_edges, dtype=torch.float32)
    batch = [g1, g2, deepcopy(g1), deepcopy(g2)]
    batch = [dgl.add_self_loop(g) for g in batch]
    bg = dgl.batch(batch)

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

    gnn_layers_kwargs = {
        "gcn": {},
        "gin": {},
        "gat": {"layer_kwargs": {"num_heads": 3}},
        "gated-gcn": {"in_dim_edges": in_dim_edges, "hidden_dims_edges": hidden_dims},
        "pna-conv": {"layer_kwargs": pna_kwargs},
        "pna-msgpass#1": {"layer_kwargs": pna_kwargs, "in_dim_edges": 0},
        "pna-msgpass#2": {"layer_kwargs": pna_kwargs, "in_dim_edges": in_dim_edges},
    }

    def test_full_network_densenet(self):

        temp_dim_1 = 5
        temp_dim_2 = 17

        pre_nn_kwargs = dict(in_dim=self.in_dim, out_dim=temp_dim_1, hidden_dims=[4, 4, 4, 4, 4])

        post_nn_kwargs = dict(in_dim=temp_dim_2, out_dim=self.out_dim, hidden_dims=[3, 3, 3, 3])

        for pooling in [["none"], ["sum"], ["mean", "s2s", "max"]]:
            for residual_skip_steps in [1, 2, 3]:
                for virtual_node in self.virtual_nodes:
                    for normalization in self.norms:
                        for layer_name, this_kwargs in self.gnn_layers_kwargs.items():
                            err_msg = f"pooling={pooling}, virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                            layer_type = layer_name.split("#")[0]

                            gnn_kwargs = dict(
                                in_dim=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,
                                pooling=pooling,
                                **this_kwargs,
                                **self.kwargs,
                            )

                            net = FullDGLNetwork(
                                pre_nn_kwargs=pre_nn_kwargs,
                                gnn_kwargs=gnn_kwargs,
                                post_nn_kwargs=post_nn_kwargs,
                            )
                            bg = deepcopy(self.bg)
                            h_out = net.forward(bg)

                            dim_1 = bg.num_nodes() if pooling == ["none"] else bg.batch_size
                            self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)


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


# tests/test_data_utils.py

import pandas as pd

import goli


def test_list_datasets():
    datasets = goli.data.utils.list_goli_datasets()
    assert isinstance(datasets, set)
    assert len(datasets) > 0


def test_download_datasets(tmpdir):
    dataset_dir = tmpdir.mkdir("goli-datasets")

    data_path = goli.data.utils.download_goli_dataset("goli-zinc-micro", output_path=dataset_dir)

    fpath = goli.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


# tests/test_datamodule.py

import pathlib
import tempfile

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

import goli


class Test_DataModule(ut.TestCase):
    def test_dglfromsmiles_dm(self):

        df = goli.data.load_tiny_zinc()
        # 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"] = ["in-ring", "bond-type-onehot"]
        featurization_args["add_self_loop"] = False
        featurization_args["use_bonds_weights"] = False
        featurization_args["explicit_H"] = False

        # Config for datamodule
        dm_args = {}
        dm_args["df"] = df
        dm_args["cache_data_path"] = None
        dm_args["featurization"] = featurization_args
        dm_args["smiles_col"] = "SMILES"
        dm_args["label_cols"] = ["SA"]
        dm_args["split_val"] = 0.2
        dm_args["split_test"] = 0.2
        dm_args["split_seed"] = 19
        dm_args["batch_size_train_val"] = 16
        dm_args["batch_size_test"] = 16
        dm_args["num_workers"] = 0
        dm_args["pin_memory"] = True
        dm_args["featurization_n_jobs"] = 16
        dm_args["featurization_progress"] = True
        dm_args["featurization_backend"] = "multiprocessing"

        dm = goli.data.DGLFromSmilesDataModule(**dm_args)

        assert dm.num_node_feats == 50
        assert dm.num_edge_feats == 6

        dm.prepare_data()
        dm.setup()

        assert len(dm) == 100
        assert len(dm.train_ds) == 60  # type: ignore
        assert len(dm.val_ds) == 20  # type: ignore
        assert 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", "smiles"}
            assert batch["labels"].shape == (16, 1)

    def test_dglfromsmiles_from_config(self):

        config = goli.load_config(name="zinc_default_fulldgl")
        df = goli.data.load_tiny_zinc()

        dm_args = dict(config.data.args)
        dm_args["df"] = df

        dm = goli.data.DGLFromSmilesDataModule(**dm_args)

        assert dm.num_node_feats == 50
        assert dm.num_edge_feats == 6

        dm.prepare_data()
        dm.setup()

        assert len(dm.train_ds) == 60  # type: ignore
        assert len(dm.val_ds) == 20  # type: ignore
        assert 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", "smiles"}
            assert batch["labels"].shape == (16, 1)

    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["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
        dm_args = {}
        dm_args["dataset_name"] = dataset_name
        dm_args["cache_data_path"] = None
        dm_args["featurization"] = featurization_args
        dm_args["batch_size_train_val"] = 16
        dm_args["batch_size_test"] = 16
        dm_args["num_workers"] = 0
        dm_args["pin_memory"] = True
        dm_args["featurization_n_jobs"] = 16
        dm_args["featurization_progress"] = True
        dm_args["featurization_backend"] = "loky"

        ds = goli.data.DGLOGBDataModule(**dm_args)

        ds.prepare_data()
        ds.setup()

        # test module
        assert ds.num_edge_feats == 5
        assert ds.num_node_feats == 50
        assert len(ds) == 642
        assert ds.dataset_name == "ogbg-molfreesolv"

        # test dataset
        assert set(ds.train_ds[0].keys()) == {"smiles", "indices", "features", "labels"}

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

    def test_datamodule_cache_invalidation(self):

        df = goli.data.load_tiny_zinc()

        cache_data_path = pathlib.Path(tempfile.mkdtemp()) / "cache.pkl"

        # 1. Build a module with specific feature arguments

        featurization_args = {}
        featurization_args["atom_property_list_float"] = ["mass", "electronegativity", "in-ring"]
        featurization_args["edge_property_list"] = ["bond-type-onehot", "stereo", "in-ring"]

        dm_args = {}
        dm_args["df"] = df
        dm_args["cache_data_path"] = cache_data_path
        dm_args["featurization"] = featurization_args
        datam = goli.data.DGLFromSmilesDataModule(**dm_args)
        datam.prepare_data()
        datam.setup()

        assert datam.num_node_feats == 3
        assert datam.num_edge_feats == 13

        # 2. Reload with the same arguments should not trigger a new preparation and give
        # the same feature's dimensions.

        datam = goli.data.DGLFromSmilesDataModule(**dm_args)
        datam.prepare_data()
        datam.setup()

        assert datam.num_node_feats == 3
        assert datam.num_edge_feats == 13

        # 3. Reloading from the same cache file should trigger a new data preparation and
        # so different feature's dimensions.

        featurization_args = {}
        featurization_args["edge_property_list"] = ["stereo", "in-ring"]
        featurization_args["atom_property_list_float"] = ["mass", "electronegativity"]

        dm_args = {}
        dm_args["df"] = df
        dm_args["cache_data_path"] = cache_data_path
        dm_args["featurization"] = featurization_args
        datam = goli.data.DGLFromSmilesDataModule(**dm_args)
        datam.prepare_data()
        datam.setup()

        assert datam.num_node_feats == 2
        assert datam.num_edge_feats == 8

    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,
            ) = goli.data.DGLFromSmilesDataModule._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)


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


# tests/test_dict_tensor.py

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

from typing import Dict
from goli.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 goli/features/featurizer.py
"""

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

from goli.features.featurizer import (
    get_mol_atomic_features_onehot,
    get_mol_atomic_features_float,
    get_mol_edge_features,
    mol_to_adj_and_features,
    mol_to_dglgraph_dict,
    mol_to_dglgraph,
)


class test_featurizer(ut.TestCase):

    smiles = [
        "C",
        "CC",
        "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",
    ]

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

    atomic_onehot_props = [
        "atomic-number",
        "valence",
        "degree",
        "implicit-valence",
        "hybridization",
        "chirality",
    ]

    atomic_float_props = [
        "atomic-number",
        "mass",
        "valence",
        "implicit-valence",
        "hybridization",
        "chirality",
        "aromatic",
        "in-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",
    ]

    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_dglgraph(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}"

                    graph = mol_to_dglgraph(
                        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(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.ndata["feat"]
                        edata = graph.edata["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_gnn_layers.py

"""
Unit tests for the different layers of goli/dgl/dgl_layers/...

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

import numpy as np
import torch
import unittest as ut
import dgl
from copy import deepcopy

from goli.nn.dgl_layers import (
    GATLayer,
    GCNLayer,
    GINLayer,
    GatedGCNLayer,
    PNAConvolutionalLayer,
    PNAMessagePassingLayer,
    DGNConvolutionalLayer,
    DGNMessagePassingLayer,
)


class test_DGL_Layers(ut.TestCase):

    in_dim = 21
    out_dim = 11
    in_dim_edges = 13
    out_dim_edges = 17

    g1 = dgl.graph((torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])))
    g2 = dgl.graph((torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])))
    g1.ndata["h"] = torch.zeros(g1.num_nodes(), in_dim, dtype=float)
    g1.edata["e"] = torch.ones(g1.num_edges(), in_dim_edges, dtype=float)
    g2.ndata["h"] = torch.ones(g2.num_nodes(), in_dim, dtype=float)
    g2.edata["e"] = torch.zeros(g2.num_edges(), in_dim_edges, dtype=float)
    bg = dgl.batch([g1, g2])
    bg = dgl.add_self_loop(bg)
    bg.ndata["pos_dir"] = torch.randn_like(bg.ndata["h"])

    kwargs = {
        "activation": "relu",
        "dropout": 0.1,
        "normalization": "batch_norm",
    }

    def test_gcnlayer(self):

        bg = deepcopy(self.bg)
        h_in = bg.ndata["h"]
        layer = GCNLayer(in_dim=self.in_dim, out_dim=self.out_dim, **self.kwargs).to(float)

        # 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
        h = layer.forward(g=bg, h=h_in)
        self.assertEqual(h.shape[0], h_in.shape[0])
        self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)

    def test_ginlayer(self):

        bg = deepcopy(self.bg)
        h_in = bg.ndata["h"]
        layer = GINLayer(in_dim=self.in_dim, out_dim=self.out_dim, **self.kwargs).to(float)

        # 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
        h = layer.forward(g=bg, h=h_in)
        self.assertEqual(h.shape[0], h_in.shape[0])
        self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)

    def test_gatlayer(self):

        num_heads = 3
        bg = deepcopy(self.bg)
        h_in = bg.ndata["h"]
        layer = GATLayer(in_dim=self.in_dim, out_dim=self.out_dim, num_heads=num_heads, **self.kwargs).to(
            float
        )

        # 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, num_heads)

        # Apply layer
        h = layer.forward(g=bg, h=h_in)
        self.assertEqual(h.shape[0], h_in.shape[0])
        self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)

    def test_gatedgcnlayer(self):

        bg = deepcopy(self.bg)
        h_in = bg.ndata["h"]
        e_in = bg.edata["e"]
        layer = GatedGCNLayer(
            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,
        ).to(float)

        # 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 without edges
        with self.assertRaises(TypeError):
            h = layer.forward(g=bg, h=h_in)

        # Apply layer with edges
        h, e = layer.forward(g=bg, h=h_in, e=e_in)
        self.assertEqual(h.shape[0], h_in.shape[0])
        self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)

    def test_pnaconvolutionallayer(self):

        bg = deepcopy(self.bg)
        h_in = bg.ndata["h"]
        e_in = bg.edata["e"]
        aggregators = ["mean", "max", "min", "lap", "std", "moment3", "moment4", "sum"]
        scalers = ["identity", "amplification", "attenuation"]

        layer = PNAConvolutionalLayer(
            in_dim=self.in_dim, out_dim=self.out_dim, aggregators=aggregators, scalers=scalers, **self.kwargs
        ).to(float)

        # Check the re-implementation of abstract methods
        self.assertTrue(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
        h = layer.forward(g=bg, h=h_in)
        self.assertEqual(h.shape[0], h_in.shape[0])
        self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)

        # Now try with edges
        layer = PNAConvolutionalLayer(
            in_dim=self.in_dim,
            out_dim=self.out_dim,
            aggregators=aggregators,
            scalers=scalers,
            in_dim_edges=self.in_dim_edges,
            **self.kwargs,
        ).to(float)

        # 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)

        # Apply layer
        with self.assertRaises(Exception):
            h = layer.forward(g=bg, h=h_in)

        h = layer.forward(g=bg, h=h_in, e=e_in)
        self.assertEqual(h.shape[0], h_in.shape[0])
        self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)

    def test_pnamessagepassinglayer(self):

        bg = deepcopy(self.bg)
        h_in = bg.ndata["h"]
        e_in = bg.edata["e"]
        aggregators = ["mean", "max", "min", "lap", "std", "moment3", "moment4", "sum"]
        scalers = ["identity", "amplification", "attenuation"]

        layer = PNAMessagePassingLayer(
            in_dim=self.in_dim, out_dim=self.out_dim, aggregators=aggregators, scalers=scalers, **self.kwargs
        ).to(float)

        # 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
        h = layer.forward(g=bg, h=h_in)
        self.assertEqual(h.shape[0], h_in.shape[0])
        self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)

        # Now try with edges
        layer = PNAMessagePassingLayer(
            in_dim=self.in_dim,
            out_dim=self.out_dim,
            aggregators=aggregators,
            scalers=scalers,
            in_dim_edges=self.in_dim_edges,
            **self.kwargs,
        ).to(float)

        # 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)

        # Apply layer
        with self.assertRaises(Exception):
            h = layer.forward(g=bg, h=h_in)

        h = layer.forward(g=bg, h=h_in, e=e_in)
        self.assertEqual(h.shape[0], h_in.shape[0])
        self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)

    def test_dgnconvolutionallayer(self):

        bg = deepcopy(self.bg)
        h_in = bg.ndata["h"]
        e_in = bg.edata["e"]
        aggregators = [
            "mean",
            "max",
            "min",
            "lap",
            "std",
            "moment3",
            "moment4",
            "sum",
            "dir0/dx_abs",
            "dir1/dx_abs",
            "dir2/dx_no_abs",
            "dir1/smooth",
            "dir1/forward",
            "dir1/backward/0.5",
            "dir4/dx_abs/5",
        ]
        scalers = ["identity", "amplification", "attenuation"]

        layer = DGNConvolutionalLayer(
            in_dim=self.in_dim, out_dim=self.out_dim, aggregators=aggregators, scalers=scalers, **self.kwargs
        ).to(float)

        # Check the re-implementation of abstract methods
        self.assertTrue(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
        h = layer.forward(g=bg, h=h_in)
        self.assertEqual(h.shape[0], h_in.shape[0])
        self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)

        # Now try with edges
        layer = DGNConvolutionalLayer(
            in_dim=self.in_dim,
            out_dim=self.out_dim,
            aggregators=aggregators,
            scalers=scalers,
            in_dim_edges=self.in_dim_edges,
            **self.kwargs,
        ).to(float)

        # 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)

        # Apply layer
        with self.assertRaises(Exception):
            h = layer.forward(g=bg, h=h_in)

        h = layer.forward(g=bg, h=h_in, e=e_in)
        self.assertEqual(h.shape[0], h_in.shape[0])
        self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)

    def test_dgnmessagepassinglayer(self):

        bg = deepcopy(self.bg)
        h_in = bg.ndata["h"]
        e_in = bg.edata["e"]
        aggregators = [
            "mean",
            "max",
            "min",
            "lap",
            "std",
            "moment3",
            "moment4",
            "sum",
            "dir0/dx_abs",
            "dir1/dx_abs",
            "dir2/dx_no_abs",
            "dir1/smooth",
            "dir1/forward",
            "dir1/backward/0.5",
            "dir4/dx_abs/5",
        ]
        scalers = ["identity", "amplification", "attenuation"]

        layer = DGNMessagePassingLayer(
            in_dim=self.in_dim, out_dim=self.out_dim, aggregators=aggregators, scalers=scalers, **self.kwargs
        ).to(float)

        # 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
        h = layer.forward(g=bg, h=h_in)
        self.assertEqual(h.shape[0], h_in.shape[0])
        self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)

        # Now try with edges
        layer = DGNMessagePassingLayer(
            in_dim=self.in_dim,
            out_dim=self.out_dim,
            aggregators=aggregators,
            scalers=scalers,
            in_dim_edges=self.in_dim_edges,
            **self.kwargs,
        ).to(float)

        # 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)

        # Apply layer
        with self.assertRaises(Exception):
            h = layer.forward(g=bg, h=h_in)

        h = layer.forward(g=bg, h=h_in, e=e_in)
        self.assertEqual(h.shape[0], h_in.shape[0])
        self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)


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


# tests/test_metrics.py

"""
Unit tests for the metrics and wrappers of goli/trainer/metrics/...
"""

import torch
import unittest as ut
import tempfile
import os
import operator as op

from goli.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-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-mean-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-flatten", "ignore-mean-label"]
        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", "target_to_int": True},
        ]

        for metric in metrics:
            for target_nan_mask in target_nan_masks:
                for kwargs in other_kwargs:
                    for threshold_kwargs in thresholds:
                        err_msg = f"{metric} - {target_nan_mask} - {kwargs} - {threshold_kwargs}"
                        metric_wrapper = MetricWrapper(
                            metric=metric,
                            threshold_kwargs=threshold_kwargs,
                            target_nan_mask=target_nan_mask,
                            **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)


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


# tests/test_multitask_datamodule.py

import pathlib
import tempfile

import unittest as ut
import numpy as np
from omegaconf import OmegaConf
import torch
import pandas as pd

import goli
from goli.data import load_micro_zinc, DGLDataset, SingleTaskDataset, MultitaskDGLDataset
from goli.data.datamodule import smiles_to_unique_mol_ids

class Test_Multitask_DataModule(ut.TestCase):
    
    def test_multitask_dglfromsmiles_dm(self):
        """Cover similar testing as for the original data module."""
        df = goli.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["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["split_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["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["split_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["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["split_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"] = "multiprocessing"
        dm_args["num_workers"] = 0
        dm_args["pin_memory"] = True
        dm_args["cache_data_path"] = None
        dm_args["batch_size_train_val"] = 16
        dm_args["batch_size_test"] = 16

        # Create the data module
        dm = goli.data.MultitaskDGLFromSmilesDataModule(**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", "smiles", "mol_ids"}
            assert len(batch["mol_ids"]) == 16
            #assert len(batch["smiles"]) == 16
            #assert len(batch["features"])                      # This is not a list, but a graph.

            #assert batch["labels"].shape == (16, 1)            # Single-task case
            assert batch["labels"]["SA"].shape == (16, 1)
            assert batch["labels"]["logp"].shape == (16, 1)
            assert batch["labels"]["score"].shape == (16, 1)

    def test_multitask_dglfromsmiles_from_config(self):

        config = goli.load_config(name="zinc_default_multitask_fulldgl")

        df = goli.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"]["df_path"] = None
        dm_args["task_specific_args"]["logp"]["df_path"] = None
        dm_args["task_specific_args"]["score"]["df_path"] = None

        dm = goli.data.MultitaskDGLFromSmilesDataModule(**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", "smiles", "mol_ids"}
            assert len(batch["mol_ids"]) == 16
            #assert len(batch["smiles"]) == 16
            #assert len(batch["features"])                      # This is not a list, but a graph.

            #assert batch["labels"].shape == (16, 1)            # Single-task case
            assert batch["labels"]["SA"].shape == (16, 1)
            assert batch["labels"]["logp"].shape == (16, 1)
            assert batch["labels"]["score"].shape == (16, 1)

    def test_multitask_dglfromsmiles_from_config_csv(self):
        config = goli.load_config(name="zinc_default_multitask_fulldgl")
     
        dm_args = OmegaConf.to_container(config.datamodule.args, resolve=True)
        dm = goli.data.MultitaskDGLFromSmilesDataModule(**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", "smiles", "mol_ids"}
            assert len(batch["mol_ids"]) == 16
            #assert len(batch["smiles"]) == 16
            #assert len(batch["features"])                      # This is not a list, but a graph.

            #assert batch["labels"].shape == (16, 1)            # Single-task case
            assert batch["labels"]["SA"].shape == (16, 1)
            assert batch["labels"]["logp"].shape == (16, 1)
            assert batch["labels"]["score"].shape == (16, 1)

class Test_Multitask_Dataset(ut.TestCase):

    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'], 
            labels=df_micro_zinc_SA.loc[:,'SA']
            )
        ds_micro_zinc_logp = SingleTaskDataset(
            smiles=df_micro_zinc_logp.loc[:,'SMILES'], 
            labels=df_micro_zinc_logp.loc[:,'logp']
            )
        ds_micro_zinc_score = SingleTaskDataset(
            smiles=df_micro_zinc_score.loc[:,'SMILES'], 
            labels=df_micro_zinc_score.loc[:,'score']
            )
        
        # Create the multitask dataset
        datasets_dict = {'SA': ds_micro_zinc_SA, 'logp': ds_micro_zinc_logp, 'score': ds_micro_zinc_score}
        multitask_microzinc = MultitaskDGLDataset(datasets_dict) # 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'], 
    #         labels=df_micro_zinc_SA.loc[:,'SA']
    #         )
    #     ds_micro_zinc_logp = SingleTaskDataset(
    #         smiles=df_micro_zinc_logp.loc[:,'SMILES'], 
    #         labels=df_micro_zinc_logp.loc[:,'logp']
    #         )
    #     ds_micro_zinc_score = SingleTaskDataset(
    #         smiles=df_micro_zinc_score.loc[:,'SMILES'], 
    #         labels=df_micro_zinc_score.loc[:,'score']
    #         )

    #     # Create the multitask dataset
    #     datasets_dict = {'SA': ds_micro_zinc_SA, 'logp': ds_micro_zinc_logp, 'score': ds_micro_zinc_score}
    #     multitask_microzinc = MultitaskDGLDataset(datasets_dict) # 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

    #         if task == 'SA':
    #             self.assertEqual(label_SA, multitask_microzinc.labels[found_idx]['SA'])
    #             self.assertFalse('score' in multitask_microzinc.labels[found_idx].keys())
    #             self.assertFalse('logp' in multitask_microzinc.labels[found_idx].keys())
    #         elif task == 'logp':
    #             self.assertEqual(label_logp, multitask_microzinc.labels[found_idx]['logp'])
    #             self.assertFalse('score' in multitask_microzinc.labels[found_idx].keys())
    #             self.assertFalse('SA' in multitask_microzinc.labels[found_idx].keys())
    #         elif task == 'score':
    #             self.assertEqual(label_score, multitask_microzinc.labels[found_idx]['score'])
    #             self.assertFalse('SA' in multitask_microzinc.labels[found_idx].keys())
    #             self.assertFalse('logp' in multitask_microzinc.labels[found_idx].keys())

    # TODO (Gabriela): Fix this test case.
    # 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]         # 200 data points in 'SA' task, but 170-199 overlap with 'logp' task
    #     df_rows_logp = df.iloc[1:4]     # 200 data points in 'logp' task, but 170-199 overlap with 'SA' task, and 370-399 overlap with 'score' task
    #     df_rows_score = df.iloc[3:5]    # 350 data points in 'score' task, but 370-399 overlap with 'logp'
    #     total_data_points = 3             # There are 750 rows, but 60 smiles overlap, giving 690 unique molecules

    #     # 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'], 
    #         labels=df_micro_zinc_SA.loc[:,'SA']
    #         )
    #     ds_micro_zinc_logp = SingleTaskDataset(
    #         smiles=df_micro_zinc_logp.loc[:,'SMILES'], 
    #         labels=df_micro_zinc_logp.loc[:,'logp']
    #         )
    #     ds_micro_zinc_score = SingleTaskDataset(
    #         smiles=df_micro_zinc_score.loc[:,'SMILES'], 
    #         labels=df_micro_zinc_score.loc[:,'score']
    #         )
   
    #     # Create the multitask dataset
    #     datasets_dict = {'SA': ds_micro_zinc_SA, 'logp': ds_micro_zinc_logp, 'score': ds_micro_zinc_score}
    #     multitask_microzinc = MultitaskDGLDataset(datasets_dict) # Can optionally have features

    #     multitask_microzinc.print_this()
    #     pprint(df)

    #     # 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__())

    # # TODO (Gabriela): After fixing case 3, implement case with bad smiles.
    # def test_multitask_dataset_case_4(self):
    #     """Case: Different tasks, semi-intersection, some bad smiles
    #         - For each task, add a few random smiles that won’t work, such as “XYZ” or simply “X”
    #         - Check that the smiles are filtered, alongside their label (ensure the right label has been filtered)
    #     """
    #     pass

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

# tests/test_positional_encoding.py

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

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

from goli.features.featurizer import (
    mol_to_adj_and_features,
    mol_to_dglgraph_dict,
)
from goli.features.positional_encoding import graph_positional_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(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}"
                    pos_enc_sign_flip, pos_enc_no_flip = graph_positional_encoder(
                        adj, pos_type="laplacian_eigvec", num_pos=num_pos, disconnected_comp=disconnected_comp
                    )
                    self.assertEqual(list(pos_enc_sign_flip.shape), [adj.shape[0], num_pos], msg=err_msg)
                    self.assertIsNone(pos_enc_no_flip)

                    # Compute eigvals and eigvecs
                    lap = np.diag(np.sum(adj, axis=1)) - adj
                    eigvals, eigvecs = np.linalg.eig(lap)
                    sort_idx = np.argsort(eigvals)
                    eigvals, eigvecs = eigvals[sort_idx], eigvecs[:, sort_idx]
                    eigvecs = eigvecs / (np.sum(eigvecs**2, axis=0, keepdims=True) + 1e-8)

                    true_num_pos = min(num_pos, len(eigvals))
                    eigvals, eigvecs = eigvals[:true_num_pos], eigvecs[:, :true_num_pos]
                    eigvecs = np.sign(eigvecs[0:1, :]) * eigvecs
                    pos_enc_sign_flip = np.sign(pos_enc_sign_flip[0:1, :]) * pos_enc_sign_flip

                    # Compare the positional encoding
                    if disconnected_comp and ("." in self.smiles[ii]):
                        self.assertGreater(np.max(np.abs(eigvecs - pos_enc_sign_flip)), 1e-3)
                    elif not ("." in self.smiles[ii]):
                        np.testing.assert_array_almost_equal(
                            eigvecs, pos_enc_sign_flip[:, :true_num_pos], decimal=6, err_msg=err_msg
                        )

    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}"
                    pos_enc_sign_flip, pos_enc_no_flip = graph_positional_encoder(
                        adj,
                        pos_type="laplacian_eigvec_eigval",
                        num_pos=num_pos,
                        disconnected_comp=disconnected_comp,
                    )
                    self.assertEqual(list(pos_enc_sign_flip.shape), [adj.shape[0], num_pos], msg=err_msg)
                    self.assertEqual(list(pos_enc_no_flip.shape), [adj.shape[0], num_pos], msg=err_msg)

                    # Compute eigvals and eigvecs
                    lap = np.diag(np.sum(adj, axis=1)) - adj
                    eigvals, eigvecs = np.linalg.eig(lap)
                    sort_idx = np.argsort(eigvals)
                    eigvals, eigvecs = eigvals[sort_idx], eigvecs[:, sort_idx]
                    eigvecs = eigvecs / (np.sum(eigvecs**2, axis=0, keepdims=True) + 1e-8)

                    true_num_pos = min(num_pos, len(eigvals))
                    eigvals, eigvecs = eigvals[:true_num_pos], eigvecs[:, :true_num_pos]
                    eigvecs = np.sign(eigvecs[0:1, :]) * eigvecs
                    pos_enc_sign_flip = np.sign(pos_enc_sign_flip[0:1, :]) * pos_enc_sign_flip

                    if not ("." in self.smiles[ii]):
                        np.testing.assert_array_almost_equal(
                            eigvecs, pos_enc_sign_flip[:, :true_num_pos], decimal=6, err_msg=err_msg
                        )
                        np.testing.assert_array_almost_equal(
                            eigvals, pos_enc_no_flip[0, :true_num_pos], decimal=6, err_msg=err_msg
                        )


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


# tests/test_predictor.py

"""
Unit tests for the file goli/trainer/predictor.py
"""

import torch
from torch.nn import BCELoss, MSELoss
import unittest as ut

from goli.trainer.predictor import PredictorModule


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 = PredictorModule.parse_loss_fun(this_loss)
            loss = loss_fun(preds, target)


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


# tests/test_residual_connections.py

"""
Unit tests for the file goli/dgl/residual_connections.py
"""

import numpy as np
import torch
import unittest as ut

from goli.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):
            h, h_prev = rc.forward(h_original[ii], h_prev, step_idx=ii)
            np.testing.assert_array_equal(h.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):
                h, 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(
                    h.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)
                h, 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
                h, 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(
                    h.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):
                h, 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(
                    h.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):
                h, 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(
                    h.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 goli/utils/...
"""

from goli.utils.tensor import nan_mad, nan_mean, nan_std, nan_var, nan_median
import torch
import numpy as np
import scipy as sp
import unittest as ut


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.4826 if normal else 1.0}
                else:
                    torch_kwargs = {"normal": normal}
                    numpy_kwargs = {"axis": dim, "nan_policy": "omit", "scale": 1.4826 if normal else 1.0}

                # Compare the nan-median
                torch_mad = nan_mad(tensor, **torch_kwargs)
                numpy_mad = sp.stats.median_absolute_deviation(tensor.numpy(), **numpy_kwargs)
                np.testing.assert_almost_equal(torch_mad.numpy(), numpy_mad, decimal=4, err_msg=err_msg)


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


# goli/cli/__init__.py

from .main import main_cli
from .data import data_cli


# goli/cli/data.py

import click

from loguru import logger

import goli

from .main import main_cli


@main_cli.group(name="data", help="Goli datasets.")
def data_cli():
    pass


@data_cli.command(name="download", help="Download a Goli dataset.")
@click.option(
    "-n",
    "--name",
    type=str,
    required=True,
    help="Name of the goli dataset to download.",
)
@click.option(
    "-o",
    "--output",
    type=str,
    required=True,
    help="Where to download the Goli 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 = goli.data.utils.download_goli_dataset(**args)

    logger.info(f"Dataset available at {fpath}.")


@data_cli.command(name="list", help="List available Goli dataset.")
def list():

    logger.info("Goli datasets:")
    logger.info(goli.data.utils.list_goli_datasets())


# goli/cli/main.py

import click


@click.group()
@click.version_option()
def main_cli():
    pass


if __name__ == "__main__":
    main_cli()


# goli/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


# goli/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("goli.config", f"{name}.yaml") as f:
        config = omegaconf.OmegaConf.load(f)

    return config


# goli/config/_loader.py

from typing import List, Dict, Union, Any

import omegaconf
from copy import deepcopy
import torch
from loguru import logger

from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning import Trainer
from goli.nn.architectures import TaskHeadParams

from goli.trainer.metrics import MetricWrapper
from goli.nn import FullDGLNetwork, FullDGLSiameseNetwork, FullDGLMultiTaskNetwork, FeedForwardNN
from goli.trainer.predictor import PredictorModule
from goli.utils.spaces import DATAMODULE_DICT

from pytorch_lightning.loggers.tensorboard import TensorBoardLogger


def load_datamodule(
    config: Union[omegaconf.DictConfig, Dict[str, Any]],
):
    module_class = DATAMODULE_DICT[config["datamodule"]["module_type"]]
    datamodule = module_class(**config["datamodule"]["args"])

    return datamodule


def load_metrics(config: Union[omegaconf.DictConfig, Dict[str, Any]]):

    metrics = {}
    cfg_metrics = deepcopy(config["metrics"])
    if cfg_metrics is None:
        return metrics

    for this_metric in cfg_metrics:
        name = this_metric.pop("name")
        metrics[name] = MetricWrapper(**this_metric)

    return metrics


def load_architecture(
    config: Union[omegaconf.DictConfig, Dict[str, Any]],
    in_dim_nodes: int,
    in_dim_edges: int,
):

    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 == "fulldglnetwork":
        model_class = FullDGLNetwork
    elif model_type == "fulldglsiamesenetwork":
        model_class = FullDGLSiameseNetwork
        kwargs["dist_method"] = cfg_arch["dist_method"]
    elif model_type == "fulldglmultitasknetwork":
        model_class = FullDGLMultiTaskNetwork
    else:
        raise ValueError(f"Unsupported model_type=`{model_type}`")

    # Prepare the various kwargs
    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"])
    post_nn_kwargs = dict(cfg_arch["post_nn"]) if cfg_arch["post_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
    task_head_params_list = []
    for params in omegaconf.OmegaConf.to_object(task_heads_kwargs): # This turns the ListConfig into List[TaskHeadParams]
        params_dict = dict(params)
        task_head_params_list.append(TaskHeadParams(**params_dict))

    # Set the input dimensions
    if pre_nn_kwargs is not None:
        pre_nn_kwargs = dict(pre_nn_kwargs)
        pre_nn_kwargs.setdefault("in_dim", in_dim_nodes)
    else:
        gnn_kwargs.setdefault("in_dim", in_dim_nodes)

    if pre_nn_edges_kwargs is not None:
        pre_nn_edges_kwargs = dict(pre_nn_edges_kwargs)
        pre_nn_edges_kwargs.setdefault("in_dim", in_dim_edges)
    else:
        gnn_kwargs.setdefault("in_dim_edges", in_dim_edges)

    # Set the parameters for the full network
    if task_heads_kwargs is None:
        model_kwargs = dict(
            gnn_kwargs=gnn_kwargs,
            pre_nn_kwargs=pre_nn_kwargs,
            pre_nn_edges_kwargs=pre_nn_edges_kwargs,
            post_nn_kwargs=post_nn_kwargs,
        )
    else:
        model_kwargs = dict(
            gnn_kwargs=gnn_kwargs,
            pre_nn_kwargs=pre_nn_kwargs,
            pre_nn_edges_kwargs=pre_nn_edges_kwargs,
            post_nn_kwargs=post_nn_kwargs,
            task_heads_kwargs=task_head_params_list,
        )

    return model_class, model_kwargs


def load_predictor(config, model_class, model_kwargs, metrics):
    # Defining the predictor

    cfg_pred = dict(deepcopy(config["predictor"]))
    predictor = PredictorModule(
        model_class=model_class,
        model_kwargs=model_kwargs,
        metrics=metrics,
        **cfg_pred,
    )

    return predictor


def load_trainer(config):
    cfg_trainer = deepcopy(config["trainer"])

    # Set the number of gpus to 0 if no GPU is available
    gpus = cfg_trainer["trainer"].pop("gpus", 0)
    num_gpus = 0
    if isinstance(gpus, int):
        num_gpus = gpus
    elif isinstance(gpus, (list, tuple)):
        num_gpus = len(gpus)
    if (num_gpus > 0) and (not torch.cuda.is_available()):
        logger.warning(
            f"Number of GPUs selected is `{num_gpus}`, but will be ignored since no GPU are available on this device"
        )
        gpus = 0

    trainer_kwargs = {}
    callbacks = []
    if "early_stopping" in cfg_trainer.keys():
        callbacks.append(EarlyStopping(**cfg_trainer["early_stopping"]))

    if "model_checkpoint" in cfg_trainer.keys():
        callbacks.append(ModelCheckpoint(**cfg_trainer["model_checkpoint"]))

    if "logger" in cfg_trainer.keys():
        trainer_kwargs["logger"] = TensorBoardLogger(**cfg_trainer["logger"], default_hp_metric=False)

    trainer_kwargs["callbacks"] = callbacks

    trainer = Trainer(
        terminate_on_nan=True,
        **cfg_trainer["trainer"],
        **trainer_kwargs,
        gpus=gpus,
    )

    return trainer


# goli/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


# goli/data/__init__.py

from .utils import load_micro_zinc
from .utils import load_tiny_zinc

from .collate import goli_collate_fn

from .datamodule import DGLOGBDataModule
from .datamodule import DGLFromSmilesDataModule
from .datamodule import DGLOGBDataModule
from .datamodule import MultitaskDGLFromSmilesDataModule

from .datamodule import DGLDataset
from .datamodule import SingleTaskDataset
from .datamodule import MultitaskDGLDataset


# goli/data/collate.py

from collections.abc import Mapping
#from pprint import pprint
import torch
from torch.utils.data.dataloader import default_collate
import dgl
from inspect import signature, _empty
from typing import Union, List, Optional, Dict, Type, Any

from goli.features import dgl_dict_to_graph, DGLGraphDict


def goli_collate_fn(
    elements,
    labels_size_dict: Optional[Dict[str, Any]] = None,
    mask_nan: Union[str, float, Type[None]] = "raise", 
    do_not_collate_keys: List[str] = []
):
    """This collate function is identical to the default
    pytorch collate function but add support for `dgl.DGLGraph`
    objects and use `dgl.batch` to batch graphs.

    Beside dgl 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.

    Important:
        Only dgl graph within a dict are currently supported. It should not be hard
        to support dgl graphs from other objects.

    Note:
        If goli 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 MultiTaskDGLDataset.
            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.

        mask_nan:
            Deal with the NaN/Inf when calling the function `dgl_dict_to_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
    """

    elem = elements[0]

    if isinstance(elem, Mapping):
        batch = {}
        for key in elem:
            # If the features are a dictionary containing DGLGraph elements,
            # Convert to DGLGraph and use the dgl batching.
            if isinstance(elem[key], DGLGraphDict):
                graphs = [d[key].make_dgl_graph(mask_nan=mask_nan) for d in elements]
                batch[key] = dgl.batch(graphs)

            # If a DGLGraph is provided, use the dgl batching
            elif isinstance(elem[key], dgl.DGLGraph):
                batch[key] = dgl.batch([d[key] for d in elements])

            # Ignore the collate for specific keys
            elif key in do_not_collate_keys:
                batch[key] = [d[key] for d in elements]

            # Multitask setting: We have to pad the missing labels
            elif key == 'labels':
                if labels_size_dict is not None:
                    for datum in elements:
                        empty_task_labels = set(labels_size_dict.keys()) - set(datum["labels"].keys())
                        for task in empty_task_labels:
                            datum['labels'][task] = torch.full((len(elements), labels_size_dict[task]), torch.nan)
                batch[key] = default_collate([datum[key] for datum in elements])
            # Otherwise, use the default torch batching
            else:
                batch[key] = default_collate([d[key] for d in elements])
#        print("The batch contains: ")
#        pprint(batch)
        return batch
    else:
        return default_collate(elements)

# goli/data/datamodule.py

from enum import unique
from queue import Empty
from typing import Type, List, Dict, Union, Any, Callable, Optional, Tuple, Iterable

import os
from functools import partial
import importlib.resources
import zipfile
from copy import deepcopy
from itertools import cycle
import random

from loguru import logger
import fsspec
import omegaconf
from sklearn import datasets
from tqdm import tqdm
from joblib import Parallel, delayed
import tempfile

import pandas as pd
import numpy as np

from sklearn.model_selection import train_test_split

import dgl
import pytorch_lightning as pl
from goli.features.featurizer import dgl_dict_to_graph

from goli.utils import fs
from goli.features import mol_to_dglgraph_dict, mol_to_dglgraph_signature, mol_to_dglgraph, DGLGraphDict
from goli.data.collate import goli_collate_fn
from goli.utils.arg_checker import check_arg_iterator

import torch
from torch.utils.data.dataloader import DataLoader, Dataset
from torch.utils.data import Subset

import datamol as dm
#from pprint import pprint

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",
    "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",
        "version": 2,
    }
)

def smiles_to_unique_mol_ids(smiles: List[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."""
    unique_mol_ids = []
    for smile in smiles:
        mol = dm.to_mol(mol=smile)
        id = dm.unique_id(mol)
        unique_mol_ids.append(id)
    return unique_mol_ids

def get_indices(data_list: List, query_list: List):
    """This function searches for data_list's elements within query_list and returns the indices of query_list where they were found."""
    not_found = set(query_list)
    datum_idx_map = {}
    for index, datum in enumerate(data_list):
        if datum in not_found:
            datum_idx_map[datum] = index
            not_found.remove(datum)
            if len(not_found) == 0:
                break
    return [datum_idx_map[d] for d in query_list]

# Modify this if the dictionaries change structure.
def concatenate_respective_lists_and_save_indices(dict_containing_lists: Dict[str, List], key_containing_list: str):
    """Given a dict with values List, this function concatenates all the lists into one, and saves the respective indices of each key's list in the concatenated list.
    
    Currently this function assumes the following structure:
        {
            "key": {
                "key_containing_list": []
                "other_key": Any
                ...
            }
            ...
        }
    """
    concatenated_list = []
    dict_respective_indices = {}
    start = 0
    for key, value in dict_containing_lists.items():
        list = value[key_containing_list]
        concatenated_list.extend(list)

        end = start + len(list) - 1

        dict_respective_indices[key] = {}
        dict_respective_indices[key]["start"] = start
        dict_respective_indices[key]["end"] = end

        start = end + 1
    return concatenated_list, dict_respective_indices

class DGLDataset(Dataset):
    def __init__(
        self,
        features: List[dgl.DGLGraph],
        labels: Union[torch.Tensor, np.ndarray],
        smiles: Optional[List[str]] = None,
        indices: Optional[List[str]] = None,
        weights: Optional[Union[torch.Tensor, np.ndarray]] = None,
    ):
        self.smiles = smiles
        self.features = features
        self.labels = labels
        self.indices = indices
        self.weights = weights

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

    def __getitem__(self, idx):
        datum = {}

        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]

        datum["features"] = self.features[idx]

        datum["labels"] = self.labels[idx]

        return datum

class SingleTaskDataset(Dataset):
    def __init__(
        self,
        labels: Union[torch.Tensor, np.ndarray],
        features: Optional[List[dgl.DGLGraph]] = None,
        smiles: Optional[List[str]] = None,
        indices: Optional[List[str]] = None,
        weights: Optional[Union[torch.Tensor, np.ndarray]] = None,
    ):
        self.features = features
        self.labels = labels
        self.smiles = smiles
        self.indices = indices
        self.weights = weights

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

    def __getitem__(self, idx):
        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]

        return datum

class MultitaskDGLDataset(Dataset):
    """This class holds the information for the multitask dataset.
  
    Several single-task datasets can be merged to create a multi-task dataset."""
    def __init__(self, datasets: Dict[str, SingleTaskDataset]):
        super().__init__()
        self.datasets = datasets
        
        task = next(iter(self.datasets))
        if "features" in datasets[task].__getitem__(0):
            self.mol_ids, self.smiles, self.labels, self.features = self.merge(self.datasets)
        else:
            self.mol_ids, self.smiles, self.labels = self.merge(self.datasets)
        #self.labels_size = self.set_label_size_dict()
 
    def __len__(self):
        return len(self.mol_ids)
 
    def __getitem__(self, idx):
        datum = {}
 
        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 merge(self, datasets):
        all_smiles = []
        all_features = []
        all_labels = []
 
        all_tasks = []
        for task, ds in datasets.items():
            # Get data from single task ds
            ds_smiles = [ds.__getitem__(i)["smiles"] for i in range(len(ds))]
            ds_labels = [ds.__getitem__(i)["labels"] for i in range(len(ds))]
            if "features" in ds.__getitem__(0):
                ds_features = [ds.__getitem__(i)["features"] for i in range(len(ds))]
            else:
                ds_features = None
 
            all_smiles.extend(ds_smiles)
            all_labels.extend(ds_labels)
            if ds_features is not None: all_features.extend(ds_features)
 
            task_list = [task] * ds.__len__()
            all_tasks.extend(task_list)
 
        mol_ids = []
        # Get all unique mol ids.
        all_mol_ids = smiles_to_unique_mol_ids(all_smiles)
        unique_mol_ids, inv = np.unique(all_mol_ids, return_inverse=True)
        mol_ids = unique_mol_ids
 
        #print("\n\n\nUnique mol ID's! ")
        #pprint(mol_ids)
        #print("\nAll mol IDs ")
        #pprint(all_mol_ids)
        #print("\nAll smiles ")
        #pprint(all_smiles)
 
        # Store the smiles.
        #smiles = [[] for i in range(len(mol_ids))]
        smiles = [[]] * len(mol_ids)
        for all_idx, unique_idx in enumerate(inv):
            smiles[unique_idx].append(all_smiles[all_idx])
 
        # Store the labels.
        #labels = [{}] * len(mol_ids)
        labels = [{} for i in range(len(mol_ids))]
        for all_idx, unique_idx in enumerate(inv):
            task = all_tasks[all_idx]
            label = all_labels[all_idx]
            labels[unique_idx][task] = label
 
        #print("All features", all_features)
        #print("length of mol_id", len(mol_ids))
        #print("length of all features", len(all_features))
 
        # Store the features
        if len(all_features) > 0:
            features = [-1 for i in range(len(mol_ids))]
            for all_idx, unique_idx in enumerate(inv):
                features[unique_idx] = all_features[all_idx]
            return mol_ids, smiles, labels, features
        else:
            return mol_ids, smiles, labels

    def set_label_size_dict(self):
        # This gives the number of labels to predict for a given task.
        task_labels_size = {}
        for task, ds in self.datasets.items():
            label = ds.__getitem__(0)["labels"]       # Assume for a fixed task, the label dimension is the same across data points, so we can choose the first data point for simplicity.
            torch_label = torch.from_numpy(label)
            #torch_label = label
            #print("\n\n\nTorch label: ", torch_label)
            #print("Type: ", type(torch_label))
            #print("The size: ", torch_label.size())
            task_labels_size[task] = torch_label.size()
        return task_labels_size

    # def print_data(self):
    #     print("\n\n mol_ids")
    #     pprint(self.mol_ids)
    #     print("\n\n smiles")
    #     pprint(self.smiles)
    #     print("\n\n labels")
    #     pprint(self.labels)

class DGLBaseDataModule(pl.LightningDataModule):
    def __init__(
        self,
        batch_size_train_val: int = 16,
        batch_size_test: int = 16,
        num_workers: int = 0,
        pin_memory: bool = True,
        persistent_workers: bool = False,
        collate_fn: Optional[Callable] = None,
    ):
        super().__init__()

        self.batch_size_train_val = batch_size_train_val
        self.batch_size_test = batch_size_test

        self.num_workers = num_workers
        self.pin_memory = pin_memory
        self.persistent_workers = persistent_workers

        if collate_fn is None:
            # Some values become `inf` when changing data type. `mask_nan` deals with that
            self.collate_fn = partial(goli_collate_fn, mask_nan=0)
            self.collate_fn.__name__ = goli_collate_fn.__name__
        else:
            self.collate_fn = collate_fn

        self.dataset = None
        self.train_ds = None
        self.val_ds = None
        self.test_ds = None
        self._predict_ds = None

    def prepare_data(self):
        raise NotImplementedError()

    def setup(self):
        raise NotImplementedError()

    def train_dataloader(self, **kwargs):
        return self._dataloader(
            dataset=self.train_ds,  # type: ignore
            batch_size=self.batch_size_train_val,
            shuffle=True,
        )

    def val_dataloader(self, **kwargs):
        return self._dataloader(
            dataset=self.val_ds,  # type: ignore
            batch_size=self.batch_size_train_val,
            shuffle=False,
        )

    def test_dataloader(self, **kwargs):

        return self._dataloader(
            dataset=self.test_ds,  # type: ignore
            batch_size=self.batch_size_test,
            shuffle=False,
        )

    def predict_dataloader(self, **kwargs):

        return self._dataloader(
            dataset=self.predict_ds,  # type: ignore
            batch_size=self.batch_size_test,
            shuffle=False,
        )

    @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

    @predict_ds.setter
    def predict_ds(self, value):
        """Set the dataset for the prediction"""
        self._predict_ds = value

    def get_first_graph(self):
        raise NotImplementedError()

    # Private methods

    @staticmethod
    def _read_csv(path, **kwargs):
        if str(path).endswith((".csv", ".csv.gz", ".csv.zip", ".csv.bz2")):
            sep = ","
        elif str(path).endswith((".tsv", ".tsv.gz", ".tsv.zip", ".tsv.bz2")):
            sep = "\t"
        else:
            raise ValueError(f"unsupported file `{path}`")
        kwargs.setdefault("sep", sep)
        df = pd.read_csv(path, **kwargs)
        return df

    def _dataloader(self, dataset: Dataset, batch_size: int, shuffle: bool):
        """Get a dataloader for a given dataset"""

        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

        # TODO (Gabriela): Develop new dataloader to handle special batching.
        loader = DataLoader(
            dataset=dataset,
            num_workers=num_workers,
            collate_fn=self.collate_fn,
            pin_memory=self.pin_memory,
            batch_size=batch_size,
            shuffle=shuffle,
            persistent_workers=self.persistent_workers,
        )
        return loader  

class DGLFromSmilesDataModule(DGLBaseDataModule):
    """
    NOTE(hadim): let's make only one class for the moment and refactor with a parent class
    once we have more concrete datamodules to implement. The class should be general enough
    to be easily refactored.

    NOTE(hadim): splitting is not very full-featured yet; only random splitting on-the-fly
    is allowed using a seed. Next is to add the availability to provide split indices data as input.

    NOTE(hadim): implement using weights. It should probably be a column in the dataframe.
    """

    def __init__(
        self,
        df: pd.DataFrame = None,
        df_path: Optional[Union[str, os.PathLike]] = None,
        cache_data_path: Optional[Union[str, os.PathLike]] = None,
        featurization: Optional[Union[Dict[str, Any], omegaconf.DictConfig]] = None,
        smiles_col: str = None,
        label_cols: List[str] = None,
        weights_col: str = None,
        weights_type: str = None,
        idx_col: str = None,
        sample_size: Union[int, float, Type[None]] = None,
        split_val: float = 0.2,
        split_test: float = 0.2,
        split_seed: int = None,
        splits_path: Optional[Union[str, os.PathLike]] = None,
        batch_size_train_val: int = 16,
        batch_size_test: int = 16,
        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 = "dgldict",
        dataset_class: type = DGLDataset,
    ):
        """

        Parameters:
            df: a dataframe.
            df_path: a path to a dataframe to load (CSV file). `df` takes precedence over
                `df_path`.
            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 DGL featurizer.
            smiles_col: 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.
            label_cols: 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`

            weights_col: Name of the column to use as sample weights. If `None`, no
                weights are used. This parameter cannot be used together with `weights_type`.
            weights_type: 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.

            idx_col: Name of the columns to use as indices. Unused if set to None.
            split_val: Ratio for the validation split.
            split_test: Ratio for the test split.
            split_seed: Seed to use for the random split. More complex splitting strategy
                should be implemented.
            splits_path: 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`.
            batch_size_train_val: batch size for training and val dataset.
            batch_size_test: 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 `goli.data.goli_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.
            prepare_dict_or_graph: Whether to preprocess all molecules as DGL graphs or dict.
                Possible options:

                - "graph": Process molecules as dgl.DGLGraph. It's slower during pre-processing
                  and requires more RAM. It is faster during training with `num_workers=0`, but
                  slower with larger `num_workers`.
                - "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
                  DGLGraphs will be created during data-loading, but faster with large
                  `num_workers`, and less likely to cause memory issues with the parallelization.
            dataset_class: The class used to create the dataset from which to sample.
        """
        super().__init__(
            batch_size_train_val=batch_size_train_val,
            batch_size_test=batch_size_test,
            num_workers=num_workers,
            pin_memory=pin_memory,
            persistent_workers=persistent_workers,
            collate_fn=collate_fn,
        )

        self.df = df
        self.df_path = df_path

        self.cache_data_path = str(cache_data_path) if cache_data_path is not None else None
        self.featurization = featurization if featurization is not None else {}

        self.smiles_col = smiles_col
        self.label_cols = self._parse_label_cols(label_cols, smiles_col)
        self.idx_col = idx_col
        self.sample_size = sample_size

        self.weights_col = weights_col
        self.weights_type = weights_type
        if self.weights_col is not None:
            assert self.weights_type is None

        self.split_val = split_val
        self.split_test = split_test
        self.split_seed = split_seed
        self.splits_path = splits_path

        self.featurization_n_jobs = featurization_n_jobs
        self.featurization_progress = featurization_progress
        self.featurization_backend = featurization_backend

        self.dataset = None
        self.train_ds = None
        self.val_ds = None
        self.test_ds = None
        self._predict_ds = None
        self.train_indices = None
        self.val_indices = None
        self.test_indices = None
        self.dataset_class = dataset_class

        # Depreciated options
        if prepare_dict_or_graph == "dict":
            logger.warning("Depreciated: Use `prepare_dict_or_graph = 'dgldict'` instead of 'dict'")
            prepare_dict_or_graph = "dgldict"
        elif prepare_dict_or_graph == "graph":
            logger.warning("Depreciated: Use `prepare_dict_or_graph = 'dglgraph'` instead of 'graph'")
            prepare_dict_or_graph = "dglgraph"

        # Whether to transform the smiles into a dglgraph or a dictionary compatible with dgl
        if prepare_dict_or_graph == "dgldict":
            self.smiles_transformer = partial(mol_to_dglgraph_dict, **featurization)
        elif prepare_dict_or_graph == "dglgraph":
            self.smiles_transformer = partial(mol_to_dglgraph, **featurization)
        else:
            raise ValueError(
                f"`prepare_dict_or_graph` should be either 'dgldict' or 'dglgraph', Provided: `{prepare_dict_or_graph}`"
            )

    def prepare_data(self): # Can create train_ds, val_ds and test_ds here instead of setup. Be careful with the featurization (to not compute several times).
        """Called only from a single process in distributed settings. Steps:

        - If cache is set and exists, reload from cache.
        - Load the dataframe if its a path.
        - Extract smiles and labels from the dataframe.
        - Compute the features.
        - Compute or set split indices.
        - Make the list of dict dataset.
        """

        if self._load_from_cache():
            return True

        if self.df is None:
            # Only load the useful columns, as some dataset can be very large
            # when loading all columns
            label_cols = check_arg_iterator(self.label_cols, enforce_type=list)
            usecols = (
                check_arg_iterator(self.smiles_col, enforce_type=list)
                + label_cols
                + check_arg_iterator(self.idx_col, enforce_type=list)
                + check_arg_iterator(self.weights_col, enforce_type=list) # Can remove the weights (Dom says it's OK)
            )
            label_dtype = {col: np.float16 for col in label_cols}

            df = self._read_csv(self.df_path, usecols=usecols, dtype=label_dtype)
        else:
            df = self.df

        df = self._sub_sample_df(df)

        logger.info(f"Prepare dataset with {len(df)} data points.")

        # Extract smiles and labels
        smiles, labels, sample_idx, extras = self._extract_smiles_labels(
            df,
            smiles_col=self.smiles_col,
            label_cols=self.label_cols,
            idx_col=self.idx_col,
            weights_col=self.weights_col,
            weights_type=self.weights_type,
        )

        # Convert SMILES to features (graphs, fingerprints, etc.)
        features, idx_none = self._featurize_molecules(smiles, sample_idx)

        # Filter the molecules, labels, etc. for the molecules that failed featurization
        df, features, smiles, labels, sample_idx, extras = self._filter_none_molecules(
            idx_none, df, features, smiles, labels, sample_idx, extras
        )

        # Get splits indices
        self.train_indices, self.val_indices, self.test_indices = self._get_split_indices(
            len(df),
            split_val=self.split_val,
            split_test=self.split_test,
            split_seed=self.split_seed,
            splits_path=self.splits_path,
            sample_idx=sample_idx,
        )

        # Make the torch datasets (mostly a wrapper there is no memory overhead here)
        self.dataset = self.dataset_class(
            smiles=smiles,
            features=features,
            labels=labels,
            **extras,
        )

        self._save_to_cache()

    def setup(self, stage: str = None): # Can get rid of setup because a single dataset will have molecules exclusively in train, val or test
        """Prepare the torch dataset. Called on every GPUs. Setting state here is ok."""

        if stage == "fit" or stage is None:
            self.train_ds = Subset(self.dataset, self.train_indices)  # type: ignore
            self.val_ds = Subset(self.dataset, self.val_indices)  # type: ignore

        if stage == "test" or stage is None:
            self.test_ds = Subset(self.dataset, self.test_indices)  # type: ignore

    def _featurize_molecules(
        self, smiles: Iterable[str], sample_idx: Iterable[int] = None
    ) -> 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
        """

        # Loop all the smiles and compute the features
        if self.featurization_n_jobs == 0:
            features = [self.smiles_transformer(s) for s in tqdm(smiles)]
        else:
            features = Parallel(n_jobs=self.featurization_n_jobs, backend=self.featurization_backend)(
                delayed(self.smiles_transformer)(s) for s in tqdm(smiles)
            )

        # Warn about None molecules
        if sample_idx is None:
            sample_idx = np.arange(len(smiles))
        idx_none = [ii for ii, feat in enumerate(features) if feat is None]
        if len(idx_none) > 0:
            mols_to_msg = [f"{sample_idx[idx]}: {smiles[idx]}" 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] = DGLFromSmilesDataModule._filter_none_molecules(idx_none, val)
            else:
                new = arg
            out.append(new)

        out = tuple(out) if len(out) > 1 else out[0]

        return out

    def _parse_label_cols(self, 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.
        """
        if self.df is None:
            # Only load the useful columns, as some dataset can be very large
            # when loading all columns
            df = self._read_csv(self.df_path, nrows=0)
        else:
            df = self.df
        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_first_graph()
        num_feats = 0
        if "feat" in graph.ndata.keys():
            num_feats += graph.ndata["feat"].shape[1]
        return num_feats

    @property
    def num_node_feats_with_positional_encoding(self):
        """Return the number of node features in the first graph
        including positional encoding features."""

        graph = self.get_first_graph()
        num_feats = 0
        if "feat" in graph.ndata.keys():
            num_feats += graph.ndata["feat"].shape[1]
        if "pos_enc_feats_sign_flip" in graph.ndata.keys():
            num_feats += graph.ndata["pos_enc_feats_sign_flip"].shape[1]
        if "pos_enc_feats_no_flip" in graph.ndata.keys():
            num_feats += graph.ndata["pos_enc_feats_no_flip"].shape[1]
        return num_feats

    @property
    def num_edge_feats(self):
        """Return the number of edge features in the first graph"""

        graph = self.get_first_graph()
        if "feat" in graph.edata.keys():
            return graph.edata["feat"].shape[1]  # type: ignore
        else:
            return 0

    def get_first_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 10 first element, then `None` is returned, otherwise the first
        graph to not fail is returned.
        """
        if self.df is None:
            df = self._read_csv(self.df_path, nrows=10)
        else:
            df = self.df.iloc[0:10, :]

        smiles, _, _, _ = self._extract_smiles_labels(df, self.smiles_col, self.label_cols)

        graph = None
        for s in smiles:
            graph = self.smiles_transformer(s)
            if graph is not None:
                break

        if isinstance(graph, DGLGraphDict):
            graph = graph.make_dgl_graph(mask_nan=0.0)

        return graph

    # Private methods

    def _save_to_cache(self):
        """Save the built dataset, indices and featurization arguments into a cache file."""

        # Cache on disk
        if self.cache_data_path is not None:
            logger.info(f"Write prepared datamodule to {self.cache_data_path}")
            cache = {}
            cache["dataset"] = self.dataset
            cache["train_indices"] = self.train_indices
            cache["val_indices"] = self.val_indices
            cache["test_indices"] = self.test_indices

            # Save featurization args used
            cache["featurization_args"] = mol_to_dglgraph_signature(dict(self.featurization or {}))

            with fsspec.open(self.cache_data_path, "wb", compression="infer") as f:
                torch.save(cache, f)

    def _load_from_cache(self):
        """Attempt to reload the data from cache. Return True if data has been
        reloaded from the cache.
        """

        if self.cache_data_path is None:
            # Cache path is not provided.
            return False

        if not fs.exists(self.cache_data_path):
            # Cache patch does not exist.
            return False

        # Reload from cache if it exists and is valid
        logger.info(f"Try reloading the data module from {self.cache_data_path}.")

        # Load cache and save it locally in a temp folder.
        # This allows loading the cache much faster if it is zipped or in the cloud
        filesystem, _ = fsspec.core.url_to_fs(self.cache_data_path, mode="rb")
        protocol = check_arg_iterator(filesystem.protocol, enforce_type=list)
        filesystem = fsspec.filesystem(
            "filecache",
            target_protocol=protocol[0],
            target_options={"anon": True},
            cache_storage=tempfile.TemporaryDirectory().name,
            compression="infer",
        )
        with filesystem.open(self.cache_data_path, "rb", compression="infer") as f:
            cache = torch.load(f)

        # Are the required keys present?
        excepted_cache_keys = {
            "dataset",
            "test_indices",
            "train_indices",
            "val_indices",
        }
        if not set(cache.keys()) != excepted_cache_keys:
            logger.info(
                f"Cache looks invalid with keys: {cache.keys()}. Excepted keys are {excepted_cache_keys}"
            )
            logger.info("Fallback to regular data preparation steps.")
            return False

        # Is the featurization signature the same?
        current_signature = mol_to_dglgraph_signature(dict(self.featurization or {}))
        cache_signature = mol_to_dglgraph_signature(cache["featurization_args"])

        if current_signature != cache_signature:
            logger.info(f"Cache featurizer arguments are different than the provided ones.")
            logger.info(f"Cache featurizer arguments: {cache_signature}")
            logger.info(f"Provided featurizer arguments: {current_signature}.")
            logger.info("Fallback to regular data preparation steps.")
            return False

        # At this point the cache can be loaded

        self.dataset = cache["dataset"]
        self.train_indices = cache["train_indices"]
        self.val_indices = cache["val_indices"]
        self.test_indices = cache["test_indices"]

        logger.info(f"Datamodule correctly reloaded from cache.")

        return True

    def _extract_smiles_labels(
        self,
        df: pd.DataFrame,
        smiles_col: str = None,
        label_cols: List[str] = [],
        idx_col: str = None,
        weights_col: str = None,
        weights_type: 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.
        """

        if smiles_col is None:
            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:
            labels = [pd.to_numeric(df[col], errors="coerce") for col in label_cols]
            labels = np.stack(labels, axis=1)
        else:
            labels = np.zeros([len(smiles), 0])

        indices = None
        if idx_col is not None:
            indices = df[idx_col].values

        sample_idx = df.index.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 = {"indices": indices, "weights": weights}
        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,
    ):
        """Compute indices of random splits."""

        if sample_idx is None:
            sample_idx = np.arange(dataset_size)

        if splits_path is None:
            # Random splitting
            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)
            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
            with fsspec.open(str(splits_path)) as f:
                splits = self._read_csv(splits_path)

            train_indices = splits["train"].dropna().astype("int").tolist()
            val_indices = splits["val"].dropna().astype("int").tolist()
            test_indices = splits["test"].dropna().astype("int").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):
        # Sub-sample the dataframe
        if isinstance(self.sample_size, int):
            n = min(self.sample_size, df.shape[0])
            df = df.sample(n=n)
        elif isinstance(self.sample_size, float):
            df = df.sample(f=self.sample_size)
        elif self.sample_size is None:
            pass
        else:
            raise ValueError(f"Wrong value for `self.sample_size`: {self.sample_size}")

        return df

    def __len__(self) -> int:
        r"""
        Returns the number of elements of the current DataModule
        """
        if self.df is None:
            df = self._read_csv(self.df_path, usecols=[self.smiles_col])
        else:
            df = self.df

        return len(df)

    def to_dict(self):
        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_train_val"] = self.batch_size_train_val
        obj_repr["batch_size_test"] = self.batch_size_test
        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_labels"] = len(self.label_cols)
        obj_repr["collate_fn"] = self.collate_fn.__name__
        obj_repr["featurization"] = self.featurization
        return obj_repr

    def __repr__(self):
        """Controls how the class is printed"""
        return omegaconf.OmegaConf.to_yaml(self.to_dict())

class DatasetParams():
   def __init__(
       self,
       task_name: str,
       df: pd.DataFrame = None,
       df_path: Optional[Union[str, os.PathLike]] = None,
       smiles_col: str = None,
       label_cols: List[str] = None,
       weights_col: str = None,
       weights_type: str = None,
       idx_col: str = None,
       sample_size: Union[int, float, Type[None]] = None,
       split_val: float = 0.2,
       split_test: float = 0.2,
       split_seed: int = None,
       splits_path: Optional[Union[str, os.PathLike]] = None,
   ):
       self.task_name = task_name
       self.df = df
       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.sample_size = sample_size
       self.split_val = split_val
       self.split_test = split_test
       self.split_seed = split_seed
       self.splits_path = splits_path
class MultitaskDGLFromSmilesDataModule(DGLBaseDataModule):
    def __init__(
        self,
        #task_df: Dict[str, pd.DataFrame] = None,
        #task_df_path: Optional[Dict[str, Union[str, os.PathLike]]] = None,
        #task_smiles_col: Dict[str, str] = None,
        #task_label_cols: Dict[str, List[str]] = None,
        #task_weights_col: Dict[str, str] = None,
        #task_weights_type: Dict[str, str] = None,
        #task_idx_col: Dict[str, str] = None,
        #task_sample_size: Dict[str, Union[int, float, None]] = None,
        #task_split_val: Dict[str, float] = None,
        #task_split_test: Dict[str, float]= None,
        #task_split_seed: Dict[str, int] = None,
        #task_splits_path: Optional[Dict[str, Union[str, os.PathLike]]] = None,
        task_specific_args: Dict[str, Any],         # Create a class that so that we can have a dictionary of class
        cache_data_path: Optional[Union[str, os.PathLike]] = None,
        featurization: Optional[Union[Dict[str, Any], omegaconf.DictConfig]] = None,
        batch_size_train_val: int = 16,
        batch_size_test: int = 16,
        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 = "dgl_dict",
        dataset_class: type = MultitaskDGLDataset,
    ):
        """
        Parameters: only for parameters beginning with task_*, we have a dictionary where the key is the task name
        and the value is specified below.
            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_split_seed: (value) Seed to use for the random split. 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 DGL featurizer.
            batch_size_train_val: batch size for training and val dataset.
            batch_size_test: 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 `goli.data.goli_collate_fn`
            prepare_dict_or_graph: Whether to preprocess all molecules as DGL graphs or dict.
                Possible options:

                - "graph": Process molecules as dgl.DGLGraph. It's slower during pre-processing
                  and requires more RAM, but faster during training.
                - "dict": Process molecules as a Dict. It's faster and requires less RAM during
                  pre-processing, but slower during training since DGLGraphs will be created
                  during data-loading.
            dataset_class: The class used to create the dataset from which to sample. 
        """
        super().__init__(
            batch_size_train_val=batch_size_train_val,
            batch_size_test=batch_size_test,
            num_workers=num_workers,
            pin_memory=pin_memory,
            persistent_workers=persistent_workers,
            collate_fn=collate_fn,
        )

        #self.task_df = task_df
        #self.task_df_path = task_df_path
        #self.task_smiles_col = task_smiles_col
        #self.task_abel_cols = {task: self._parse_label_cols(task_df[task], task_df_path[task], task_label_cols[task], task_smiles_col[task]) for task in self.task_df.keys()} # Make sure this is not none
        #self.idx_col = task_idx_col
        #self.task_weights_col = task_weights_col      # Can remove this
        #self.task_weights_type = task_weights_type    # Can remove this
        #self.task_idx_col = task_idx_col
        #self.task_sample_size = task_sample_size
        #self.task_split_val = task_split_val
        #self.task_split_test = task_split_test
        #self.task_spit_seed = task_split_seed
        #self.task_splits_path = task_splits_path
        self.task_specific_args = task_specific_args

        self.featurization_n_jobs = featurization_n_jobs
        self.featurization_progress = featurization_progress
        self.featurization_backend = featurization_backend

        self.dataset_class = dataset_class

        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.dataset = None
        self.train_ds = None
        self.val_ds = None
        self.test_ds = None

        if prepare_dict_or_graph == "dgl_dict":         # Changed based on Dom's comments. Compare with version in this code to see if the change works or not.
            self.smiles_transformer = partial(mol_to_dglgraph_dict, **featurization)
        elif prepare_dict_or_graph == "graph":
            self.smiles_transformer = partial(mol_to_dglgraph, **featurization)
        else:
            raise ValueError(
                f"`prepare_dict_or_graph` should be either 'dict' or 'graph', Provided: `{prepare_dict_or_graph}`"
            )

    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
        """
        # TODO (Gabriela): Implement the ability to load from cache.

        """Load all single-task dataframes."""
        task_df = {}
        for task, args in self.task_specific_args.items():
            if args["df"] is None:
                # Only load the useful columns, as some datasets can be very large when loading all columns.
                label_cols = check_arg_iterator(args["label_cols"], enforce_type=list)
                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.float16 for col in label_cols}

                task_df[task] = self._read_csv(args["df_path"], usecols=usecols, dtype=label_dtype)
            else:
                task_df[task] = args["df"]

        """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_specific_args[task]["sample_size"]
            df = self._sub_sample_df(df, sample_size)

            logger.info(f"Prepare single-task dataset for task '{task}' with {len(df)} data points.")

            # Extract smiles, labels, extras
            args = self.task_specific_args[task]
            smiles, labels, sample_idx, extras = self._extract_smiles_labels(
                df,
                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 = []
        for task, dataset_args in task_dataset_args.items():
            all_smiles.extend(dataset_args["smiles"])
            for count in range(len(dataset_args["smiles"])):
                all_tasks.append(task)
        # Get all unique mol ids
        all_mol_ids = smiles_to_unique_mol_ids(all_smiles)
        unique_mol_ids, inv = np.unique(all_mol_ids, return_inverse=True)
        mol_ids = unique_mol_ids
        # Get smiles to be used for featurization with those mol ids.
        smiles_to_featurize = []
        for id in range(len(mol_ids)):
            for all_idx, unique_idx in enumerate(inv):
                if unique_idx == id:  # If we found this unique mol id
                    smiles_to_featurize.append(all_smiles[all_idx])
                    break # No need to check if another smiles matches
        assert len(mol_ids) == len(smiles_to_featurize)

        # Convert SMILES to features
        features, idx_none = self._featurize_molecules(smiles_to_featurize) # sample_idx is removed ... might need to add it again later in another way
        
        # 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 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])
        # Update idx_none per task for later filtering
        for task, args in task_dataset_args.items():
            for idx, feat in enumerate(args["features"]):
                if feat == None: args["idx_none"].append(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():
            df, features, smiles, labels, sample_idx, extras = self._filter_none_molecules(
                args["idx_none"], 
                task_df[task],
                args["features"],
                args["smiles"], 
                args["labels"],
                args["sample_idx"], 
                args["extras"]
            )

            # Update the data
            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] = SingleTaskDataset(
                features=task_dataset_args[task]["features"],
                labels=task_dataset_args[task]["labels"],
                smiles=task_dataset_args[task]["smiles"],
                **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 = {}

        self.train_singletask_datasets = {}
        self.val_singletask_datasets = {}
        self.test_singletask_datasets = {}
        for task, df in task_df.items():
            train_indices, val_indices, test_indices = self._get_split_indices(
                len(df),
                split_val=self.task_specific_args[task]["split_val"],
                split_test=self.task_specific_args[task]["split_test"],
                split_seed=self.task_specific_args[task]["split_seed"],
                splits_path=self.task_specific_args[task]["splits_path"],
                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

            #train_singletask_datasets[task] = Subset()
            self.train_singletask_datasets[task] = Subset(self.single_task_datasets[task], train_indices)
            self.val_singletask_datasets[task] = Subset(self.single_task_datasets[task], val_indices)
            self.test_singletask_datasets[task] = Subset(self.single_task_datasets[task], test_indices)


        # TODO (Gabriela): Implement the ability to save to cache. 

    def setup(self, stage: str = None): # Can possibly get rid of setup because a single dataset will have molecules exclusively in train, val or test
        """Prepare the torch dataset. Called on every GPUs. Setting state here is ok."""

        if stage == "fit" or stage is None:
            self.train_ds = MultitaskDGLDataset(self.train_singletask_datasets)  # type: ignore
            self.val_ds = MultitaskDGLDataset(self.val_singletask_datasets)  # type: ignore

        if stage == "test" or stage is None:
            self.test_ds = MultitaskDGLDataset(self.test_singletask_datasets)  # type: ignore

        # Produce the label sizes
        train_label_sizes = self.train_ds.set_label_size_dict()
        val_label_sizes = self.val_ds.set_label_size_dict()
        test_label_sizes = self.test_ds.set_label_size_dict()

        label_sizes = train_label_sizes
        label_sizes.update(val_label_sizes)
        label_sizes.update(test_label_sizes)

    # 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
        """

        # Loop all the smiles and compute the features
        if self.featurization_n_jobs == 0:
            features = [self.smiles_transformer(s) for s in tqdm(smiles)]
        else:
            features = Parallel(n_jobs=self.featurization_n_jobs, backend=self.featurization_backend)(
                delayed(self.smiles_transformer)(s) for s in tqdm(smiles)
            )

        # Warn about None molecules
        idx_none = [ii for ii, feat in enumerate(features) if feat is None]
        if len(idx_none) > 0:
            mols_to_msg = [f"{smiles[idx]}" 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] = DGLFromSmilesDataModule._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: str, 
        str: 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.
        """
        if df is None:
            # Only load the useful columns, as some dataset can be very large
            # when loading all columns
            data_frame = self._read_csv(df_path, nrows=0)
        else:
            data_frame = df
        cols = list(data_frame.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_first_graph()
        num_feats = 0
        if "feat" in graph.ndata.keys():
            num_feats += graph.ndata["feat"].shape[1]
        return num_feats

    @property
    def num_node_feats_with_positional_encoding(self):
        """Return the number of node features in the first graph
        including positional encoding features."""

        graph = self.get_first_graph()
        num_feats = 0
        if "feat" in graph.ndata.keys():
            num_feats += graph.ndata["feat"].shape[1]
        if "pos_enc_feats_sign_flip" in graph.ndata.keys():
            num_feats += graph.ndata["pos_enc_feats_sign_flip"].shape[1]
        if "pos_enc_feats_no_flip" in graph.ndata.keys():
            num_feats += graph.ndata["pos_enc_feats_no_flip"].shape[1]
        return num_feats

    @property
    def num_edge_feats(self):
        """Return the number of edge features in the first graph"""

        graph = self.get_first_graph()
        if "feat" in graph.edata.keys():
            return graph.edata["feat"].shape[1]  # type: ignore
        else:
            return 0

    def get_first_graph(self):      # Fix this
        """
        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 10 first element, then `None` is returned, otherwise the first
        graph to not fail is returned.
        """
        if self.df is None:
            df = self._read_csv(self.df_path, nrows=10)
        else:
            df = self.df.iloc[0:10, :]

        smiles, _, _, _ = self._extract_smiles_labels(df, self.smiles_col, self.label_cols)

        graph = None
        for s in smiles:
            graph = self.smiles_transformer(s)
            if graph is not None:
                break

        if isinstance(graph, dict):
            graph = dgl_dict_to_graph(**graph, mask_nan=0.0)

        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,
        smiles_col: str = None,
        label_cols: List[str] = [],
        idx_col: str = None,
        weights_col: str = None,
        weights_type: 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.
        """

        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:
            labels = [pd.to_numeric(df[col], errors="coerce") for col in label_cols]
            labels = np.stack(labels, axis=1)
        else:
            labels = np.zeros([len(smiles), 0])

        indices = None                      # What are indices for?
        if idx_col is not None:
            indices = df[idx_col].values

        sample_idx = df.index.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 = {"indices": indices, "weights": weights}
        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,
    ):
        """Compute indices of random splits."""

        if sample_idx is None:
            sample_idx = np.arange(dataset_size)

        if splits_path is None:
            # Random splitting
            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)
            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
            with fsspec.open(str(splits_path)) as f:
                splits = self._read_csv(splits_path)

            train_indices = splits["train"].dropna().astype("int").tolist()
            val_indices = splits["val"].dropna().astype("int").tolist()
            test_indices = splits["test"].dropna().astype("int").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, sample_size):
        # Sub-sample the dataframe
        if isinstance(sample_size, int):
            n = min(self.sample_size, df.shape[0])
            df = df.sample(n=n)
        elif isinstance(sample_size, float):
            df = df.sample(f=sample_size)
        elif sample_size is None:
            pass
        else:
            raise ValueError(f"Wrong value for `self.sample_size`: {self.sample_size}") # Maybe specify which task it was for?

        return df

    ########################## len  ###################################################
    def __len__(self) -> int:
        r"""
        Returns the number of elements of the current DataModule
        """
        num_elements = 0
        for task, args in self.task_specific_args.items():
            if args["df"] is None:
                df = self._read_csv(args["df_path"], usecols=args["smiles_col"])
                num_elements += len(df)
            else:
                num_elements += len(args["df"])
        return num_elements

    ########################### to dict ###############################################
    # Modify these
    def to_dict(self):
        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_train_val"] = self.batch_size_train_val
        obj_repr["batch_size_test"] = self.batch_size_test
        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_labels"] = len(self.label_cols)
        obj_repr["collate_fn"] = self.collate_fn.__name__
        obj_repr["featurization"] = self.featurization
        return obj_repr

    def __repr__(self):
        """Controls how the class is printed"""
        return omegaconf.OmegaConf.to_yaml(self.to_dict())

class DGLOGBDataModule(DGLFromSmilesDataModule):
    """Load an OGB GraphProp dataset."""

    def __init__(
        self,
        dataset_name: str,
        cache_data_path: Optional[Union[str, os.PathLike]] = None,
        featurization: Optional[Union[Dict[str, Any], omegaconf.DictConfig]] = None,
        weights_col: str = None,
        weights_type: str = None,
        sample_size: Union[int, float, Type[None]] = None,
        batch_size_train_val: int = 16,
        batch_size_test: int = 16,
        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,
    ):
        """

        Parameters:
            dataset_name: Name of the OGB dataset to load. Examples of possible datasets are
                "ogbg-molhiv", "ogbg-molpcba", "ogbg-moltox21", "ogbg-molfreesolv".
            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 DGL featurizer.
            batch_size_train_val: batch size for training and val dataset.
            batch_size_test: 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 `goli.data.goli_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.
        """

        self.dataset_name = dataset_name

        # Get OGB metadata
        self.metadata = self._get_dataset_metadata(self.dataset_name)

        # Get dataset
        df, idx_col, smiles_col, label_cols, splits_path = self._load_dataset(self.metadata)

        # Config for datamodule
        dm_args = {}
        dm_args["df"] = df
        dm_args["cache_data_path"] = cache_data_path
        dm_args["featurization"] = featurization
        dm_args["smiles_col"] = smiles_col
        dm_args["label_cols"] = label_cols
        dm_args["idx_col"] = idx_col
        dm_args["splits_path"] = splits_path
        dm_args["batch_size_train_val"] = batch_size_train_val
        dm_args["batch_size_test"] = batch_size_test
        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["weights_col"] = weights_col
        dm_args["weights_type"] = weights_type
        dm_args["sample_size"] = sample_size

        # Init DGLFromSmilesDataModule
        super().__init__(**dm_args)

    def to_dict(self):
        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):
        """Download, extract and load an OGB dataset."""

        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)

        # 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"
            logger.info(f"Saving splits to {splits_path}")
            splits.to_csv(splits_path, index=None)
        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"):
            idx_col = df.columns[0]
            smiles_col = df.columns[-2]
            label_cols = df.columns[-1:].to_list()
        else:
            idx_col = df.columns[-1]
            smiles_col = df.columns[-2]
            label_cols = df.columns[:-2].to_list()

        return df, idx_col, smiles_col, label_cols, splits_path

    def _get_dataset_metadata(self, dataset_name: str):
        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 = ogb_metadata.append(pd.DataFrame(PCQM4M_meta, index=["ogbg-lsc-pcqm4m"]))
        ogb_metadata = ogb_metadata.append(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


# goli/data/utils.py

import importlib.resources
import zipfile

import pandas as pd

import goli

GOLI_DATASETS_BASE_URL = "gcs://goli-public/datasets"
GOLI_DATASETS = {
    "goli-zinc-micro": "goli-zinc-micro.zip",
    "goli-zinc-bench-gnn": "goli-zinc-bench-gnn.zip",
    "goli-htsfp": "goli-htsfp.csv.gz",
}


def load_micro_zinc() -> pd.DataFrame:
    """Return a dataframe of micro ZINC (1000 data points)."""

    with importlib.resources.open_text("goli.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)."""

    with importlib.resources.open_text("goli.data.micro_ZINC", "micro_ZINC.csv") as f:
        df = pd.read_csv(f, nrows=100)

    return df  # type: ignore


def list_goli_datasets():
    """List Goli datasets available to download."""
    return set(GOLI_DATASETS.keys())


def download_goli_dataset(name: str, output_path: str, extract_zip: bool = True, progress: bool = False):
    """Download a Goli dataset to a specified location.

    Args:
        name: Name of the Goli dataset from `goli.data.utils.get_goli_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.
    """

    if name not in GOLI_DATASETS:
        raise ValueError(f"'{name}' is not a valid Goli dataset name. Choose from {GOLI_DATASETS}")

    fname = GOLI_DATASETS[name]

    dataset_path_source = goli.utils.fs.join(GOLI_DATASETS_BASE_URL, fname)
    dataset_path_destination = goli.utils.fs.join(output_path, fname)

    if not goli.utils.fs.exists(dataset_path_destination):
        goli.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


# goli/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_dglgraph_dict
from .featurizer import mol_to_dglgraph
from .featurizer import dgl_dict_to_graph
from .featurizer import mol_to_dglgraph_signature
from .featurizer import DGLGraphDict


# goli/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 csr_matrix, issparse, coo_matrix
import torch
import dgl

from rdkit import Chem
from rdkit.Chem.rdmolops import GetAdjacencyMatrix
from datamol import to_mol

from goli.features import nmp
from goli.utils.tensor import one_of_k_encoding
from goli.features.positional_encoding import get_all_positional_encoding


def _to_dense_array(array, dtype=None):
    # Assign the node data
    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 _mask_nans_inf(mask_nan, array, array_name):

    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: Chem.rdchem.Mol, property_list: List[str]) -> Dict[str, np.ndarray]:
    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"

    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"))]
            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_atomic_features_float(
    mol: Chem.rdchem.Mol,
    property_list: Union[List[str], List[Callable]],
    offset_carbon: bool = True,
    mask_nan: Union[str, float, type(None)] = "raise",
) -> Dict[str, np.ndarray]:
    """
    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"
            - "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"

        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

            if isinstance(prop, str):

                prop = prop.lower()
                prop_name = prop

                if prop in ["atomic-number"]:
                    val = (atom.GetAtomicNum() - (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 ["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(
                        atom.GetAtomicNum()
                    ) - 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[atom.GetAtomicNum() - 1] - offC * nmp.FIRST_IONIZATION[C_num - 1]
                    ) / 5
                elif prop in ["melting-point"]:
                    val = (
                        nmp.MELTING_POINT[atom.GetAtomicNum() - 1] - offC * nmp.MELTING_POINT[C_num - 1]
                    ) / 200
                elif prop in ["metal"]:
                    val = nmp.METAL[atom.GetAtomicNum() - 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: Chem.rdchem.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: Chem.rdchem.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: Chem.rdchem.Mol, property_list: List[str], mask_nan: Union[str, float, type(None)] = "raise"
):
    r"""
    Get the following set of features for any given bond
    See `goli.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:
            prop_dict[prop] = np.array([])

    return prop_dict


def mol_to_adj_and_features(
    mol: Union[str, Chem.rdchem.Mol],
    atom_property_list_onehot: List[str] = [],
    atom_property_list_float: List[Union[str, Callable]] = [],
    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,
    pos_encoding_as_directions: Dict[str, Any] = None,
    dtype: np.dtype = np.float16,
    mask_nan: Union[str, float, type(None)] = "raise",
) -> Union[csr_matrix, Union[np.ndarray, None], Union[np.ndarray, None]]:
    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`

        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.

        pos_encoding_as_directions: keyword arguments for function `graph_positional_encoder`
            to generate positional encoding for directional features.

        dtype:
            The numpy 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:
            Scipy 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`

        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.

        pos_enc_dir:
            Node positional encoding used to define directions. This can thus
            be used for relative position inference.
            This is used, for example, by `DGNConvolutionalLayer` to define
            the direction of the messages.

    """

    if isinstance(mol, str):
        mol = to_mol(mol)

    # Add or remove explicit hydrogens
    if explicit_H:
        mol = Chem.AddHs(mol)
    else:
        mol = Chem.RemoveHs(mol)

    # Get the adjacency matrix
    adj = GetAdjacencyMatrix(mol, useBO=use_bonds_weights, force=True)
    if add_self_loop:
        adj = adj + np.eye(adj.shape[0], dtype=np.int8)
    adj = csr_matrix(adj, dtype=np.int8)

    # 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)
    ndata = list(atom_features_float.values()) + list(atom_features_onehot.values())
    ndata = [np.expand_dims(d, axis=1) if d.ndim == 1 else d for d in ndata]

    if len(ndata) > 0:
        ndata = np.concatenate(ndata, axis=1).astype(dtype)
        ndata = coo_matrix(ndata)
    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)
        edata = coo_matrix(edata)
    else:
        edata = None

    pos_enc_feats_sign_flip, pos_enc_feats_no_flip, pos_enc_dir = get_all_positional_encoding(
        adj, pos_encoding_as_features, pos_encoding_as_directions
    )

    # Mask the NaNs
    pos_enc_feats_sign_flip = _mask_nans_inf(mask_nan, pos_enc_feats_sign_flip, "pos_enc_feats_sign_flip")
    pos_enc_feats_no_flip = _mask_nans_inf(mask_nan, pos_enc_feats_no_flip, "pos_enc_feats_no_flip")
    pos_enc_dir = _mask_nans_inf(mask_nan, pos_enc_dir, "pos_enc_dir")

    return adj, ndata, edata, pos_enc_feats_sign_flip, pos_enc_feats_no_flip, pos_enc_dir


class DGLGraphDict(dict):
    def __init__(
        self,
        dic: Dict,
    ):
        """
        Store the parameters required to initialize a `dgl.DGLGraph`, but
        as a dictionary to reduce memory consumption.

        Possible keys for the dictionary:

        - adj: A numpy array containing the adjacency matrix

        - ndata: A dictionnary containing different keys and numpy
            arrays associated to the node features `DGLGraph.ndata`.

        - edata: A dictionnary containing different keys and numpy
            arrays associated to the edge features `DGLGraph.edata`.

        - dtype: The numpy dtype for the floating data. The arrays
            will be converted to `torch.Tensor` when building 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
        """
        default_dic = {
            "dtype": np.float16,
            "mask_nan": "raise",
        }
        default_dic.update(dic)
        super().__init__(default_dic)

    def make_dgl_graph(self, **kwargs) -> dgl.DGLGraph:
        """
        Convert the current dictionary of parameters, containing an adjacency matrix with node/edge data
        into a `dgl.DGLGraph` of torch Tensors.

        `**kwargs` can be used to overwrite any parameter from the current dictionary. See `DGLGraphDict.__init__`
        for a list of parameters
        """
        this_dict = self.copy()
        for key, val in kwargs.items():
            this_dict[key] = val
        return dgl_dict_to_graph(**this_dict)

    @property
    def ndata(self):
        return self.dgl_dict["ndata"]

    @property
    def edata(self):
        return self.dgl_dict["edata"]

    @property
    def adj(self):
        return self.dgl_dict["adj"]

    @property
    def dtype(self):
        return self.dgl_dict["dtype"]

    @property
    def mask_nan(self):
        return self.dgl_dict["mask_nan"]


def mol_to_dglgraph_dict(
    mol: Chem.rdchem.Mol,
    atom_property_list_onehot: List[str] = [],
    atom_property_list_float: List[Union[str, Callable]] = [],
    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,
    pos_encoding_as_directions: Dict[str, Any] = None,
    dtype: np.dtype = np.float16,
    on_error: str = "ignore",
    mask_nan: Union[str, float, type(None)] = "raise",
) -> DGLGraphDict:
    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 `DGLGraph` object.

    Compared to `mol_to_dglgraph`, 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`

        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.

        pos_encoding_as_directions: keyword arguments for function `graph_positional_encoder`
            to generate positional encoding for directional 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 None
            - "ignore": Ignore the error and return None

        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:

        graph:
            A dictionary containing the keys required to build a DGLGraph fromt the function
            `mol_to_dglgraph_dict`. The keys are:

            - "adj": A sparse int-array containing the adjacency matrix

            - "ndata": A dictionnary containing different keys and numpy
              arrays associated to the node features.

            - "edata": A dictionnary containing different keys and numpy
              arrays associated to the edge features.

            - "dtype": The numpy dtype for the floating data.
    """

    input_mol = mol

    try:

        if isinstance(mol, str):
            mol = to_mol(mol)
        if explicit_H:
            mol = Chem.AddHs(mol)
        else:
            mol = Chem.RemoveHs(mol)

        # Get the adjacency, node features and edge features
        (
            adj,
            ndata,
            edata,
            pos_enc_feats_sign_flip,
            pos_enc_feats_no_flip,
            pos_enc_dir,
        ) = mol_to_adj_and_features(
            mol=mol,
            atom_property_list_onehot=atom_property_list_onehot,
            atom_property_list_float=atom_property_list_float,
            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,
            pos_encoding_as_directions=pos_encoding_as_directions,
            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, Chem.rdchem.Mol):
                smiles = Chem.MolToSmiles(input_mol)

            msg = str(e) + "\nIgnoring following molecule:" + smiles
            logger.warning(msg)
            return None
        elif on_error.lower() == "ignore":
            return None

    dgl_dict = {"adj": adj, "edata": {}, "ndata": {}, "dtype": dtype}

    # Assign the node data
    if ndata is not None:
        dgl_dict["ndata"]["feat"] = ndata

    # Assign the edge data. Due to DGL only supporting Hetero-graphs, we
    # need to duplicate each edge information for its 2 entries
    if edata is not None:
        edata = _to_dense_array(edata, dtype=dtype)
        src_ids, dst_ids = np.argwhere(adj).transpose()
        hetero_edata = np.zeros(shape=(edata.shape[0] * 2, edata.shape[1]), dtype=edata.dtype)
        for ii in range(mol.GetNumBonds()):
            bond = mol.GetBondWithIdx(ii)
            src, dst = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
            id1 = np.where((src == src_ids) & (dst == dst_ids))[0]
            id2 = np.where((dst == src_ids) & (src == dst_ids))[0]
            hetero_edata[id1, :] = edata[ii, :]
            hetero_edata[id2, :] = edata[ii, :]

        dgl_dict["edata"]["feat"] = coo_matrix(hetero_edata)

    # Add sign-flip positional encoding
    if pos_enc_feats_sign_flip is not None:
        dgl_dict["ndata"]["pos_enc_feats_sign_flip"] = pos_enc_feats_sign_flip

    # Add non-sign-flip positional encoding
    if pos_enc_feats_no_flip is not None:
        dgl_dict["ndata"]["pos_enc_feats_no_flip"] = pos_enc_feats_no_flip

    # Add positional encoding for directional use
    if pos_enc_dir is not None:
        dgl_dict["ndata"]["pos_dir"] = pos_enc_dir

    dgl_dict = DGLGraphDict(dgl_dict)

    return dgl_dict


def mol_to_dglgraph(
    mol: Chem.rdchem.Mol,
    atom_property_list_onehot: List[str] = [],
    atom_property_list_float: List[Union[str, Callable]] = [],
    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,
    pos_encoding_as_directions: Dict[str, Any] = None,
    dtype: np.dtype = np.float16,
    on_error: str = "ignore",
    mask_nan: Union[str, float, type(None)] = "raise",
) -> dgl.DGLGraph:
    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
    `DGLGraph` 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`

        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.

        pos_encoding_as_directions: keyword arguments for function `graph_positional_encoder`
            to generate positional encoding for directional 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 None
            - "ignore": Ignore the error and return None

        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

    Returns:

        graph:
            DGL graph, with `graph.ndata['feat']` corresponding to the concatenated
            node data from `atom_property_list_onehot` and `atom_property_list_float`,
            `graph.edata['feat']` corresponding to the concatenated edge data from `edge_property_list`.
            There are also additional entries for the positional encodings.

    """

    dgl_dict = mol_to_dglgraph_dict(
        mol=mol,
        atom_property_list_onehot=atom_property_list_onehot,
        atom_property_list_float=atom_property_list_float,
        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,
        pos_encoding_as_directions=pos_encoding_as_directions,
        dtype=dtype,
        on_error=on_error,
        mask_nan=mask_nan,
    )

    if dgl_dict is not None:
        return dgl_dict_to_graph(**dgl_dict)

    return None


def dgl_dict_to_graph(
    adj,
    ndata: Dict,
    edata: Dict,
    dtype: np.dtype = None,
    mask_nan: Union[str, float, type(None)] = "raise",
) -> dgl.DGLGraph:
    """
    Convert an adjacency matrix with node/edge data into a `DGLGraph`
    of torch Tensors.

    Parameters:

        adj: A numpy array containing the adjacency matrix

        ndata: A dictionnary containing different keys and numpy
            arrays associated to the node features `DGLGraph.ndata`.

        edata: A dictionnary containing different keys and numpy
            arrays associated to the edge features `DGLGraph.edata`.

        dtype: The numpy dtype for the floating data. The arrays
            will be converted to `torch.Tensor` when building 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

    """

    # Transform the matrix and data into a DGLGraph object
    graph = dgl.from_scipy(adj, idtype=torch.int32)

    if ndata is not None:
        for key, val in ndata.items():
            graph.ndata[key] = torch.as_tensor(
                _mask_nans_inf(mask_nan=mask_nan, array=_to_dense_array(val, dtype=dtype), array_name="ndata")
            )

    if edata is not None:
        for key, val in edata.items():
            graph.edata[key] = torch.as_tensor(
                _mask_nans_inf(mask_nan=mask_nan, array=_to_dense_array(val, dtype=dtype), array_name="edata")
            )

    return graph


def mol_to_dglgraph_signature(featurizer_args: Dict[str, Any] = None):
    """Get the default arguments of `mol_to_dglgraph_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.
    """

    # Get the signature of `mol_to_dglgraph_dict`
    signature = inspect.signature(mol_to_dglgraph_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


# goli/features/nmp.py

import importlib.resources

import pandas as pd
from rdkit import Chem

# NOTE(hadim): usually it's best to embed this in a function.
with importlib.resources.open_text("goli.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):
    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()

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,
]


# goli/features/positional_encoding.py

from typing import Tuple, Optional, Dict, Union
import numpy as np
from scipy.sparse import spmatrix
import torch

from goli.features.spectral import compute_laplacian_positional_eigvecs


def get_all_positional_encoding(
    adj: Union[np.ndarray, spmatrix],
    pos_encoding_as_features: Optional[Dict] = None,
    pos_encoding_as_directions: Optional[Dict] = None,
) -> Tuple[np.ndarray, np.ndarray]:
    r"""
    Get features positional encoding and direction positional encoding.

    Parameters:
        adj: Adjacency matrix of the graph
        pos_encoding_as_features: keyword arguments for function `graph_positional_encoder`
            to generate positional encoding for node features.
        pos_encoding_as_directions: keyword arguments for function `graph_positional_encoder`
            to generate positional encoding for directional features.
    """

    pos_enc_feats_sign_flip, pos_enc_feats_no_flip, pos_enc_dir = None, None, None
    pos_encoding_as_features = {} if pos_encoding_as_features is None else pos_encoding_as_features
    pos_encoding_as_directions = {} if pos_encoding_as_directions is None else pos_encoding_as_directions

    # Get the positional encoding for the features
    if len(pos_encoding_as_features) > 0:
        pos_enc_feats_sign_flip, pos_enc_feats_no_flip = graph_positional_encoder(
            adj, **pos_encoding_as_features
        )

    # Get the positional encoding for the directions
    if len(pos_encoding_as_directions) > 0:
        if pos_encoding_as_directions == pos_encoding_as_features:

            # Concatenate the sign-flip and non-sign-flip positional encodings
            if pos_enc_feats_sign_flip is None:
                pos_enc_dir = pos_enc_feats_no_flip
            elif pos_enc_feats_no_flip is None:
                pos_enc_dir = pos_enc_feats_sign_flip
            else:
                pos_enc_dir = np.concatenate((pos_enc_feats_sign_flip, pos_enc_feats_no_flip), axis=1)

        else:
            pos_enc_dir1, pos_enc_dir2 = graph_positional_encoder(adj, **pos_encoding_as_directions)
            # Concatenate both positional encodings
            if pos_enc_dir1 is None:
                pos_enc_dir = pos_enc_dir2
            elif pos_enc_dir2 is None:
                pos_enc_dir = pos_enc_dir1
            else:
                pos_enc_dir = np.concatenate((pos_enc_dir1, pos_enc_dir2), axis=1)

    return pos_enc_feats_sign_flip, pos_enc_feats_no_flip, pos_enc_dir


def graph_positional_encoder(
    adj: Union[np.ndarray, spmatrix], pos_type: str, num_pos: int, disconnected_comp: bool = True, **kwargs
) -> np.ndarray:
    r"""
    Get a positional encoding that depends on the parameters.

    Parameters:

        adj: Adjacency matrix of the graph

        pos_type: The type of positional encoding to use. Supported types are:

            - laplacian_eigvec: the
            - laplacian_eigvec_eigval

    """

    pos_type = pos_type.lower()
    pos_enc_sign_flip, pos_enc_no_flip = None, None

    if pos_type == "laplacian_eigvec":
        _, eigvecs = compute_laplacian_positional_eigvecs(
            adj=adj, num_pos=num_pos, disconnected_comp=disconnected_comp, **kwargs
        )
        pos_enc_sign_flip = eigvecs

    elif pos_type == "laplacian_eigvec_eigval":
        eigvals_tile, eigvecs = compute_laplacian_positional_eigvecs(
            adj=adj, num_pos=num_pos, disconnected_comp=disconnected_comp, **kwargs
        )
        pos_enc_sign_flip = eigvecs
        pos_enc_no_flip = eigvals_tile
    else:
        raise ValueError(f"Unknown `pos_type`: {pos_type}")

    if pos_enc_sign_flip is not None:
        pos_enc_sign_flip = np.real(pos_enc_sign_flip).astype(np.float32)

    if pos_enc_no_flip is not None:
        pos_enc_no_flip = np.real(pos_enc_no_flip).astype(np.float32)

    return pos_enc_sign_flip, pos_enc_no_flip


# goli/features/properties.py

from typing import Union, List, Callable

import numpy as np
import datamol as dm

from rdkit import Chem
from rdkit.Chem import rdMolDescriptors as rdMD

from mordred import Calculator, descriptors


def get_prop_or_none(prop, n, *args, **kwargs):
    r"""
    return properties. If error, return list of `None` with lenght `n`.
    """
    try:
        return prop(*args, **kwargs)
    except RuntimeError:
        return [None] * n


def get_props_from_mol(mol: Union[Chem.rdchem.Mol, str], properties: Union[List[str], str] = "autocorr3d"):
    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 = Chem.AddHs(mol)  # type: ignore

    if ("descriptors" in properties) or ("all" in properties):
        # Calculate the descriptors of the molecule
        for desc in descriptors.all:
            classes_names.append(desc.__name__.replace("mordred.", ""))
            classes_start_idx.append(len(props))
            calc = Calculator(desc, ignore_3D=True)
            props.extend(calc(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


# goli/features/spectral.py

from typing import Tuple, Union

from scipy.sparse.linalg import eigs
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 goli.utils.tensor import is_dtype_torch_tensor, is_dtype_numpy_array


def compute_laplacian_positional_eigvecs(
    adj: Union[np.ndarray, spmatrix],
    num_pos: int,
    disconnected_comp: bool = True,
    normalization: str = "none",
) -> Tuple[np.ndarray, np.ndarray]:

    # Sparsify the adjacency patrix
    if issparse(adj):
        adj = adj.astype(np.float64)
    else:
        adj = csr_matrix(adj, dtype=np.float64)

    # Compute tha Laplacian, and normalize it
    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)

    components = []
    if disconnected_comp:
        # Get the list of connected components
        components = list(nx.connected_components(nx.from_scipy_sparse_matrix(adj)))

    # Compute the eigenvectors for each connected component, and stack them together
    if len(components) > 1:
        eigvals_tile = np.zeros((L_norm.shape[0], num_pos), dtype=np.float64)
        eigvecs = np.zeros_like(eigvals_tile)
        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)
            eigvecs[comp, :] = np.real(this_eigvecs)
            eigvals_tile[comp, :] = np.real(this_eigvals)
    else:
        eigvals, eigvecs = _get_positional_eigvecs(L, num_pos=num_pos)
        eigvals_tile = np.tile(eigvals, (L_norm.shape[0], 1))

    # Eigenvalues previously set to infinite are now set to 0
    # Any NaN in the eigvals or eigvecs will be set to 0
    eigvecs[~np.isfinite(eigvecs)] = 0.0
    eigvals_tile[~np.isfinite(eigvals_tile)] = 0.0

    return eigvals_tile, eigvecs


def _get_positional_eigvecs(matrix, num_pos: int):

    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, degree_vector=None, normalization: str = None):
    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


# goli/nn/__init__.py

from .architectures import FullDGLNetwork
from .architectures import FullDGLSiameseNetwork
from .architectures import FeedForwardNN
from .architectures import FullDGLMultiTaskNetwork


# goli/nn/architectures.py

from torch import nn
import torch
import dgl
from typing import Iterable, List, Dict, Tuple, Union, Callable, Any, Optional, Type
import inspect

from goli.nn.base_layers import FCLayer, get_activation
from goli.nn.dgl_layers import BaseDGLLayer
from goli.nn.residual_connections import (
    ResidualConnectionBase,
    ResidualConnectionWeighted,
    ResidualConnectionRandom,
)
from goli.nn.dgl_layers.pooling import parse_pooling_layer, VirtualNode


class FeedForwardNN(nn.Module):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        hidden_dims: Union[List[int], int],
        depth: Optional[int] = None,
        activation: Union[str, Callable] = "relu",
        last_activation: Union[str, Callable] = "none",
        dropout: float = 0.0,
        last_dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
        last_normalization: Union[str, Callable] = "none",
        residual_type: str = "none",
        residual_skip_steps: int = 1,
        name: str = "LNN",
        layer_type: Union[str, nn.Module] = "fc",
        layer_kwargs: Optional[Dict] = None,
    ):
        r"""
        A flexible neural network architecture, with variable hidden dimensions,
        support for multiple layer types, and support for different residual
        connections.

        Parameters:

            in_dim:
                Input feature dimensions of the layer

            out_dim:
                Output feature dimensions of the layer

            hidden_dims:
                Either an integer specifying all the hidden dimensions,
                or a list of dimensions in the hidden layers.
                Be careful, the "simple" residual type only supports
                hidden dimensions of the same value.

            depth:
                If `hidden_dims` is an integer, `depth` is 1 + the number of
                hidden layers to use. If `hidden_dims` is a `list`, `depth` must
                be `None`.

            activation:
                activation function to use in the hidden layers.

            last_activation:
                activation function to use in the last layer.

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

            last_dropout:
                The ratio of units to dropout for the last_layer. 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

            last_normalization:
                Whether to use batch normalization in the last layer

            residual_type:
                - "none": No residual connection
                - "simple": Residual connection similar to the ResNet architecture.
                  See class `ResidualConnectionSimple`
                - "weighted": Residual connection similar to the Resnet architecture,
                  but with weights applied before the summation. See class `ResidualConnectionWeighted`
                - "concat": Residual connection where the residual is concatenated instead
                  of being added.
                - "densenet": Residual connection where the residual of all previous layers
                  are concatenated. This leads to a strong increase in the number of parameters
                  if there are multiple hidden layers.

            residual_skip_steps:
                The number of steps to skip between each residual connection.
                If `1`, all the layers are connected. If `2`, half of the
                layers are connected.

            name:
                Name attributed to the current network, for display and printing
                purposes.

            layer_type:
                The type of layers to use in the network.
                Either "fc" as the `FCLayer`, or a class representing the `nn.Module`
                to use.

            layer_kwargs:
                The arguments to be used in the initialization of the layer provided by `layer_type`

        """

        super().__init__()

        # Set the class attributes
        self.in_dim = in_dim
        self.out_dim = out_dim
        if isinstance(hidden_dims, int):
            self.hidden_dims = [hidden_dims] * (depth - 1)
        else:
            self.hidden_dims = list(hidden_dims)
            assert depth is None
        self.depth = len(self.hidden_dims) + 1
        self.activation = get_activation(activation)
        self.last_activation = get_activation(last_activation)
        self.dropout = dropout
        self.last_dropout = last_dropout
        self.normalization = normalization
        self.last_normalization = last_normalization
        self.residual_type = None if residual_type is None else residual_type.lower()
        self.residual_skip_steps = residual_skip_steps
        self.layer_kwargs = layer_kwargs if layer_kwargs is not None else {}
        self.name = name

        # Parse the layer and residuals
        from goli.utils.spaces import LAYERS_DICT, RESIDUALS_DICT

        self.layer_class, self.layer_name = self._parse_class_from_dict(layer_type, LAYERS_DICT)
        self.residual_class, self.residual_name = self._parse_class_from_dict(residual_type, RESIDUALS_DICT)

        self.full_dims = [self.in_dim] + self.hidden_dims + [self.out_dim]
        self._create_layers()
        self._check_bad_arguments()

    def _check_bad_arguments(self):
        r"""
        Raise comprehensive errors if the arguments seem wrong
        """
        if (self.residual_type == "simple") and not (self.hidden_dims[:-1] == self.hidden_dims[1:]):
            raise ValueError(
                f"When using the residual_type={self.residual_type}"
                + f", all elements in the hidden_dims must be equal. Provided:{self.hidden_dims}"
            )

    def _parse_class_from_dict(
        self, name_or_class: Union[type, str], class_dict: Dict[str, type]
    ) -> Tuple[type, str]:
        r"""
        Register the hyperparameters for tracking by Pytorch-lightning
        """
        if isinstance(name_or_class, str):
            obj_name = name_or_class.lower()
            obj_class = class_dict[obj_name]
        elif callable(name_or_class):
            obj_name = str(name_or_class)
            obj_class = name_or_class
        else:
            raise TypeError(f"`name_or_class` must be str or callable, provided: {type(name_or_class)}")

        return obj_class, obj_name

    def _create_residual_connection(self, out_dims: List[int]) -> Tuple[ResidualConnectionBase, List[int]]:
        r"""
        Create the residual connection classes.
        The out_dims is only used if the residual classes requires weights
        """
        if self.residual_class == ResidualConnectionWeighted:
            # if self.residual_class.has_weights:
            residual_layer = self.residual_class(
                skip_steps=self.residual_skip_steps,
                out_dims=out_dims,
                dropout=self.dropout,
                activation=self.activation,
                normalization=self.normalization,
                bias=False,
            )
        elif self.residual_class == ResidualConnectionRandom:
            residual_layer = self.residual_class(
                out_dims=out_dims,
                skip_steps=self.residual_skip_steps,
            )
        else:
            residual_layer = self.residual_class(skip_steps=self.residual_skip_steps)

        residual_out_dims = residual_layer.get_true_out_dims(self.full_dims[1:])

        return residual_layer, residual_out_dims

    def _create_layers(self):
        r"""
        Create all the necessary layers for the network.
        It's a bit complicated to explain what's going on in this function,
        but it must manage the varying features sizes caused by:

        - The presence of different types of residual connections
        """

        self.residual_layer, residual_out_dims = self._create_residual_connection(out_dims=self.full_dims[1:])

        # Create a ModuleList of the GNN layers
        self.layers = nn.ModuleList()
        this_in_dim = self.full_dims[0]
        this_activation = self.activation
        this_norm = self.normalization
        this_dropout = self.dropout

        for ii in range(self.depth):
            this_out_dim = self.full_dims[ii + 1]
            if ii == self.depth - 1:
                this_activation = self.last_activation
                this_norm = self.last_normalization
                this_dropout = self.last_dropout

            # Create the layer
            self.layers.append(
                self.layer_class(
                    in_dim=this_in_dim,
                    out_dim=this_out_dim,
                    activation=this_activation,
                    dropout=this_dropout,
                    normalization=this_norm,
                    **self.layer_kwargs,
                )
            )

            if ii < len(residual_out_dims):
                this_in_dim = residual_out_dims[ii]

    def forward(self, h: torch.Tensor) -> torch.Tensor:
        r"""
        Apply the neural network on the input features.

        Parameters:

            h: `torch.Tensor[..., Din]`:
                Input feature tensor, before the network.
                `Din` is the number of input features

        Returns:

            `torch.Tensor[..., Dout]`:
                Output feature tensor, after the network.
                `Dout` is the number of output features

        """
        h_prev = None
        for ii, layer in enumerate(self.layers):
            h = layer.forward(h)
            if ii < len(self.layers) - 1:
                h, h_prev = self.residual_layer.forward(h, h_prev, step_idx=ii)

        return h

    def __repr__(self):
        r"""
        Controls how the class is printed
        """
        class_str = f"{self.name}(depth={self.depth}, {self.residual_layer})\n    "
        layer_str = f"[{self.layer_class.__name__}[{' -> '.join(map(str, self.full_dims))}]"

        return class_str + layer_str


class FeedForwardDGL(FeedForwardNN):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        hidden_dims: List[int],
        depth: Optional[int] = None,
        activation: Union[str, Callable] = "relu",
        last_activation: Union[str, Callable] = "none",
        dropout: float = 0.0,
        last_dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
        last_normalization: Union[str, Callable] = "none",
        residual_type: str = "none",
        residual_skip_steps: int = 1,
        in_dim_edges: int = 0,
        hidden_dims_edges: List[int] = [],
        pooling: Union[List[str], List[Callable]] = ["sum"],
        name: str = "GNN",
        layer_type: Union[str, nn.Module] = "gcn",
        layer_kwargs: Optional[Dict] = None,
        virtual_node: str = "none",
    ):
        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 DGL-based graph neural networks
        layers. Any layer must inherit from `goli.nn.dgl_layers.BaseDGLLayer`.

        Parameters:

            in_dim:
                Input feature dimensions of the layer

            out_dim:
                Output feature dimensions of the layer

            hidden_dims:
                List of dimensions in the hidden layers.
                Be careful, the "simple" residual type only supports
                hidden dimensions of the same value.

            depth:
                If `hidden_dims` is an integer, `depth` is 1 + the number of
                hidden layers to use. If `hidden_dims` is a `list`, `depth` must
                be `None`.

            activation:
                activation function to use in the hidden layers.

            last_activation:
                activation function to use in the last layer.

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

            last_dropout:
                The ratio of units to dropout for the last layer. 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

            last_normalization:
                Whether to use batch normalization in the last layer

            residual_type:
                - "none": No residual connection
                - "simple": Residual connection similar to the ResNet architecture.
                  See class `ResidualConnectionSimple`
                - "weighted": Residual connection similar to the Resnet architecture,
                  but with weights applied before the summation. See class `ResidualConnectionWeighted`
                - "concat": Residual connection where the residual is concatenated instead
                  of being added.
                - "densenet": Residual connection where the residual of all previous layers
                  are concatenated. This leads to a strong increase in the number of parameters
                  if there are multiple hidden layers.

            residual_skip_steps:
                The number of steps to skip between each residual connection.
                If `1`, all the layers are connected. If `2`, half of the
                layers are connected.

            in_dim_edges:
                Input edge-feature dimensions of the network. Keep at 0 if not using
                edge features, or if the layer doesn't support edges.

            hidden_dims_edges:
                Hidden dimensions for the edges. Most models don't support it, so it
                should only be used for those that do, i.e. `GatedGCNLayer`

            pooling:
                The pooling types to use. Multiple pooling can be used, and their
                results will be concatenated.
                For node feature predictions, use `["none"]`.
                For graph feature predictions see `goli.nn.dgl_layers.pooling.parse_pooling_layer`.
                The list must either contain Callables, or the string below

                - "none": No pooling is applied
                - "sum": `SumPooling`
                - "mean": `MeanPooling`
                - "max": `MaxPooling`
                - "min": `MinPooling`
                - "std": `StdPooling`
                - "s2s": `Set2Set`

            name:
                Name attributed to the current network, for display and printing
                purposes.

            layer_type:
                The type of layers to use in the network.
                Either a class that inherits from `goli.nn.dgl_layers.BaseDGLLayer`,
                or one of the following strings

                - "gcn": `GCNLayer`
                - "gin": `GINLayer`
                - "gat": `GATLayer`
                - "gated-gcn": `GatedGCNLayer`
                - "pna-conv": `PNAConvolutionalLayer`
                - "pna-msgpass": `PNAMessagePassingLayer`
                - "dgn-conv": `DGNConvolutionalLayer`
                - "dgn-msgpass": `DGNMessagePassingLayer`

            layer_kwargs:
                The arguments to be used in the initialization of the layer provided by `layer_type`

            virtual_node:
                A string associated to the type of virtual node to use,
                either `None`, "none", "mean", "sum", "max", "logsum".
                See `goli.nn.dgl_layers.VirtualNode`.

                The virtual node will not use any residual connection if `residual_type`
                is "none". Otherwise, it will use a simple ResNet like residual
                connection.

        """

        # Initialize the additional attributes
        self.in_dim_edges = in_dim_edges
        if isinstance(hidden_dims_edges, int):
            self.hidden_dims_edges = [hidden_dims_edges] * (depth - 1)
        elif len(hidden_dims_edges) == 0:
            self.hidden_dims_edges = []
        else:
            self.hidden_dims_edges = list(hidden_dims_edges)
            assert depth is None
        self.full_dims_edges = None
        if len(self.hidden_dims_edges) > 0:
            self.full_dims_edges = [self.in_dim_edges] + self.hidden_dims_edges + [self.hidden_dims_edges[-1]]

        self.virtual_node = virtual_node.lower() if virtual_node is not None else "none"
        self.pooling = pooling

        # Initialize the parent `FeedForwardNN`
        super().__init__(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            depth=depth,
            activation=activation,
            last_activation=last_activation,
            normalization=normalization,
            last_normalization=last_normalization,
            residual_type=residual_type,
            residual_skip_steps=residual_skip_steps,
            name=name,
            layer_type=layer_type,
            dropout=dropout,
            last_dropout=last_dropout,
            layer_kwargs=layer_kwargs,
        )

    def _check_bad_arguments(self):
        r"""
        Raise comprehensive errors if the arguments seem wrong
        """
        super()._check_bad_arguments()
        if (
            (self.in_dim_edges > 0) or (self.full_dims_edges is not None)
        ) and not self.layer_class.layer_supports_edges:
            raise ValueError(f"Cannot use edge features with class `{self.layer_class}`")

    def _create_layers(self):
        r"""
        Create all the necessary layers for the network.
        It's a bit complicated to explain what's going on in this function,
        but it must manage the varying features sizes caused by:

        - The presence of different types of residual connections
        - The presence or absence of edges
        - The output dimensions varying for different networks i.e. `GatLayer` outputs different feature sizes according to the number of heads
        - The presence or absence of virtual nodes
        - The different possible pooling, and the concatenation of multiple pooling together.
        """

        residual_layer_temp, residual_out_dims = self._create_residual_connection(out_dims=self.full_dims[1:])

        # Create a ModuleList of the GNN layers
        self.layers = nn.ModuleList()
        self.virtual_node_layers = nn.ModuleList()
        this_in_dim = self.full_dims[0]
        this_activation = self.activation
        this_norm = self.normalization
        this_dropout = self.dropout

        # Find the appropriate edge dimensions, depending if edges are used,
        # And if the residual is required for the edges
        this_in_dim_edges, this_out_dim_edges = None, None
        if self.full_dims_edges is not None:
            this_in_dim_edges, this_out_dim_edges = self.full_dims_edges[0:2]
            residual_out_dims_edges = residual_layer_temp.get_true_out_dims(self.full_dims_edges[1:])
        elif self.in_dim_edges > 0:
            this_in_dim_edges = self.in_dim_edges
        layer_out_dims_edges = []

        # Create all the layers in a loop
        for ii in range(self.depth):
            this_out_dim = self.full_dims[ii + 1]

            if ii == self.depth - 1:
                this_activation = self.last_activation

            # Find the edge key-word arguments depending on the layer type and residual connection
            this_edge_kwargs = {}
            if self.layer_class.layer_supports_edges and self.in_dim_edges > 0:
                this_edge_kwargs["in_dim_edges"] = this_in_dim_edges
                if "out_dim_edges" in inspect.signature(self.layer_class.__init__).parameters.keys():
                    layer_out_dims_edges.append(self.full_dims_edges[ii + 1])
                    this_edge_kwargs["out_dim_edges"] = layer_out_dims_edges[-1]

            # Create the GNN layer
            self.layers.append(
                self.layer_class(
                    in_dim=this_in_dim,
                    out_dim=this_out_dim,
                    activation=this_activation,
                    dropout=this_dropout,
                    normalization=this_norm,
                    **self.layer_kwargs,
                    **this_edge_kwargs,
                )
            )

            # Create the Virtual Node layer, except at the last layer
            if ii < len(residual_out_dims):
                self.virtual_node_layers.append(
                    VirtualNode(
                        dim=this_out_dim * self.layers[-1].out_dim_factor,
                        activation=this_activation,
                        dropout=this_dropout,
                        normalization=this_norm,
                        bias=True,
                        vn_type=self.virtual_node,
                        residual=self.residual_type is not None,
                    )
                )

            # Get the true input dimension of the next layer,
            # by factoring both the residual connection and GNN layer type
            if ii < len(residual_out_dims):
                this_in_dim = residual_out_dims[ii] * self.layers[ii - 1].out_dim_factor
                if self.full_dims_edges is not None:
                    this_in_dim_edges = residual_out_dims_edges[ii] * self.layers[ii - 1].out_dim_factor

        layer_out_dims = [layer.out_dim_factor * layer.out_dim for layer in self.layers]

        # Initialize residual and pooling layers
        self.residual_layer, _ = self._create_residual_connection(out_dims=layer_out_dims)
        if len(layer_out_dims_edges) > 0:
            self.residual_edges_layer, _ = self._create_residual_connection(out_dims=layer_out_dims_edges)
        else:
            self.residual_edges_layer = None
        self.global_pool_layer, out_pool_dim = parse_pooling_layer(layer_out_dims[-1], self.pooling)

        # Output linear layer
        self.out_linear = FCLayer(
            in_dim=out_pool_dim,
            out_dim=self.out_dim,
            activation="none",
            dropout=self.dropout,
            normalization=self.normalization,
        )

    def _pool_layer_forward(self, g, h):
        r"""
        Apply the graph pooling layer, followed by the linear output layer.

        Parameters:

            g: dgl.DGLGraph
                graph on which the convolution is done

            h (torch.Tensor[..., N, Din]):
                Node feature tensor, before convolution.
                `N` is the number of nodes, `Din` is the output size of the last DGL layer

        Returns:

            torch.Tensor[..., M, Din] or torch.Tensor[..., N, Din]:
                Node feature tensor, after convolution.
                `N` is the number of nodes, `M` is the number of graphs, `Dout` is the output dimension ``self.out_dim``
                If the pooling is `None`, the dimension is `N`, otherwise it is `M`

        """

        if len(self.global_pool_layer) > 0:
            pooled_h = []
            for this_pool in self.global_pool_layer:
                pooled_h.append(this_pool(g, h))
            pooled_h = torch.cat(pooled_h, dim=-1)
        else:
            pooled_h = h

        pooled_h = self.out_linear(pooled_h)

        return pooled_h

    def _dgl_layer_forward(
        self,
        layer: BaseDGLLayer,
        g: dgl.DGLGraph,
        h: torch.Tensor,
        e: Union[torch.Tensor, None],
        h_prev: Union[torch.Tensor, None],
        e_prev: Union[torch.Tensor, None],
        step_idx: int,
    ) -> Tuple[torch.Tensor, Union[torch.Tensor, None], Union[torch.Tensor, None], Union[torch.Tensor, None]]:
        r"""
        Apply the *i-th* DGL graph layer, where *i* is the index given by `step_idx`.
        The layer is applied differently depending if there are edge features or not.

        Then, the residual is also applied on both the features and the edges (if applicable)

        Parameters:

            layer:
                The DGL layer used for the convolution

            g:
                graph on which the convolution is done

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

            e (torch.Tensor[..., N, Ein]):
                Edge feature tensor, before convolution.
                `N` is the number of nodes, `Ein` is the input edge features

            h_prev:
                Node feature of the previous residual connection, or `None`

            e_prev:
                Edge feature of the previous residual connection, or `None`

            step_idx:
                The current step idx in the forward loop

        Returns:

            h (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

            e:
                Edge feature tensor, after convolution and residual.
                `N` is the number of nodes, `Ein` is the input edge features

            h_prev:
                Node feature tensor to be used at the next residual connection, or `None`

            e_prev:
                Edge feature tensor to be used at the next residual connection, or `None`

        """

        # Apply the GNN layer with the right inputs/outputs
        if layer.layer_inputs_edges and layer.layer_outputs_edges:
            h, e = layer(g=g, h=h, e=e)
        elif layer.layer_inputs_edges:
            h = layer(g=g, h=h, e=e)
        elif layer.layer_outputs_edges:
            h, e = layer(g=g, h=h)
        else:
            h = layer(g=g, h=h)

        # Apply the residual layers on the features and edges (if applicable)
        if step_idx < len(self.layers) - 1:
            h, h_prev = self.residual_layer.forward(h, h_prev, step_idx=step_idx)
            if (self.residual_edges_layer is not None) and (layer.layer_outputs_edges):
                e, e_prev = self.residual_edges_layer.forward(e, e_prev, step_idx=step_idx)

        return h, e, h_prev, e_prev

    def _virtual_node_forward(
        self, g: dgl.DGLGraph, h: torch.Tensor, vn_h: torch.Tensor, step_idx: int
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        r"""
        Apply the *i-th* virtual node layer, where *i* is the index given by `step_idx`.

        Parameters:

            g:
                graph on which the convolution is done

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

            vn_h (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

            step_idx:
                The current step idx in the forward loop

        Returns:

            `h = 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_h = 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

        """

        if step_idx == 0:
            vn_h = 0.0
        if step_idx < len(self.virtual_node_layers):
            h, vn_h = self.virtual_node_layers[step_idx].forward(g=g, h=h, vn_h=vn_h)

        return h, vn_h

    def forward(self, g: dgl.DGLGraph) -> torch.Tensor:
        r"""
        Apply the full graph neural network on the input graph and node features.

        Parameters:

            g:
                graph on which the convolution is done.
                Must contain the following elements:

                - `g.ndata["h"]`: `torch.Tensor[..., N, Din]`.
                  Input node feature tensor, before the network.
                  `N` is the number of nodes, `Din` is the input features

                - `g.edata["e"]`: `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`.

        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.out_dim``
                If the `self.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`

        """

        # Get node and edge features
        h = g.ndata["h"]
        e = g.edata["e"] if (self.in_dim_edges > 0) else None

        # Initialize values of the residuals and virtual node
        h_prev = None
        e_prev = None
        vn_h = 0

        # Apply the forward loop of the layers, residuals and virtual nodes
        for ii, layer in enumerate(self.layers):
            h, e, h_prev, e_prev = self._dgl_layer_forward(
                layer=layer, g=g, h=h, e=e, h_prev=h_prev, e_prev=e_prev, step_idx=ii
            )
            h, vn_h = self._virtual_node_forward(g=g, h=h, vn_h=vn_h, step_idx=ii)

        pooled_h = self._pool_layer_forward(g=g, h=h)

        return pooled_h

    def __repr__(self):
        r"""
        Controls how the class is printed
        """
        class_str = f"{self.name}(depth={self.depth}, {self.residual_layer})\n    "
        layer_str = f"{self.layer_class.__name__}[{' -> '.join(map(str, self.full_dims))}]\n    "
        pool_str = f"-> Pooling({self.pooling})"
        out_str = f" -> {self.out_linear}"

        return class_str + layer_str + pool_str + out_str


class FullDGLNetwork(nn.Module):
    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,
        post_nn_kwargs: Optional[Dict[str, Any]] = None,
        num_inference_to_average: int = 1,
        name: str = "DGL_GNN",
    ):
        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 `FeedForwardDGL`.
                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 post_nn_kwargs["in_dim"]

            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.

            post_nn_kwargs:
                key-word arguments to use for the initialization of the post-processing
                MLP network after the GNN, using the class `FeedForwardNN`.
                If `None`, there won't be a post-processing MLP.

            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.

            name:
                Name attributed to the current network, for display and printing
                purposes.
        """

        super().__init__()
        self.name = name
        self.num_inference_to_average = num_inference_to_average
        self._concat_last_layers = None

        # Initialize the networks
        self.pre_nn, self.post_nn, self.pre_nn_edges = None, None, None
        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"]

        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"]

        name = gnn_kwargs.pop("name", "GNN")
        self.gnn = FeedForwardDGL(**gnn_kwargs, name=name)
        next_in_dim = self.gnn.out_dim

        if post_nn_kwargs is not None:
            name = post_nn_kwargs.pop("name", "post-NN")
            post_nn_kwargs.setdefault("in_dim", next_in_dim)
            self.post_nn = FeedForwardNN(**post_nn_kwargs, name=name)
            assert next_in_dim == self.post_nn.in_dim

    def _check_bad_arguments(self):
        r"""
        Raise comprehensive errors if the arguments seem wrong
        """
        if self.pre_nn is not None:
            if self.pre_nn["out_dim"] != self.gnn["in_dim"]:
                raise ValueError(
                    f"`self.pre_nn.out_dim` must be equal to `self.gnn.in_dim`."
                    + 'Provided" {self.pre_nn.out_dim} and {self.gnn.in_dim}'
                )

        if self.post_nn is not None:
            if self.gnn["out_dim"] != self.post_nn["in_dim"]:
                raise ValueError(
                    f"`self.gnn.out_dim` must be equal to `self.post_nn.in_dim`."
                    + 'Provided" {self.gnn.out_dim} and {self.post_nn.in_dim}'
                )

    def drop_post_nn_layers(self, num_layers_to_drop: int) -> None:
        r"""
        Remove the last layers of the model. Useful for Transfer Learning.

        Parameters:
            num_layers_to_drop: The number of layers to drop from the `self.post_nn` network.

        """

        assert num_layers_to_drop >= 0
        assert num_layers_to_drop <= len(self.post_nn.layers)

        if num_layers_to_drop > 0:
            self.post_nn.layers = self.post_nn.layers[:-num_layers_to_drop]

    def extend_post_nn_layers(self, layers: nn.ModuleList):
        r"""
        Add layers at the end of the model. Useful for Transfer Learning.

        Parameters:
            layers: A ModuleList of all the layers to extend

        """

        assert isinstance(layers, nn.ModuleList)
        if len(self.post_nn.layers) > 0:
            assert layers[0].in_dim == self.post_nn.layers.out_dim[-1]

        self.post_nn.extend(layers)

    def forward(self, g: dgl.DGLGraph) -> torch.Tensor:
        r"""
        Apply the pre-processing neural network, the graph neural network,
        and the post-processing neural network on the graph features.

        Parameters:

            g:
                graph on which the convolution is done.
                Must contain the following elements:

                - `g.ndata["h"]`: `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``

                - `g.edata["e"]`: `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`.

        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.post_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`

        """

        if self.training:
            return self._forward(g, flip_pos_enc="random")
        else:
            # If in test mode, try different sign flips according to `self.num_inference_to_average` and average them together
            h = [self._forward(g, flip_pos_enc="no-flip")]
            if ("pos_enc_feats_sign_flip" in g.ndata.keys()) and self.num_inference_to_average > 1:
                h.append(self._forward(g, flip_pos_enc="sign-flip"))
                for _ in range(2, self.num_inference_to_average):
                    h.append(self._forward(g, flip_pos_enc="random"))
            return torch.mean(torch.stack(h, dim=-1), dim=-1)

    def _forward(self, g: dgl.DGLGraph, flip_pos_enc: str) -> torch.Tensor:
        # Get the node features and positional embedding
        h = g.ndata["feat"]
        if "pos_enc_feats_sign_flip" in g.ndata.keys():
            pos_enc = g.ndata["pos_enc_feats_sign_flip"]
            if flip_pos_enc == "random":
                rand_sign_shape = ([1] * (pos_enc.ndim - 1)) + [pos_enc.shape[-1]]
                rand_sign = torch.sign(torch.randn(rand_sign_shape, dtype=h.dtype, device=h.device))
                pos_enc = pos_enc * rand_sign
            elif flip_pos_enc == "no-flip":
                pass
            elif flip_pos_enc == "sign-flip":
                pos_enc = -pos_enc
            h = torch.cat((h, pos_enc), dim=-1)
        if "pos_enc_feats_no_flip" in g.ndata.keys():
            pos_enc = g.ndata["pos_enc_feats_no_flip"]
            h = torch.cat((h, pos_enc), dim=-1)

        g.ndata["h"] = h

        if "feat" in g.edata.keys():
            g.edata["e"] = g.edata["feat"]

        # Run the pre-processing network on node features
        if self.pre_nn is not None:
            h = g.ndata["h"]
            h = self.pre_nn.forward(h)
            g.ndata["h"] = h

        # 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.edata["e"]
            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.edata["e"] = e

        # Run the graph neural network
        h = self.gnn.forward(g)

        # Run the output network
        if self.post_nn is not None:
            if self.concat_last_layers is None:
                h = self.post_nn.forward(h)
            else:
                # Concatenate the output of the last layers according to `self._concat_last_layers``.
                # Useful for generating fingerprints
                h = [h]
                for ii in range(len(self.post_nn.layers)):
                    h.insert(0, self.post_nn.layers[ii].forward(h[0]))  # Append in reverse order
                h = torch.cat([h[ii] for ii in self._concat_last_layers], dim=-1)

        return h

    @property
    def concat_last_layers(self) -> Union[Type[None], Iterable[int]]:
        """
        Property to control the output of the `self.forward`.
        If set to a list of integer, the `forward` function will
        concatenate the output of different layers.

        If set to `None`, the output of the last layer is returned.

        NOTE: The indexes are inverted. 0 is the last layer, 1 is the second last, etc.
        """
        return self._concat_last_layers

    @concat_last_layers.setter
    def concat_last_layers(self, value: Union[Type[None], int, Iterable[int]]) -> None:
        """
        Set the property to control the output of the `self.forward`.
        If set to a list of integer, the `forward` function will
        concatenate the output of different layers.
        If a single integer is provided, it will output that specific layer.

        If set to `None`, the output of the last layer is returned.

        NOTE: The indexes are inverted. 0 is the last layer, 1 is the second last, etc.

        Parameters:
            value: Output layers to concatenate, in reverse order (`0` is the last layer)
        """
        if (value is not None) and not isinstance(value, Iterable):
            value = [value]
        self._concat_last_layers = value

    def __repr__(self):
        r"""
        Controls how the class is printed
        """
        pre_nn_str, post_nn_str, pre_nn_edges_str = "", "", ""
        if self.pre_nn is not None:
            pre_nn_str = self.pre_nn.__repr__() + "\n\n"
        if self.pre_nn_edges is not None:
            pre_nn_edges_str = self.pre_nn_edges.__repr__() + "\n\n"
        gnn_str = self.gnn.__repr__() + "\n\n"
        if self.post_nn is not None:
            post_nn_str = self.post_nn.__repr__()

        child_str = "    " + pre_nn_str + pre_nn_edges_str + gnn_str + post_nn_str
        child_str = "    ".join(child_str.splitlines(True))

        full_str = self.name + "\n" + "-" * (len(self.name) + 2) + "\n" + child_str

        return full_str

    @property
    def in_dim(self):
        r"""
        Returns the input dimension of the network
        """
        if self.pre_nn is not None:
            return self.pre_nn.in_dim
        else:
            return self.gnn.in_dim

    @property
    def out_dim(self):
        r"""
        Returns the output dimension of the network
        """
        if self.pre_nn is not None:
            return self.post_nn.out_dim
        else:
            return self.gnn.out_dim

    @property
    def in_dim_edges(self):
        r"""
        Returns the input edge dimension of the network
        """
        return self.gnn.in_dim_edges


class FullDGLSiameseNetwork(FullDGLNetwork):
    def __init__(self, pre_nn_kwargs, gnn_kwargs, post_nn_kwargs, dist_method, name="Siamese_DGL_GNN"):

        # Initialize the parent nn.Module
        super().__init__(
            pre_nn_kwargs=pre_nn_kwargs,
            gnn_kwargs=gnn_kwargs,
            post_nn_kwargs=post_nn_kwargs,
            name=name,
        )

        self.dist_method = dist_method.lower()

    def forward(self, graphs):
        graph_1, graph_2 = graphs

        out_1 = super().forward(graph_1)
        out_2 = super().forward(graph_2)

        if self.dist_method == "manhattan":
            # Normalized L1 distance
            out_1 = out_1 / torch.mean(out_1.abs(), dim=-1, keepdim=True)
            out_2 = out_2 / torch.mean(out_2.abs(), dim=-1, keepdim=True)
            dist = torch.abs(out_1 - out_2)
            out = torch.mean(dist, dim=-1)

        elif self.dist_method == "euclidean":
            # Normalized Euclidean distance
            out_1 = out_1 / torch.norm(out_1, dim=-1, keepdim=True)
            out_2 = out_2 / torch.norm(out_2, dim=-1, keepdim=True)
            out = torch.norm(out_1 - out_2, dim=-1)
        elif self.dist_method == "cosine":
            # Cosine distance
            out = torch.sum(out_1 * out_2, dim=-1) / (torch.norm(out_1, dim=-1) * torch.norm(out_2, dim=-1))
        else:
            raise ValueError(f"Unsupported `dist_method`: {self.dist_method}")

        return out


class TaskHeadParams():
    def __init__(
        self,
        task_name: str,
        #in_dim: int,                           # We don't need to specify the input dimension since the TaskHeads class will know which dimension to use.
        out_dim: int,
        hidden_dims: Union[List[int], int],     # Should this only be List? See FeedForwardNN vs. FeedForwardDGL
        depth: Optional[int] = None,
        activation: Union[str, Callable] = "relu",
        last_activation: Union[str, Callable] = "none",
        dropout: float = 0.0,
        last_dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
        last_normalization: Union[str, Callable] = "none",
        residual_type: str = "none",
        residual_skip_steps: int = 1,
        name: str = "LNN",
        layer_type: Union[str, nn.Module] = "fc",
        layer_kwargs: Optional[Dict] = None,
    ):
        r"""
        This class holds all the information to instatiate a TaskHead.
        """

        self.task_name = task_name
        #self.in_dim = in_dim
        self.nn_params = {}
        self.nn_params["out_dim"] = out_dim
        self.nn_params["hidden_dims"] = hidden_dims
        self.nn_params["depth"] = depth
        self.nn_params["activation"] = activation
        self.nn_params["last_activation"] = last_activation
        self.nn_params["dropout"] = dropout
        self.nn_params["last_dropout"] = last_dropout
        self.nn_params["normalization"] = normalization
        self.nn_params["last_normalization"] = last_normalization
        self.nn_params["residual_type"] = residual_type
        self.nn_params["residual_skip_steps"] = residual_skip_steps
        self.nn_params["name"] = name
        self.nn_params["layer_type"] = layer_type
        self.nn_params["layer_kwargs"] = layer_kwargs

class TaskHead(FeedForwardNN):
    def __init__(
        self,
        task_name: str,                         # The name matters for per-task analysis
        in_dim: int,
        out_dim: int,
        hidden_dims: Union[List[int], int],     # Should this only be List? See FeedForwardNN vs. FeedForwardDGL
        depth: Optional[int] = None,
        activation: Union[str, Callable] = "relu",
        last_activation: Union[str, Callable] = "none",
        dropout: float = 0.0,
        last_dropout: float = 0.0,
        normalization: Union[str, Callable] = "none",
        last_normalization: Union[str, Callable] = "none",
        residual_type: str = "none",
        residual_skip_steps: int = 1,
        name: str = "LNN",
        layer_type: Union[str, nn.Module] = "fc",
        layer_kwargs: Optional[Dict] = None,
    ):
        r"""
        This class instantiates a task head, and it is identical to the FeedForwardNN class with the addition of a `task_name` attribute.

        Parameters:

            task_name:
                The name of the task for which the current output head performs predictions.

            in_dim:
                Input feature dimensions of the layer

            out_dim:
                Output feature dimensions of the layer

            hidden_dims:
                Either an integer specifying all the hidden dimensions,
                or a list of dimensions in the hidden layers.
                Be careful, the "simple" residual type only supports
                hidden dimensions of the same value.

            depth:
                If `hidden_dims` is an integer, `depth` is 1 + the number of
                hidden layers to use. If `hidden_dims` is a `list`, `depth` must
                be `None`.

            activation:
                activation function to use in the hidden layers.

            last_activation:
                activation function to use in the last layer.

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

            last_dropout:
                The ratio of units to dropout for the last_layer. 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

            last_normalization:
                Whether to use batch normalization in the last layer

            residual_type:
                - "none": No residual connection
                - "simple": Residual connection similar to the ResNet architecture.
                  See class `ResidualConnectionSimple`
                - "weighted": Residual connection similar to the Resnet architecture,
                  but with weights applied before the summation. See class `ResidualConnectionWeighted`
                - "concat": Residual connection where the residual is concatenated instead
                  of being added.
                - "densenet": Residual connection where the residual of all previous layers
                  are concatenated. This leads to a strong increase in the number of parameters
                  if there are multiple hidden layers.

            residual_skip_steps:
                The number of steps to skip between each residual connection.
                If `1`, all the layers are connected. If `2`, half of the
                layers are connected.

            name:
                Name attributed to the current network, for display and printing
                purposes.

            layer_type:
                The type of layers to use in the network.
                Either "fc" as the `FCLayer`, or a class representing the `nn.Module`
                to use.

            layer_kwargs:
                The arguments to be used in the initialization of the layer provided by `layer_type`  
        """

        self.task_name = task_name

        # Initialize the FeedForwardNN
        super().__init__(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            depth=depth,
            activation=activation,
            last_activation=last_activation,
            normalization=normalization,
            last_normalization=last_normalization,
            residual_type=residual_type,
            residual_skip_steps=residual_skip_steps,
            name=name,
            layer_type=layer_type,
            dropout=dropout,
            last_dropout=last_dropout,
            layer_kwargs=layer_kwargs,
        )

class TaskHeads(nn.Module):
    def __init__(
        self,
        in_dim: int,
        task_heads_params: List[Type[TaskHeadParams]],
    ):
        r"""
        Class that groups all multi-task output heads together to provide the task-specific outputs.

        Parameters:

            in_dim:
                Input feature dimensions of the layer

            task_heads_params:
                This argument is of type List[TaskHeadParams]. Each argument is used to
                initialize a task-specific MLP.
        """

        super().__init__()

        self.in_dim = in_dim

        self.task_heads = nn.ModuleDict()
        for head_params in task_heads_params:
            self.task_heads[head_params.task_name] = TaskHead(
                task_name=head_params.task_name,
                in_dim=self.in_dim,
                **head_params.nn_params
            )
    
    # Return a dictionary: Dict[task_name, Tensor]
    def forward(self, h: torch.Tensor):
        task_head_outputs = {}

        for task, head in self.task_heads.items():
            task_head_outputs[task] = head.forward(h)

        return task_head_outputs

class FullDGLMultiTaskNetwork(nn.Module):
    def __init__(
        self,
        task_heads_kwargs: List[Type[TaskHeadParams]],
        gnn_kwargs: Dict[str, Any],
        pre_nn_kwargs: Optional[Dict[str, Any]] = None,
        pre_nn_edges_kwargs: Optional[Dict[str, Any]] = None,
        post_nn_kwargs: Optional[Dict[str, Any]] = None,
        num_inference_to_average: int = 1,                          # Should this be changed?
        name: str = "MTL_DGL_GNN",
    ):
        r"""
        Class that allows to implement a full multi-task graph neural network architecture,
        including the pre-processing MLP, post-processing MLP and the task-specific heads.

        In this model, the tasks share a full DGL network as a "trunk", and additionally have task-specific MLPs. 
        Each molecular graph is associated with a variety of tasks, so the network outputs the task-specific preedictions for a graph.

        Parameters:

            task_heads_kwargs:
                This argument is of type List[TaskHeadParams]. Each argument is used to
                initialize a task-specific MLP.

            gnn_kwargs:
                key-word arguments to use for the initialization of the pre-processing
                GNN network using the class `FeedForwardDGL`.
                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 post_nn_kwargs["in_dim"]

            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.

            post_nn_kwargs:
                key-word arguments to use for the initialization of the post-processing
                MLP network after the GNN, using the class `FeedForwardNN`.
                If `None`, there won't be a post-processing MLP.

            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.

            name:
                Name attributed to the current network, for display and printing
                purposes.
        """

        super().__init__()
        self.name = name

        # Initialize the FullDGLNetwork
        self.full_dgl_network = FullDGLNetwork(
            gnn_kwargs=gnn_kwargs,
            pre_nn_kwargs=pre_nn_kwargs,
            pre_nn_edges_kwargs=pre_nn_edges_kwargs,
            post_nn_kwargs=post_nn_kwargs,
            num_inference_to_average=num_inference_to_average,
        )
        
        # Initialize TaskHeads
        if self.full_dgl_network.post_nn is not None:                   # FullDGLNetwork uses .pre_nn which seems wrong...
            task_heads_in_dim = self.full_dgl_network.post_nn.out_dim
        else:
            task_heads_in_dim =  self.full_dgl_network.gnn.out_dim
        self.task_heads = TaskHeads(
            #in_dim=self.full_dgl_network.out_dim(),
            in_dim=task_heads_in_dim,
            task_heads_params=task_heads_kwargs
        )

        self.model = nn.Sequential()
        self.model.add_module(
            'FullDGLNetwork',
            self.full_dgl_network
        )
        self.model.add_module(
            'TaskHeads',
            self.task_heads
        )

    def forward(self, g: dgl.DGLGraph):
        h_out_full_dgl_network = self.full_dgl_network.forward(g)
        task_outputs = self.task_heads.forward(h_out_full_dgl_network)

        return task_outputs

# goli/nn/base_layers.py

import torch
import torch.nn as nn
import torch.nn.functional as F

from typing import Union, Callable

SUPPORTED_ACTIVATION_MAP = {"ReLU", "Sigmoid", "Tanh", "ELU", "SELU", "GLU", "LeakyReLU", "Softplus", "None"}


def get_activation(activation: Union[type(None), str, Callable]) -> Union[type(None), 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", "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), "Unhandled activation function"
    activation = activation[0]

    return vars(torch.nn.modules.activation)[activation]()


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: Union[type(None), Callable] = None,
    ):

        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}}$$

        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"]
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.bias = bias
        self.linear = nn.Linear(in_dim, out_dim, bias=bias)
        self.dropout = None
        self.normalization = self._parse_norm(normalization)

        if dropout:
            self.dropout = nn.Dropout(p=dropout)
        self.activation = get_activation(activation)
        self.init_fn = nn.init.xavier_uniform_

        self.reset_parameters()

    def _parse_norm(self, normalization):

        parsed_norm = None
        if normalization is None or normalization == "none":
            pass
        elif callable(normalization):
            parsed_norm = normalization
        elif normalization == "batch_norm":
            parsed_norm = nn.BatchNorm1d(self.out_dim)
        elif normalization == "layer_norm":
            parsed_norm = nn.LayerNorm(self.out_dim)
        else:
            raise ValueError(
                f"Undefined normalization `{normalization}`, must be `None`, `Callable`, 'batch_norm', 'layer_norm', 'none'"
            )
        return parsed_norm

    def reset_parameters(self, init_fn=None):
        init_fn = init_fn or self.init_fn
        if init_fn is not None:
            init_fn(self.linear.weight, 1 / self.in_dim)
        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)

        return h

    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_dim: int,
        out_dim: int,
        layers: int,
        activation: Union[str, Callable] = "relu",
        last_activation: Union[str, Callable] = "none",
        dropout=0.0,
        normalization="none",
        last_normalization="none",
    ):
        r"""
        Simple multi-layer perceptron, built of a series of FCLayers

        Parameters:
            in_dim:
                Input dimension of the MLP
            hidden_dim:
                Hidden dimension of the MLP. All hidden dimensions will have
                the same number of parameters
            out_dim:
                Output dimension of the MLP.
            layers:
                Number of hidden layers
            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:
                Whether to use batch normalization in the last layer

        """

        super().__init__()

        self.in_dim = in_dim
        self.hidden_dim = hidden_dim
        self.out_dim = out_dim

        self.fully_connected = nn.ModuleList()
        if layers <= 1:
            self.fully_connected.append(
                FCLayer(
                    in_dim,
                    out_dim,
                    activation=last_activation,
                    normalization=last_normalization,
                    dropout=dropout,
                )
            )
        else:
            self.fully_connected.append(
                FCLayer(
                    in_dim,
                    hidden_dim,
                    activation=activation,
                    normalization=normalization,
                    dropout=dropout,
                )
            )
            for _ in range(layers - 2):
                self.fully_connected.append(
                    FCLayer(
                        hidden_dim,
                        hidden_dim,
                        activation=activation,
                        normalization=normalization,
                        dropout=dropout,
                    )
                )
            self.fully_connected.append(
                FCLayer(
                    hidden_dim,
                    out_dim,
                    activation=last_activation,
                    normalization=last_normalization,
                    dropout=dropout,
                )
            )

    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

        """
        for fc in self.fully_connected:
            h = fc(h)
        return h

    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


# goli/nn/dgn_operations.py

from typing import Optional
import torch
from torch import Tensor


EPS = 1e-8


def get_grad_of_pos(
    source_pos: Tensor, dest_pos: Tensor, dir_idx: int, temperature: Optional[float] = None
) -> Tensor:
    r"""
    Get the vector field associated to the gradient of the positional
    encoding.

    $$F_k = \nabla pos_k$$

    or, if a temperature $T$ is provided

    $$F_k = softmax((\nabla pos_k)^+) - softmax((\nabla pos_k)^-)$$

    Where $F_k$ is the *k-th* directional field associated to the $k-th$ positional
    encoding.

    Parameters:

        source_pos: The positional encoding at the source node, used to compute the directional field
        dest_pos: The positional encoding at the destination node, used to compute the directional field
        h_in: The input features of the layer, before any operation.
        dir_idx: The index of the positional encoding ($k$ in the equation above)
        temperature: The temperature to use in the softmax of the directional field.
            If `None`, then the softmax is not applied on the field

    """

    grad = source_pos[:, :, dir_idx] - dest_pos[:, :, dir_idx]
    if temperature is not None:
        grad_pos, grad_neg = grad >= 0, grad <= 0
        grad_plus, grad_minus = grad, -grad

        # Compute softmax, considering one sign at a time.
        grad_plus[grad_neg] = -float("inf")
        grad_minus[grad_pos] = -float("inf")
        grad_plus = torch.nn.Softmax(1)(temperature * grad_plus)
        grad_minus = torch.nn.Softmax(1)(temperature * grad_minus)
        grad_plus[grad_neg] = 0
        grad_minus[grad_pos] = 0

        grad = grad_plus - grad_minus

    return grad


def aggregate_dir_smooth(
    h: Tensor,
    source_pos: Tensor,
    dest_pos: Tensor,
    h_in: Tensor,
    dir_idx: int,
    temperature: Optional[float] = None,
    **kwargs,
) -> Tensor:
    r"""
    The aggregation is the following:

    $$y^{(l)} = |\hat{F}_k| h^{(l)}$$

    - $\hat{F}^+_k$ is the normalized directional field *k-th* directional field $F_k$
    - $y^{(l)}$ is the returned aggregated result at the *l-th* layer.
    - $h^{(l)}$ is the node features at the *l-th* layer.

    Parameters:

        h: The features to aggregate $h^{(l)}$
        source_pos: The positional encoding at the source node, used to compute the directional field
        dest_pos: The positional encoding at the destination node, used to compute the directional field
        h_in: The input features of the layer, before any operation.
        dir_idx: The index of the positional encoding ($k$ in the equation above)
        temperature: The temperature to use in the softmax of the directional field.
            If `None`, then the softmax is not applied on the field

    Returns:
        h_mod: The aggregated features $y^{(l)}$

    """
    grad = get_grad_of_pos(source_pos=source_pos, dest_pos=dest_pos, dir_idx=dir_idx, temperature=temperature)
    h_mod = h * (grad.abs() / (torch.sum(grad.abs(), keepdim=True, dim=1) + EPS)).unsqueeze(-1)
    return torch.sum(h_mod, dim=1)


def aggregate_dir_dx_abs(
    h: Tensor,
    source_pos: Tensor,
    dest_pos: Tensor,
    h_in: Tensor,
    dir_idx: int,
    temperature: Optional[float] = None,
    **kwargs,
) -> Tensor:
    r"""
    The aggregation is the following:

    $$y^{(l)} = |B_{dx}^k h^{(l)}|$$

    $$B_{dx}^k = \hat{F}_k - diag \left(\sum_j{\hat{F}_{k_{(:, j)}}} \right)$$

    - $\hat{F}^+_k$ is the normalized positive component of the directional field *k-th* directional field $F_k$
    - $y^{(l)}$ is the returned aggregated result at the *l-th* layer.
    - $h^{(l)}$ is the node features at the *l-th* layer.

    Parameters:

        h: The features to aggregate $h^{(l)}$
        source_pos: The positional encoding at the source node, used to compute the directional field
        dest_pos: The positional encoding at the destination node, used to compute the directional field
        h_in: The input features of the layer, before any operation.
        dir_idx: The index of the positional encoding ($k$ in the equation above)
        temperature: The temperature to use in the softmax of the directional field.
            If `None`, then the softmax is not applied on the field

    Returns:
        h_mod: The aggregated features $y^{(l)}$

    """
    return torch.abs(aggregate_dir_dx_no_abs(h, source_pos, dest_pos, h_in, dir_idx, temperature, **kwargs))


def aggregate_dir_dx_no_abs(
    h: Tensor,
    source_pos: Tensor,
    dest_pos: Tensor,
    h_in: Tensor,
    dir_idx: int,
    temperature: Optional[float] = None,
    **kwargs,
) -> Tensor:
    r"""
    The aggregation is the following:

    $$y^{(l)} = B_{dx}^k h^{(l)}$$

    $$B_{dx}^k = \hat{F}_k - diag \left(\sum_j{\hat{F}_{k_{(:, j)}}} \right)$$

    - $\hat{F}^+_k$ is the normalized positive component of the directional field *k-th* directional field $F_k$
    - $y^{(l)}$ is the returned aggregated result at the *l-th* layer.
    - $h^{(l)}$ is the node features at the *l-th* layer.

    Parameters:

        h: The features to aggregate $h^{(l)}$
        source_pos: The positional encoding at the source node, used to compute the directional field
        dest_pos: The positional encoding at the destination node, used to compute the directional field
        h_in: The input features of the layer, before any operation.
        dir_idx: The index of the positional encoding ($k$ in the equation above)
        temperature: The temperature to use in the softmax of the directional field.
            If `None`, then the softmax is not applied on the field

    Returns:
        h_mod: The aggregated features $y^{(l)}$

    """
    grad = get_grad_of_pos(source_pos=source_pos, dest_pos=dest_pos, dir_idx=dir_idx, temperature=temperature)
    dir_weight = (grad / (torch.sum(grad.abs(), keepdim=True, dim=1) + EPS)).unsqueeze(-1)
    h_mod = h * dir_weight

    h_dx = torch.sum(h_mod, dim=1)
    h_self = -torch.sum(dir_weight, dim=1) * h_in

    # In case h_in has more parameters than h (for example when concatenating edges),
    # the derivative is only computed for the features contained in h_in.
    h_dx[..., : h_in.shape[-1]] = h_dx[..., : h_in.shape[-1]] + h_self
    return h_dx


def aggregate_dir_dx_abs_balanced(
    h: Tensor,
    source_pos: Tensor,
    dest_pos: Tensor,
    h_in: Tensor,
    dir_idx: int,
    temperature: Optional[float] = None,
    **kwargs,
) -> Tensor:
    r"""
    The aggregation is the same as `aggregate_dir_dx_no_abs`, but the positive and
    negative parts of the field are normalized separately.


    Parameters:

        h: The features to aggregate $h^{(l)}$
        source_pos: The positional encoding at the source node, used to compute the directional field
        dest_pos: The positional encoding at the destination node, used to compute the directional field
        h_in: The input features of the layer, before any operation.
        dir_idx: The index of the positional encoding ($k$ in the equation above)
        temperature: The temperature to use in the softmax of the directional field.
            If `None`, then the softmax is not applied on the field

    Returns:
        h_mod: The aggregated features $y^{(l)}$

    """
    grad = get_grad_of_pos(source_pos=source_pos, dest_pos=dest_pos, dir_idx=dir_idx, temperature=temperature)
    eig_front = torch.relu(grad) / (torch.sum(torch.relu(grad), keepdim=True, dim=1) + EPS)
    eig_back = torch.relu(-grad) / (torch.sum(torch.relu(-grad), keepdim=True, dim=1) + EPS)

    dir_weight = (eig_front.unsqueeze(-1) + eig_back.unsqueeze(-1)) / 2
    h_mod = h * dir_weight

    h_dx = torch.sum(h_mod, dim=1)
    h_self = -torch.sum(dir_weight, dim=1) * h_in

    # In case h_in has more parameters than h (for example when concatenating edges),
    # the derivative is only computed for the features contained in h_in.
    h_dx[..., : h_in.shape[-1]] = h_dx[..., : h_in.shape[-1]] + h_self
    return torch.abs(h_dx)


def aggregate_dir_forward(
    h: Tensor,
    source_pos: Tensor,
    dest_pos: Tensor,
    h_in: Tensor,
    dir_idx: int,
    temperature: Optional[float] = None,
    **kwargs,
) -> Tensor:
    r"""
    The aggregation is the following:

    $$y^{(l)} = \hat{F}^+_k h^{(l)}$$

    - $\hat{F}^+_k$ is the normalized positive component of the directional field *k-th* directional field $F_k$
    - $y^{(l)}$ is the returned aggregated result at the *l-th* layer.
    - $h^{(l)}$ is the node features at the *l-th* layer.

    Parameters:

        h: The features to aggregate $h^{(l)}$
        source_pos: The positional encoding at the source node, used to compute the directional field
        dest_pos: The positional encoding at the destination node, used to compute the directional field
        h_in: The input features of the layer, before any operation.
        dir_idx: The index of the positional encoding ($k$ in the equation above)
        temperature: The temperature to use in the softmax of the directional field.
            If `None`, then the softmax is not applied on the field

    Returns:
        h_mod: The aggregated features $y^{(l)}$
    """
    grad = get_grad_of_pos(source_pos=source_pos, dest_pos=dest_pos, dir_idx=dir_idx, temperature=temperature)
    eig_front = torch.relu(grad) / (torch.sum(torch.relu(grad), keepdim=True, dim=1) + EPS)
    h_mod = h * eig_front.unsqueeze(-1)
    return torch.sum(h_mod, dim=1)


def aggregate_dir_backward(
    h: Tensor,
    source_pos: Tensor,
    dest_pos: Tensor,
    h_in: Tensor,
    dir_idx: int,
    temperature: Optional[float] = None,
    **kwargs,
) -> Tensor:
    r"""
    The aggregation is the following:

    $$y^{(l)} = \hat{F}^-_k h^{(l)}$$

    - $\hat{F}^-_k$ is the normalized positive component of the directional field *k-th* directional field $F_k$
    - $y^{(l)}$ is the returned aggregated result at the *l-th* layer.
    - $h^{(l)}$ is the node features at the *l-th* layer.

    Parameters:

        h: The features to aggregate $h^{(l)}$
        source_pos: The positional encoding at the source node, used to compute the directional field
        dest_pos: The positional encoding at the destination node, used to compute the directional field
        h_in: The input features of the layer, before any operation.
        dir_idx: The index of the positional encoding ($k$ in the equation above)
        temperature: The temperature to use in the softmax of the directional field.
            If `None`, then the softmax is not applied on the field

    Returns:
        h_mod: The aggregated features $y^{(l)}$
    """
    return aggregate_dir_forward(h, -source_pos, -dest_pos, h_in, dir_idx, temperature, **kwargs)


DGN_AGGREGATORS = {
    "smooth": aggregate_dir_smooth,
    "dx_abs": aggregate_dir_dx_abs,
    "dx_no_abs": aggregate_dir_dx_no_abs,
    "dx_abs_balanced": aggregate_dir_dx_abs_balanced,
    "forward": aggregate_dir_forward,
    "backward": aggregate_dir_backward,
}


# goli/nn/pna_operations.py

import torch
import numpy as np
from functools import partial

EPS = 1e-5


def aggregate_mean(h, **kwargs):
    return torch.mean(h, dim=-2)


def aggregate_max(h, **kwargs):
    return torch.max(h, dim=-2)[0]


def aggregate_min(h, **kwargs):
    return torch.min(h, dim=-2)[0]


def aggregate_std(h, **kwargs):
    return torch.sqrt(aggregate_var(h) + EPS)


def aggregate_mean_laplacian(h, h_in, **kwargs):
    # In case h_in has more parameters than h (for example when concatenating edges),
    # the laplacian is only computed for the features contained in h_in.
    lap = -aggregate_mean(h, **kwargs)
    lap[..., : h_in.shape[-1]] = h_in + lap[..., : h_in.shape[-1]]
    return lap


def aggregate_var(h, **kwargs):
    h_mean_squares = torch.mean(h * h, dim=-2)
    h_mean = torch.mean(h, dim=-2)
    var = torch.relu(h_mean_squares - h_mean * h_mean)
    return var


def aggregate_moment(h, n=3, **kwargs):
    # for each node (E[(X-E[X])^n])^{1/n}
    # EPS is added to the absolute value of expectation before taking the nth root for stability
    h_mean = torch.mean(h, dim=-2, keepdim=True)
    h_n = torch.mean(torch.pow(h - h_mean, n), dim=-2)
    rooted_h_n = torch.sign(h_n) * torch.pow(torch.abs(h_n) + EPS, 1.0 / n)
    return rooted_h_n


def aggregate_sum(h, **kwargs):
    return torch.sum(h, dim=-2)


# each scaler is a function that takes as input X (B x N x Din), adj (B x N x N) and
# avg_d (dictionary containing averages over training set) and returns X_scaled (B x N x Din) as output


def scale_identity(h, D=None, avg_d=None):
    return h


def scale_amplification(h, D, avg_d):
    # log(D + 1) / d * h     where d is the average of the ``log(D + 1)`` in the training set
    return h * (np.log(D + 1) / avg_d["log"])


def scale_attenuation(h, D, avg_d):
    # (log(D + 1))^-1 / d * X     where d is the average of the ``log(D + 1))^-1`` in the training set
    return h * (avg_d["log"] / np.log(D + 1))


PNA_AGGREGATORS = {
    "mean": aggregate_mean,
    "sum": aggregate_sum,
    "max": aggregate_max,
    "min": aggregate_min,
    "std": aggregate_std,
    "var": aggregate_var,
    "lap": aggregate_mean_laplacian,
    "moment3": partial(aggregate_moment, n=3),
    "moment4": partial(aggregate_moment, n=4),
    "moment5": partial(aggregate_moment, n=5),
}


PNA_SCALERS = {
    "identity": scale_identity,
    "amplification": scale_amplification,
    "attenuation": scale_attenuation,
}


# goli/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 goli.nn.base_layers import FCLayer
from goli.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):

        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


# goli/utils/__init__.py

from . import fs
from . import tensor
from . import dict_tensor


# goli/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


# goli/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)()


# goli/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


# goli/utils/fs.py

from typing import Union
from typing import Optional

import os
import io
import appdirs
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(appdirs.user_cache_dir(appname="goli"))

    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 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()


# goli/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 goli.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


# goli/utils/spaces.py

from copy import deepcopy
import torch
import torch.optim.lr_scheduler as sc
import torchmetrics.functional as met

from goli.nn.base_layers import FCLayer
from goli.data.datamodule import DGLFromSmilesDataModule, DGLOGBDataModule

from goli.nn.dgl_layers import (
    GATLayer,
    GCNLayer,
    GINLayer,
    GatedGCNLayer,
    PNAConvolutionalLayer,
    PNAMessagePassingLayer,
    DGNConvolutionalLayer,
    DGNMessagePassingLayer,
)

from goli.nn.residual_connections import (
    ResidualConnectionConcat,
    ResidualConnectionDenseNet,
    ResidualConnectionNone,
    ResidualConnectionSimple,
    ResidualConnectionWeighted,
    ResidualConnectionRandom,
)


FC_LAYERS_DICT = {
    "fc": FCLayer,
}

DGL_LAYERS_DICT = {
    "gcn": GCNLayer,
    "gin": GINLayer,
    "gat": GATLayer,
    "gated-gcn": GatedGCNLayer,
    "pna-conv": PNAConvolutionalLayer,
    "pna-msgpass": PNAMessagePassingLayer,
    "dgn-conv": DGNConvolutionalLayer,
    "dgn-msgpass": DGNMessagePassingLayer,
}

LAYERS_DICT = deepcopy(DGL_LAYERS_DICT)
LAYERS_DICT.update(deepcopy(FC_LAYERS_DICT))


RESIDUALS_DICT = {
    "none": ResidualConnectionNone,
    "simple": ResidualConnectionSimple,
    "weighted": ResidualConnectionWeighted,
    "concat": ResidualConnectionConcat,
    "densenet": ResidualConnectionDenseNet,
    "random": ResidualConnectionRandom,
}

LOSS_DICT = {
    "mse": torch.nn.MSELoss(),
    "bce": torch.nn.BCELoss(),
    "l1": torch.nn.L1Loss(),
    "mae": torch.nn.L1Loss(),
}

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,
}

METRICS_CLASSIFICATION = {
    "accuracy": met.accuracy,
    "averageprecision": met.average_precision,
    "auroc": met.auroc,
    "confusionmatrix": met.confusion_matrix,
    "f1": met.f1_score,
    "fbeta": met.fbeta_score,
    "precisionrecallcurve": met.precision_recall_curve,
    "precision": met.precision,
    "recall": met.recall,
    "mcc": met.matthews_corrcoef,
}

METRICS_REGRESSION = {
    "mae": met.mean_absolute_error,
    "mape": met.mean_absolute_percentage_error,
    "mse": met.mean_squared_error,
    "msle": met.mean_squared_log_error,
    "pearsonr": met.pearson_corrcoef,
    "spearmanr": met.spearman_corrcoef,
    "r2": met.r2_score,
    "cosine": met.cosine_similarity,
}

METRICS_DICT = deepcopy(METRICS_CLASSIFICATION)
METRICS_DICT.update(METRICS_REGRESSION)


DATAMODULE_DICT = {
    "DGLFromSmilesDataModule": DGLFromSmilesDataModule,
    "DGLOGBDataModule": DGLOGBDataModule,
}


# goli/utils/tensor.py

import os
import torch
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
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]) -> np.ndarray:
    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 = [0] * (len(classes) + 1)
    found = False
    for i, v in enumerate(classes):
        if v == val:
            encoding[i] = 1
            found = True
            break
    if not found:
        encoding[-1] = 1
    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, **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, **kwargs)
    num = torch.sum(~torch.isnan(input), **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())
        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
    """

    return torch.sqrt(nan_var(input=input, unbiased=unbiased, **kwargs))


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 ModuleListConcat(torch.nn.ModuleList):
    def __init__(self, dim: int = -1):
        super().__init__()
        self.dim = dim

    def forward(self, *args, **kwargs) -> Tensor:
        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])


# goli/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)


# setup.py

from setuptools import setup
from setuptools import find_packages

# Sync the env.yml file here
install_requires = [
    "click",
    "tqdm",
    "loguru",
    "omegaconf",
    "appdirs",
    "fsspec>=2021.6",
    "s3fs>=2021.6",
    "gcsfs>=2021.6",
    "pandas",
    "numpy",
    "scipy",
    "scikit-learn",
    "torch>=1.7,<1.9",
    "torchvision",
    "tensorboard",
    "pytorch-lightning>=1.3",
    "torchmetrics>=0.2",
    "dgl>=0.5.2",
    "ogb",
    "datamol",
    "mordred",
    "umap-learn",
]

setup(
    name="goli",
    version="0.3.1",
    author="Valence Discovery",
    author_email="dominique@valencediscovery.com",
    url="https://github.com/valence-discovery/goli",
    description="A deep learning library focused on graph representation learning for real-world chemical tasks.",
    long_description=open("README.md", encoding="utf8").read(),
    long_description_content_type="text/markdown",
    project_urls={
        "Bug Tracker": "https://github.com/valence-discovery/goli/issues",
        "Documentation": "https://valence-discovery.github.io/goli",
        "Source Code": "https://github.com/valence-discovery/goli",
    },
    python_requires=">=3.7",
    # install_requires=install_requires,
    packages=find_packages(),
    include_package_data=True,
    entry_points={
        "console_scripts": ["goli=goli.cli:main_cli"],
    },
    classifiers=[
        "Development Status :: 5 - Production/Stable",
        "Intended Audience :: Developers",
        "Intended Audience :: Healthcare Industry",
        "Intended Audience :: Science/Research",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
        "Topic :: Scientific/Engineering :: Bio-Informatics",
        "Topic :: Scientific/Engineering :: Information Analysis",
        "Topic :: Scientific/Engineering :: Medical Science Apps.",
        "Natural Language :: English",
        "Operating System :: OS Independent",
        "Programming Language :: Python",
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.7",
        "Programming Language :: Python :: 3.8",
    ],
)


# goli/__init__.py

try:
    import dgl
except ImportError:
    raise ImportError(
        "DGL not installed. Please install it following the official documentation at https://github.com/dmlc/dgl/#installation."
    )


from ._version import __version__

from .config import load_config

from . import utils
from . import features
from . import data
from . import nn
from . import trainer


# goli/_version.py

__version__ = "0.3.1"


# goli/trainer/__init__.py

from . import predictor
from . import metrics

from .predictor import PredictorModule


# goli/trainer/metrics.py

from os import stat
from typing import Union, Callable, Optional, Dict, Any

import sys

import torch
import operator as op

from torchmetrics.utilities.distributed import reduce
import torchmetrics.functional.regression.mae

from goli.utils.tensor import nan_mean

# NOTE(hadim): the below is a fix to be able to import previously saved Goli 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/goli/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,
        target_to_int: bool = False,
    ):

        # Basic params
        self.threshold = threshold
        self.th_on_target = th_on_target
        self.th_on_preds = th_on_preds
        self.target_to_int = target_to_int
        self.operator, self.op_str = self._get_operator(operator)

    def compute(self, preds: torch.Tensor, target: torch.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)

        if self.target_to_int:
            target = target.to(int)

        return preds, target

    def __call__(self, preds: torch.Tensor, target: torch.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
        state["target_to_int"] = self.target_to_int

        # 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.target_to_int == obj.target_to_int,
            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,
        **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-flatten': The Tensor will be reduced to a vector without the NaN values.

                - 'ignore-mean-label': NaNs will be ignored when computing the loss. Note that each column
                  has a different number of NaNs, so the metric will be computed separately
                  on each column, and the metric result will be averaged over all columns.
                  *This option might slowdown the computation if there are too many labels*

            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 = target_nan_mask
        self.kwargs = kwargs

    @staticmethod
    def _get_metric(metric):
        from goli.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: torch.Tensor, target: torch.Tensor) -> torch.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)

        target_nans = torch.isnan(target)

        # Threshold the prediction
        if self.thresholder is not None:
            preds, target = self.thresholder(preds, target)

        # Manage the NaNs
        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-flatten":
            target = target[~target_nans]
            preds = preds[~target_nans]
        elif self.target_nan_mask == "ignore-mean-label":
            target_list = [target[..., ii][~target_nans[..., ii]] for ii in range(target.shape[-1])]
            preds_list = [preds[..., ii][~target_nans[..., ii]] for ii in range(preds.shape[-1])]
            target = target_list
            preds = preds_list
        else:
            raise ValueError(f"Invalid option `{self.target_nan_mask}`")

        if self.target_nan_mask == "ignore-mean-label":

            # Compute the metric for each column, and output nan if there's an error on a given column
            metric_val = []
            for ii in range(len(target)):
                try:
                    metric_val.append(self.metric(preds[ii], target[ii], **self.kwargs))
                except:
                    pass

            # Average the metric
            metric_val = nan_mean(torch.stack(metric_val))

        else:
            metric_val = self.metric(preds, target, **self.kwargs)
        return metric_val

    def __call__(self, preds: torch.Tensor, target: torch.Tensor) -> torch.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.kwargs == obj.kwargs,
        ]
        return 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["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)


# goli/trainer/predictor.py

from goli.trainer.metrics import MetricWrapper
from typing import Dict, List, Any, Union, Any, Callable, Tuple, Type, Optional
import os
import numpy as np
from copy import deepcopy
import yaml
import dgl
from loguru import logger
from inspect import signature

import torch
from torch import nn, Tensor

import pytorch_lightning as pl
from pytorch_lightning import _logger as log
from pytorch_lightning.utilities.exceptions import MisconfigurationException

from goli.config.config_convert import recursive_config_reformating
from goli.utils.tensor import nan_mean, nan_std, nan_median
from goli.utils.fs import mkdir
from goli.utils.spaces import LOSS_DICT, SCHEDULER_DICT

GOLI_PRETRAINED_MODELS = {
    "goli-zinc-micro-dummy-test": "gcs://goli-public/pretrained-models/goli-zinc-micro-dummy-test/model.ckpt"
}


class EpochSummary:
    r"""Container for collecting epoch-wise results"""

    def __init__(self, monitor="loss", mode: str = "min", metrics_on_progress_bar=[]):
        self.monitor = monitor
        self.mode = mode
        self.metrics_on_progress_bar = metrics_on_progress_bar
        self.summaries = {}
        self.best_summaries = {}

    class Results:
        def __init__(
            self,
            targets: Tensor,
            predictions: Tensor,
            loss: float,
            metrics: dict,
            monitored_metric: str,
            n_epochs: int,
        ):
            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

    def set_results(self, name, targets, predictions, loss, metrics, n_epochs) -> float:
        metrics[f"loss/{name}"] = loss
        self.summaries[name] = EpochSummary.Results(
            targets=targets,
            predictions=predictions,
            loss=loss,
            metrics=metrics,
            monitored_metric=f"{self.monitor}/{name}",
            n_epochs=n_epochs,
        )
        if self.is_best_epoch(name, loss, metrics):
            self.best_summaries[name] = self.summaries[name]

    def is_best_epoch(self, name, loss, metrics):
        if not (name in self.best_summaries.keys()):
            return True

        metrics[f"loss/{name}"] = loss
        monitor_name = f"{self.monitor}/{name}"
        if not monitor_name in self.best_summaries.keys():
            return True

        if self.mode == "max":
            return metrics[monitor_name] > self.best_summaries[name].monitored
        elif self.mode == "min":
            return metrics[monitor_name] < self.best_summaries[name].monitored
        else:
            ValueError(f"Mode must be 'min' or 'max', provided `{self.mode}`")

    def get_results(self, name):
        return self.summaries[name]

    def get_best_results(self, name):
        return self.best_summaries[name]

    def get_results_on_progress_bar(self, name):
        results = self.summaries[name]
        results_prog = {
            f"{kk}/{name}": results.metrics[f"{kk}/{name}"] for kk in self.metrics_on_progress_bar
        }
        return results_prog

    def get_dict_summary(self):
        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


class PredictorModule(pl.LightningModule):
    def __init__(
        self,
        model_class: Type[nn.Module],
        model_kwargs: Dict[str, Any],
        loss_fun: Union[str, Callable],
        random_seed: int = 42,
        optim_kwargs: Optional[Dict[str, Any]] = None,
        lr_reduce_on_plateau_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[int, float, str]] = None,
        metrics: Dict[str, Callable] = None,
        metrics_on_progress_bar: List[str] = [],
        metrics_on_training_set: Optional[List[str]] = None,
        flag_kwargs: Dict[str, Any] = None,
    ):
        r"""
        A class that allows to use regression or classification models easily
        with Pytorch-Lightning.

        Parameters:
            model_class:
                pytorch module used to create a model

            model_kwargs:
                Key-word arguments used to initialize the model from `model_class`.

            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()`.

            random_seed:
                The random seed used by Pytorch to initialize random tensors.

            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.`)

            lr_reduce_on_plateau_kwargs:
                DEPRECIATED. USE `torch_scheduler_kwargs` instead.
                Dictionnary for the reduction of learning rate when reaching plateau, with possible keys below.

                - factor `float`: Factor by which to reduce the learning rate (Default=`0.5`)
                - patience `int`: Number of epochs without improvement to wait before reducing
                  the learning rate (Default=`10`)
                - mode `str`: One of min, max. In min mode, lr will be reduced when the quantity
                  monitored has stopped decreasing; in max mode it will be reduced when the quantity
                  monitored has stopped increasing. (Default=`"min"`).
                - min_lr `float`: A scalar or a list of scalars. A lower bound on the learning rate
                  of all param groups or each group respectively (Default=`1e-4`)

            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`)

            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-flatten': The Tensor will be reduced to a vector without the NaN values.

                - 'ignore-mean-label': NaNs will be ignored when computing the loss. Note that each column
                  has a different number of NaNs, so the metric will be computed separately
                  on each column, and the metric result will be averaged over all columns.
                  *This option might slowdown the computation if there are too many labels*

            metrics:
                A dictionnary of metrics to compute on the prediction, other than the loss function.
                These metrics will be logged into TensorBoard.

            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.

            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
        """

        self.save_hyperparameters()

        torch.random.manual_seed(random_seed)
        np.random.seed(random_seed)

        super().__init__()
        self.model = model_class(**model_kwargs)

        # Basic attributes
        self.loss_fun = self.parse_loss_fun(loss_fun)
        self.random_seed = random_seed
        self.target_nan_mask = target_nan_mask
        self.metrics = metrics if metrics is not None else {}
        self.metrics_on_progress_bar = metrics_on_progress_bar
        self.metrics_on_training_set = (
            list(self.metrics.keys()) if metrics_on_training_set is None else 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_kwargs = {"alpha": 0.01, "n_steps": 0}
        if flag_kwargs is not None:
            self.flag_kwargs.update(flag_kwargs)
        assert isinstance(self.flag_kwargs["n_steps"], int) and (self.flag_kwargs["n_steps"] >= 0)
        assert self.flag_kwargs["alpha"] > 0

        # Set the parameters and default value for the optimizer, and check values
        self.optim_kwargs = {"lr": 1e-3, "weight_decay": 0.0}
        if optim_kwargs is not None:
            self.optim_kwargs.update(optim_kwargs)
        assert self.optim_kwargs["lr"] > 0
        assert self.optim_kwargs["weight_decay"] >= 0

        # Set the lightning scheduler
        self.scheduler_kwargs = {
            "interval": "epoch",
            "monitor": "loss/val",
            "mode": "min",
            "frequency": 1,
            "strict": True,
        }
        if scheduler_kwargs is not None:
            self.scheduler_kwargs.update(scheduler_kwargs)

        # Set the depreciated arguments, if provided
        if lr_reduce_on_plateau_kwargs is not None:
            logger.warning(
                "`lr_reduce_on_plateau_kwargs` is depreciated, use `torch_scheduler_kwargs` instead."
            )
            if torch_scheduler_kwargs is not None:
                raise ValueError(
                    "Ambiguous. Both `lr_reduce_on_plateau_kwargs` and `torch_scheduler_kwargs` are provided"
                )
            torch_scheduler_kwargs = {
                "module_type": "ReduceLROnPlateau",
                **lr_reduce_on_plateau_kwargs,
            }

        # Set the pytorch scheduler arguments
        if torch_scheduler_kwargs is None:
            self.torch_scheduler_kwargs = {}
        else:
            self.torch_scheduler_kwargs = torch_scheduler_kwargs
        self.torch_scheduler_kwargs.setdefault("module_type", "ReduceLROnPlateau")

        # Initialize the epoch summary
        monitor = self.scheduler_kwargs["monitor"].split("/")[0]
        mode = self.scheduler_kwargs["mode"]
        self.epoch_summary = EpochSummary(
            monitor, mode=mode, metrics_on_progress_bar=self.metrics_on_progress_bar
        )

        # This helps avoid a bug when saving hparams to yaml with different
        # dict or str formats
        self._set_hparams(recursive_config_reformating(self.hparams))

    @staticmethod
    def parse_loss_fun(loss_fun: Union[str, Callable]) -> Callable:
        r"""
        Parse the loss function from a string

        Parameters:
            loss_fun:
                A callable corresponding to the loss function or a string
                specifying the loss function from `LOSS_DICT`. Accepted strings are:
                "mse", "bce", "l1", "mae", "cosine".

        Returns:
            Callable:
                Function or callable to compute the loss, takes `preds` and `targets` as inputs.
        """

        if isinstance(loss_fun, str):
            loss_fun = LOSS_DICT[loss_fun]
        elif not callable(loss_fun):
            raise ValueError(f"`loss_fun` must be `str` or `callable`. Provided: {type(loss_fun)}")

        return loss_fun

    def forward(self, inputs: Dict) -> Dict[str, Union[torch.Tensor, Any]]:
        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"])
        out = self.model.forward(feats)

        # Convert the output of the model to a dictionary
        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, dgl.DGLHeteroGraph):
            for key, val in feats.ndata.items():
                if isinstance(val, torch.Tensor):
                    feats.ndata[key] = val.to(dtype=self.dtype)
            for key, val in feats.edata.items():
                if isinstance(val, torch.Tensor):
                    feats.edata[key] = val.to(dtype=self.dtype)

        return feats

    def configure_optimizers(self):
        # Configure the parameters for the schedulers
        sc_kwargs = deepcopy(self.torch_scheduler_kwargs)
        scheduler_class = SCHEDULER_DICT[sc_kwargs.pop("module_type")]
        sig = signature(scheduler_class.__init__)
        key_args = [p.name for p in sig.parameters.values()]
        if "monitor" in key_args:
            sc_kwargs.setdefault("monitor", self.scheduler_kwargs["monitor"])
        if "mode" in key_args:
            sc_kwargs.setdefault("mode", self.scheduler_kwargs["mode"])

        # Define the optimizer and schedulers
        optimiser = torch.optim.Adam(self.parameters(), **self.optim_kwargs)
        torch_scheduler = scheduler_class(optimizer=optimiser, **sc_kwargs)
        scheduler = {
            "scheduler": torch_scheduler,
            **self.scheduler_kwargs,
        }
        return [optimiser], [scheduler]

    @staticmethod
    def compute_loss(
        preds: Tensor,
        targets: Tensor,
        weights: Optional[Tensor],
        loss_fun: Callable,
        target_nan_mask: Union[Type, str] = "ignore",
    ) -> 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-flatten': The Tensor will be reduced to a vector without the NaN values.

                - 'ignore-mean-label': NaNs will be ignored when computing the loss. Note that each column
                  has a different number of NaNs, so the metric will be computed separately
                  on each column, and the metric result will be averaged over all columns.
                  *This option might slowdown the computation if there are too many labels*

            loss_fun:
                Loss function to use

        Returns:
            Tensor:
                Resulting loss
        """

        wrapped_loss_fun = MetricWrapper(
            metric=loss_fun, threshold_kwargs=None, target_nan_mask=target_nan_mask
        )
        if weights is not None:
            raise NotImplementedError("Weights are no longer supported in the loss")
        loss = wrapped_loss_fun(preds=preds, target=targets)

        return loss

    def get_metrics_logs(
        self, preds: Tensor, targets: Tensor, weights: Optional[Tensor], step_name: str, loss: Tensor
    ) -> Dict[str, Any]:
        r"""
        Get the logs for the loss and the different metrics, in a format compatible with
        Pytorch-Lightning.

        Parameters:
            preds:
                Predicted values

            targets:
                Target values

            step_name:
                A string to mention whether the metric is computed on the training,
                validation or test set.

                - "train": On the training set
                - "val": On the validation set
                - "test": On the test set

        """

        targets = targets.to(dtype=preds.dtype, device=preds.device)

        # Compute the metrics always used in regression tasks
        metric_logs = {}
        metric_logs[f"mean_pred/{step_name}"] = nan_mean(preds)
        metric_logs[f"std_pred/{step_name}"] = nan_std(preds)
        metric_logs[f"median_pred/{step_name}"] = nan_median(preds)
        metric_logs[f"mean_target/{step_name}"] = nan_mean(targets)
        metric_logs[f"std_target/{step_name}"] = nan_std(targets)
        metric_logs[f"median_target/{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 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 = f"{key}/{step_name}"
            try:
                metric_logs[metric_name] = metric(preds, targets)
            except Exception as e:
                metric_logs[metric_name] = torch.as_tensor(float("nan"))

        # Convert all metrics to CPU, except for the loss
        metric_logs[f"{self.loss_fun._get_name()}/{step_name}"] = loss.detach().cpu()
        metric_logs = {key: metric.detach().cpu() for key, metric in metric_logs.items()}

        return metric_logs

    def _general_step(
        self, batch: Dict[str, Tensor], batch_idx: int, step_name: str, to_cpu: bool
    ) -> Dict[str, Any]:
        r"""Common code for training_step, validation_step and testing_step"""
        preds = self.forward(batch)["preds"]
        targets = batch.pop("labels").to(dtype=preds.dtype)
        weights = batch.pop("weights", None)

        loss = self.compute_loss(
            preds=preds,
            targets=targets,
            weights=weights,
            target_nan_mask=self.target_nan_mask,
            loss_fun=self.loss_fun,
        )

        device = "cpu" if to_cpu else None
        preds = preds.detach().to(device=device)
        targets = targets.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"{self.loss_fun._get_name()}/{step_name}"] = loss.detach().cpu()
        return loss, step_dict

    def flag_step(
        self, batch: Dict[str, Tensor], batch_idx: int, 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.ndata["feat"].shape

        pert = torch.FloatTensor(X_shape).uniform_(-alpha, alpha).to(device=X.device)
        pert.requires_grad = True

        # Perturb the features
        pert_batch = deepcopy(batch)
        pert_batch["features"].ndata["feat"] = batch["features"].ndata["feat"] + pert

        preds = self.forward(pert_batch)["preds"]
        targets = batch.pop("labels").to(dtype=preds.dtype)
        weights = batch.pop("weights", None)
        loss = (
            self.compute_loss(
                preds=preds,
                targets=targets,
                weights=weights,
                target_nan_mask=self.target_nan_mask,
                loss_fun=self.loss_fun,
            )
            / 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
            pert_batch["features"].ndata["feat"] = batch["features"].ndata["feat"] + pert
            preds = self.forward(pert_batch)["preds"]
            loss = (
                self.compute_loss(
                    preds=preds,
                    targets=targets,
                    weights=weights,
                    target_nan_mask=self.target_nan_mask,
                    loss_fun=self.loss_fun,
                )
                / n_steps
            )

        device = "cpu" if to_cpu else None
        preds = preds.detach().to(device=device)
        targets = targets.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"{self.loss_fun._get_name()}/{step_name}"] = loss.detach().cpu()
        return loss, step_dict

    def training_step(self, batch: Dict[str, Tensor], batch_idx: int) -> Dict[str, Any]:
        loss = None
        step_dict = None

        # Train using FLAG
        if self.flag_kwargs["n_steps"] > 0:
            loss, step_dict = self.flag_step(batch=batch, batch_idx=batch_idx, step_name="train", to_cpu=True)
        # Train normally, without using FLAG
        elif self.flag_kwargs["n_steps"] == 0:
            loss, step_dict = self._general_step(
                batch=batch, batch_idx=batch_idx, step_name="train", to_cpu=True
            )

        metrics_logs = self.get_metrics_logs(
            preds=step_dict["preds"],
            targets=step_dict["targets"],
            weights=step_dict["weights"],
            step_name="train",
            loss=loss,
        )

        step_dict.update(metrics_logs)
        step_dict["loss"] = loss

        self.logger.log_metrics(metrics_logs, step=self.global_step)

        # Predictions and targets are no longer needed after the step.
        # Keeping them will increase memory usage significantly for large datasets.
        step_dict.pop("preds")
        step_dict.pop("targets")
        step_dict.pop("weights")

        return step_dict

    def validation_step(self, batch: Dict[str, Tensor], batch_idx: int) -> Dict[str, Any]:
        return self._general_step(batch=batch, batch_idx=batch_idx, step_name="val", to_cpu=True)[1]

    def test_step(self, batch: Dict[str, Tensor], batch_idx: int) -> Dict[str, Any]:
        return self._general_step(batch=batch, batch_idx=batch_idx, step_name="test", to_cpu=True)[1]

    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 = torch.cat([out["preds"] for out in outputs], dim=0)
        targets = torch.cat([out["targets"] for out in outputs], dim=0)
        if outputs[0]["weights"] is not None:
            weights = torch.cat([out["weights"] for out in outputs], dim=0)
        else:
            weights = None
        loss = self.compute_loss(
            preds=preds,
            targets=targets,
            weights=weights,
            target_nan_mask=self.target_nan_mask,
            loss_fun=self.loss_fun,
        )
        metrics_logs = self.get_metrics_logs(
            preds=preds, targets=targets, weights=weights, step_name=step_name, loss=loss
        )

        self.epoch_summary.set_results(
            name=step_name,
            predictions=preds,
            targets=targets,
            loss=loss,
            metrics=metrics_logs,
            n_epochs=self.current_epoch,
        )

        return metrics_logs

    def training_epoch_end(self, outputs: Dict):
        """
        Nothing happens at the end of the training epoch.
        It serves no purpose to do a general step for the training,
        but it can explode the RAM when using a large dataset.
        """
        pass

    def validation_epoch_end(self, outputs: Dict[str, Any]):

        metrics_logs = self._general_epoch_end(outputs=outputs, step_name="val")

        lr = self.optimizers().param_groups[0]["lr"]
        metrics_logs["lr"] = lr
        metrics_logs["n_epochs"] = self.current_epoch
        self.log_dict(metrics_logs)

        # Save yaml file with the metrics summaries
        full_dict = {}
        full_dict.update(self.epoch_summary.get_dict_summary())
        tb_path = self.logger.log_dir

        # Write the YAML file with the metrics
        if self.current_epoch >= 0:
            mkdir(tb_path)
            with open(os.path.join(tb_path, "metrics.yaml"), "w") as file:
                yaml.dump(full_dict, file)

    def test_epoch_end(self, outputs: Dict[str, Any]):

        metrics_logs = self._general_epoch_end(outputs=outputs, step_name="test")
        self.log_dict(metrics_logs)

        # Save yaml file with the metrics summaries
        full_dict = {}
        full_dict.update(self.epoch_summary.get_dict_summary())
        tb_path = self.logger.log_dir
        os.makedirs(tb_path, exist_ok=True)
        with open(f"{tb_path}/metrics.yaml", "w") as file:
            yaml.dump(full_dict, file)

    def on_train_start(self):
        hparams_log = deepcopy(self.hparams)
        hparams_log["n_params"] = self.n_params
        self.logger.log_hyperparams(hparams_log, self.epoch_summary.get_results("val").metrics)

    def get_progress_bar_dict(self) -> Dict[str, float]:
        prog_dict = super().get_progress_bar_dict()
        results_on_progress_bar = self.epoch_summary.get_results_on_progress_bar("val")
        prog_dict["loss/val"] = self.epoch_summary.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 GOLI_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 `goli.trainer.PredictorModule.list_pretrained_models()`.
        """

        if name not in GOLI_PRETRAINED_MODELS:
            raise ValueError(
                f"The model '{name}' is not available. Choose from {set(GOLI_PRETRAINED_MODELS.keys())}."
            )

        return PredictorModule.load_from_checkpoint(GOLI_PRETRAINED_MODELS[name])


# goli/trainer/refactor_predictor_mtl.py
from goli.trainer.metrics import MetricWrapper
from typing import Dict, List, Any, Union, Any, Callable, Tuple, Type, Optional
import os
import numpy as np
from copy import deepcopy
import yaml
import dgl
from loguru import logger
from inspect import signature
from pprint import pprint

import torch
from torch import nn, Tensor

import pytorch_lightning as pl
from pytorch_lightning import _logger as log
from pytorch_lightning.utilities.exceptions import MisconfigurationException

from goli.config.config_convert import recursive_config_reformating
from goli.trainer.predictor_options import EvalOptions, FlagOptions, ModelOptions, OptimOptions
from goli.trainer.predictor_summaries import Summary, TaskSummaries
from goli.utils.tensor import nan_mean, nan_std, nan_median
from goli.utils.fs import mkdir
from goli.utils.spaces import SCHEDULER_DICT

GOLI_PRETRAINED_MODELS = {
    "goli-zinc-micro-dummy-test": "gcs://goli-public/pretrained-models/goli-zinc-micro-dummy-test/model.ckpt"
}


class PredictorModule(pl.LightningModule):
    def __init__(
        self,
        model_class: Type[nn.Module],                                   # Leave
        model_kwargs: Dict[str, Any],                                   # Leave
        loss_fun: Dict[str, Union[str, Callable]],                      # Task-specific
        random_seed: int = 42,                                          # Leave        
        optim_kwargs: Optional[Dict[str, Any]] = None,                  # Leave for now
        lr_reduce_on_plateau_kwargs: Optional[Dict[str, Any]] = None,   # Leave for now
        torch_scheduler_kwargs: Optional[Dict[str, Any]] = None,        # Leave for now
        scheduler_kwargs: Optional[Dict[str, Any]] = None,              # Leave for now
        target_nan_mask: Optional[Union[int, float, str]] = None,       # Leave
        metrics: Dict[str, Callable] = None,                            # Task-specific
        metrics_on_progress_bar: List[str] = [],                        # Task-specific
        metrics_on_training_set: Optional[List[str]] = None,            # Task-specific
        flag_kwargs: Dict[str, Any] = None,
    ):
        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

        super().__init__()

        model_options = ModelOptions(
            model_class=model_class, 
            model_kwargs=model_kwargs
        )
        optim_options = OptimOptions(
            optim_kwargs=optim_kwargs,
            lr_reduce_on_plateau_kwargs=lr_reduce_on_plateau_kwargs,
            torch_scheduler_kwargs=torch_scheduler_kwargs,
            scheduler_kwargs=scheduler_kwargs,
        )
        eval_options = EvalOptions(
            loss_fun=loss_fun,
            metrics=metrics,
            metrics_on_progress_bar=metrics_on_progress_bar,
            metrics_on_training_set=metrics_on_training_set
        )
        flag_options = FlagOptions(
            flag_kwargs=flag_kwargs
        )

        self._model_options = model_options
        self._optim_options = optim_options
        self._eval_options = eval_options
        self._flag_options = flag_options

        self.model = self._model_options.model_class(**self._model_options.model_kwargs)
        self.tasks = list(loss_fun.keys())

        # Basic attributes
        self.loss_fun = {"micro_zinc": EvalOptions.parse_loss_fun(loss_fun["micro_zinc"])}
        #self.loss_fun = self._eval_options.parse_loss_fun(self._eval_options.loss_fun)              # Look into this
        #self.metrics = self._eval_options.metrics if self._eval_options.metrics is not None else {}
        #self.metrics_on_progress_bar = self._eval_options.metrics_on_progress_bar
        #self.metrics_on_training_set = (
        #    list(self.metrics.keys()) if self._eval_options.metrics_on_training_set is None else self._eval_options.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()
        # Set the parameters and default value for the optimizer, and check values
        self.optim_kwargs = self._optim_options.optim_kwargs
        # Set the lightning scheduler
        self.scheduler_kwargs = self._optim_options.scheduler_kwargs
        # Set the pytorch scheduler arguments
        self.torch_scheduler_kwargs = self._optim_options.torch_scheduler_kwargs

        # Initialize the epoch summary
        monitor = self.scheduler_kwargs["monitor"].split("/")[0]
        mode = self.scheduler_kwargs["mode"]
        #self.epoch_summary = Summary(
        #    loss_fun=self.loss_fun,
        #    metrics=self.metrics,
        #    metrics_on_training_set=self.metrics_on_training_set,
        #    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))

    def forward(self, inputs: Dict) -> Dict[str, Union[torch.Tensor, Any]]:
        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.

        """
        #print("INSIDE FORWARD!! \n")
        #pprint(inputs.keys())
        #pprint(inputs['labels'])
        #print("\n\n")

        # Convert to the right dtype and run the model
        feats = self._convert_features_dtype(inputs["features"])
        out = self.model.forward(feats)
        #pprint(out)
        #print("\n\n")

        # Convert the output of the model to a dictionary
        if isinstance(out, dict) and ("preds" in out.keys()):
            out_dict = out
        else:
            out_dict = {"preds": out}

        #pprint(out_dict)
        #print("DONE FORWARD!! \n")
        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, dgl.DGLHeteroGraph):
            for key, val in feats.ndata.items():
                if isinstance(val, torch.Tensor):
                    feats.ndata[key] = val.to(dtype=self.dtype)
            for key, val in feats.edata.items():
                if isinstance(val, torch.Tensor):
                    feats.edata[key] = val.to(dtype=self.dtype)

        return feats

    def configure_optimizers(self):
        # Configure the parameters for the schedulers
        sc_kwargs = deepcopy(self.torch_scheduler_kwargs)
        scheduler_class = SCHEDULER_DICT[sc_kwargs.pop("module_type")]
        sig = signature(scheduler_class.__init__)
        key_args = [p.name for p in sig.parameters.values()]
        if "monitor" in key_args:
            sc_kwargs.setdefault("monitor", self.scheduler_kwargs["monitor"])
        if "mode" in key_args:
            sc_kwargs.setdefault("mode", self.scheduler_kwargs["mode"])

        # Define the optimizer and schedulers
        optimiser = torch.optim.Adam(self.parameters(), **self.optim_kwargs)
        torch_scheduler = scheduler_class(optimizer=optimiser, **sc_kwargs)
        scheduler = {
            "scheduler": torch_scheduler,
            **self.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],                  # Was Callable
        target_nan_mask: Union[Type, str] = "ignore",
    ) -> 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-flatten': The Tensor will be reduced to a vector without the NaN values.

                - 'ignore-mean-label': NaNs will be ignored when computing the loss. Note that each column
                  has a different number of NaNs, so the metric will be computed separately
                  on each column, and the metric result will be averaged over all columns.
                  *This option might slowdown the computation if there are too many labels*

            loss_fun:
                Loss function to use

        Returns:
            Tensor:
                Resulting loss
        """            
        #wrapped_loss_fun = MetricWrapper(metric=loss_fun, threshold_kwargs=None, target_nan_mask=target_nan_mask)
        wrapped_loss_fun_dict = {task: MetricWrapper(metric=loss, threshold_kwargs=None, target_nan_mask=target_nan_mask) for task, loss in loss_fun.items()}
        # They all have the same loss function. Get the first wrapped_loss_fun
        #wrapped_loss_fun = next(iter(wrapped_loss_fun_dict.values()))

        if weights is not None:
            raise NotImplementedError("Weights are no longer supported in the loss")
        #loss = wrapped_loss_fun(preds=preds, target=targets)
        all_task_losses = {task: wrapped(preds=preds[task], target=targets[task]) for task, wrapped in wrapped_loss_fun_dict.items()}
        #total_loss = torch.sum(all_task_losses.values(), dim=0)
        total_loss = torch.zeros_like(next(iter(all_task_losses.values())), requires_grad=True)
        for task, loss in all_task_losses.items():
            torch.add(total_loss, loss)
        num_tasks = len(all_task_losses.keys())
        weighted_loss = torch.div(total_loss, num_tasks)
        #weighted_loss = Variable(weighted_loss, requires_grad = True)
        return weighted_loss

    def _general_step(
        self, batch: Dict[str, Tensor], batch_idx: int, step_name: str, to_cpu: bool
    ) -> Dict[str, Any]:
        r"""Common code for training_step, validation_step and testing_step"""
        preds = self.forward(batch)["preds"]                    # The dictionary of predictions
        #targets = batch.pop("labels").to(dtype=preds.dtype)
        targets_dict = batch.pop("labels")
        for task, pred in preds.items():
            targets_dict[task] = targets_dict[task].to(dtype=pred.dtype)
        weights = batch.pop("weights", None)

        loss = self.compute_loss(
            preds=preds,
            targets=targets_dict,
            weights=weights,
            target_nan_mask=self.target_nan_mask,
            loss_fun=self.loss_fun,                         # This is a dictionary
        )

        device = "cpu" if to_cpu else None
        #preds = preds.detach().to(device=device)
        #targets = targets.detach().to(device=device)
        for task in preds:
            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()

        #step_dict[f"weighted_loss/{step_name}"] = loss.detach().cpu()
        #step_dict[f"{self.loss_fun[self.tasks[0]]._get_name()}/{step_name}"] = loss.detach().cpu()
        step_dict[f"loss/{step_name}"] = loss.detach().cpu()

        return loss, step_dict

    def get_metrics_logs(
        self, preds: Dict[str, Tensor], targets: Dict[str, Tensor], weights: Optional[Tensor], step_name: str, loss: Tensor
    ) -> Dict[str, Any]:
        r"""
        Get the logs for the loss and the different metrics, in a format compatible with
        Pytorch-Lightning.

        Parameters:
            preds:
                Predicted values

            targets:
                Target values

            step_name:
                A string to mention whether the metric is computed on the training,
                validation or test set.

                - "train": On the training set
                - "val": On the validation set
                - "test": On the test set
        """
        targets_dict = {}
        for task in self.tasks:
            targets_dict[task] = targets[task].to(dtype=preds[task].dtype, device=preds[task].device)
        #targets = targets.to(dtype=preds.dtype, device=preds.device)

        # Compute the metrics always used in regression tasks
        task = self.tasks[0]
        metric_logs = {}
        metric_logs[f"loss/{step_name}"] = nan_mean(loss)           # To get it running
        metric_logs[f"mean_pred/{step_name}"] = nan_mean(preds[task])
        metric_logs[f"std_pred/{step_name}"] = nan_std(preds[task])
        metric_logs[f"median_pred/{step_name}"] = nan_median(preds[task])
        metric_logs[f"mean_target/{step_name}"] = nan_mean(targets_dict[task])
        metric_logs[f"std_target/{step_name}"] = nan_std(targets_dict[task])
        metric_logs[f"median_target/{step_name}"] = nan_median(targets_dict[task])
#        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 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 = f"{key}/{step_name}"
#            try:
#                metric_logs[metric_name] = metric(preds, targets)
#            except Exception as e:
#                metric_logs[metric_name] = torch.as_tensor(float("nan"))

        # Convert all metrics to CPU, except for the loss
#        metric_logs[f"{self.loss_fun._get_name()}/{step_name}"] = loss.detach().cpu()
#        metric_logs = {key: metric.detach().cpu() for key, metric in metric_logs.items()}

        return metric_logs

    def flag_step(
        self, batch: Dict[str, Tensor], batch_idx: int, 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.ndata["feat"].shape

        pert = torch.FloatTensor(X_shape).uniform_(-alpha, alpha).to(device=X.device)
        pert.requires_grad = True

        # Perturb the features
        pert_batch = deepcopy(batch)
        pert_batch["features"].ndata["feat"] = batch["features"].ndata["feat"] + pert

        preds = self.forward(pert_batch)["preds"]
        targets = batch.pop("labels").to(dtype=preds.dtype)
        weights = batch.pop("weights", None)
        loss = (
            self.compute_loss(
                preds=preds,
                targets=targets,
                weights=weights,
                target_nan_mask=self.target_nan_mask,
                loss_fun=self.loss_fun,
            )
            / 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
            pert_batch["features"].ndata["feat"] = batch["features"].ndata["feat"] + pert
            preds = self.forward(pert_batch)["preds"]
            loss = (
                self.compute_loss(
                    preds=preds,
                    targets=targets,
                    weights=weights,
                    target_nan_mask=self.target_nan_mask,
                    loss_fun=self.loss_fun,
                )
                / n_steps
            )

        device = "cpu" if to_cpu else None
        preds = preds.detach().to(device=device)
        targets = targets.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"{self.loss_fun._get_name()}/{step_name}"] = loss.detach().cpu()
        return loss, step_dict

    def training_step(self, batch: Dict[str, Tensor], batch_idx: int) -> Dict[str, Any]:
        loss = None
        step_dict = None

        # Train using FLAG
        if self.flag_kwargs["n_steps"] > 0:
            loss, step_dict = self.flag_step(batch=batch, batch_idx=batch_idx, step_name="train", to_cpu=True)
        # Train normally, without using FLAG
        elif self.flag_kwargs["n_steps"] == 0:
            loss, step_dict = self._general_step(
                batch=batch, batch_idx=batch_idx, step_name="train", to_cpu=True
            )

        #self.epoch_summary.update_predictor_state(
        #    step_name="train",
        #    targets=step_dict["targets"],
        #    predictions=step_dict["preds"],
        #    loss=loss,
        #    n_epochs=self.current_epoch,
        #)
        #metrics_logs = self.epoch_summary.get_metrics_logs()
        metrics_logs = self.get_metrics_logs(preds=step_dict["preds"], targets=step_dict["targets"], weights=None, step_name="train", loss=loss)

        step_dict.update(metrics_logs)
        step_dict["loss"] = loss

        self.logger.log_metrics(metrics_logs, step=self.global_step)            # This is a pytorch lightning function call

        # Predictions and targets are no longer needed after the step.
        # Keeping them will increase memory usage significantly for large datasets.
        step_dict.pop("preds")
        step_dict.pop("targets")
        step_dict.pop("weights")

        return step_dict

    def validation_step(self, batch: Dict[str, Tensor], batch_idx: int) -> Dict[str, Any]:
        return self._general_step(batch=batch, batch_idx=batch_idx, step_name="val", to_cpu=True)[1]

    def test_step(self, batch: Dict[str, Tensor], batch_idx: int) -> Dict[str, Any]:
        return self._general_step(batch=batch, batch_idx=batch_idx, step_name="test", to_cpu=True)[1]

    # So it's over the entire dataset
    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"""
        # epoch_end returns a list of all the output from the _step
        # Transform the list of dict of dict, into a dict of list of dict
        #preds = torch.cat([out["preds"] for out in outputs], dim=0)         # outputs is a list of the output
        #targets = torch.cat([out["targets"] for out in outputs], dim=0)
        preds = {}
        targets = {}
        for task in self.tasks:
            preds[task] = torch.cat([out["preds"][task] for out in outputs], dim=0)
            targets[task] = torch.cat([out["targets"][task] for out in outputs], dim=0)
        if outputs[0]["weights"] is not None:
            weights = torch.cat([out["weights"] for out in outputs], dim=0)
        else:
            weights = None
        loss = self.compute_loss(
            preds=preds,
            targets=targets,
            weights=weights,
            target_nan_mask=self.target_nan_mask,
            loss_fun=self.loss_fun,
        )
        """Problem: we need to update the predictor results to make the metrics_logs, but we need the metrics_logs to update the predictor_result"""
        #self.epoch_summary.update_predictor_state(            # After adding predictor results
        #    step_name=step_name,
        #    predictions=preds,
        #    targets=targets,
        #    loss=loss,
        #    n_epochs=self.current_epoch,        
        #)
        #metrics_logs = self.epoch_summary.get_metrics_logs()
        #self.epoch_summary.set_results(metrics=metrics_logs)
        metrics_logs = self.get_metrics_logs(preds=preds, targets=targets, weights=weights, step_name=step_name, loss=loss)

        return metrics_logs

    def training_epoch_end(self, outputs: Dict):
        """
        Nothing happens at the end of the training epoch.
        It serves no purpose to do a general step for the training,
        but it can explode the RAM when using a large dataset.
        """
        pass

    def validation_epoch_end(self, outputs: Dict[str, Any]):

        metrics_logs = self._general_epoch_end(outputs=outputs, step_name="val")

        lr = self.optimizers().param_groups[0]["lr"]
        metrics_logs["lr"] = lr
        metrics_logs["n_epochs"] = self.current_epoch
        self.log_dict(metrics_logs)

        # Save yaml file with the metrics summaries
        #full_dict = {}
        #full_dict.update(self.epoch_summary.get_dict_summary())
        #tb_path = self.logger.log_dir

        # Write the YAML file with the metrics
        #if self.current_epoch >= 0:
        #    mkdir(tb_path)
        #    with open(os.path.join(tb_path, "metrics.yaml"), "w") as file:
        #        yaml.dump(full_dict, file)

    def test_epoch_end(self, outputs: Dict[str, Any]):

        metrics_logs = self._general_epoch_end(outputs=outputs, step_name="test")
        self.log_dict(metrics_logs)

        # Save yaml file with the metrics summaries
        full_dict = {}
        #full_dict.update(self.epoch_summary.get_dict_summary())
        #tb_path = self.logger.log_dir
        #os.makedirs(tb_path, exist_ok=True)
        #with open(f"{tb_path}/metrics.yaml", "w") as file:
        #    yaml.dump(full_dict, file)

    def on_train_start(self):
        hparams_log = deepcopy(self.hparams)
        hparams_log["n_params"] = self.n_params
        #self.logger.log_hyperparams(hparams_log, self.epoch_summary.get_results("val").metrics)

    def get_progress_bar_dict(self) -> Dict[str, float]:
        prog_dict = super().get_progress_bar_dict()
        #results_on_progress_bar = self.epoch_summary.get_results_on_progress_bar("val")
        #prog_dict["loss/val"] = self.epoch_summary.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 GOLI_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 `goli.trainer.PredictorModule.list_pretrained_models()`.
        """

        if name not in GOLI_PRETRAINED_MODELS:
            raise ValueError(
                f"The model '{name}' is not available. Choose from {set(GOLI_PRETRAINED_MODELS.keys())}."
            )

        return PredictorModule.load_from_checkpoint(GOLI_PRETRAINED_MODELS[name])
